[WIP] On-disk Lloyd clusterization

src/CMakeLists.txt

@@ -28,6 +28,7 @@ add_executable(urukrama
     in_memory_graph.cpp
     on_disk_graph.cpp
     faiss.cpp
+    on_disk_fvecs.cpp
 )
 
 target_link_libraries(urukrama

src/faiss.cpp

@@ -5,6 +5,8 @@
 #include <faiss/IndexPQ.h>
 #include <faiss/impl/ProductQuantizer.h>
 
+namespace urukrama {
+
 using faiss::idx_t;
 
 std::tuple<std::vector<float>, std::vector<int64_t>, std::vector<float>> ComputeClustersCPU(
@@ -65,4 +67,6 @@ faiss::IndexPQ BuildIndexPQ(const std::vector<float>& vecs, size_t dim, size_t M
     index_pq.add(idx_t(vecs.size() / dim), vecs.data());
 
     return index_pq;
-}
+}
+
+}  // namespace urukrama

src/faiss.hpp

@@ -5,6 +5,8 @@
 #include <cstdint>
 #include <vector>
 
+namespace urukrama {
+
 std::tuple<std::vector<float>, std::vector<int64_t>, std::vector<float>> ComputeClusters(
     const std::vector<float>& vecs,
     size_t dim,
@@ -13,4 +15,6 @@ std::tuple<std::vector<float>, std::vector<int64_t>, std::vector<float>> Compute
     size_t number_of_iterations = 20,
     bool verbose = true);
 
-faiss::IndexPQ BuildIndexPQ(const std::vector<float>& vecs, size_t dim, size_t M = 32, size_t nbits = 8);
+faiss::IndexPQ BuildIndexPQ(const std::vector<float>& vecs, size_t dim, size_t M = 32, size_t nbits = 8);
+
+}  // namespace urukrama

src/in_memory_graph.cpp

@@ -5,7 +5,6 @@
 #include "lib/logger/logger.hpp"
 
 #include <tbb/parallel_for.h>
-#include <tbb/parallel_for_each.h>
 
 #include <algorithm>
 #include <random>

src/in_memory_graph.hpp

@@ -2,17 +2,13 @@
 
 #include "types.hpp"
 
-#include <boost/noncopyable.hpp>
-
 
 namespace urukrama {
 
 template <typename T>
-class OnDiskGraph;
-
-template <typename T>
-class InMemoryGraph: boost::noncopyable {
-    friend OnDiskGraph<T>;
+class InMemoryGraph {
+    template <typename, bool>
+    friend class OnDiskGraph;
 
 public:
     InMemoryGraph(std::vector<Point<T>>&& points, size_t R, size_t L);
@@ -27,7 +23,6 @@ class InMemoryGraph: boost::noncopyable {
 
     [[nodiscard]] std::vector<HashSet<size_t>> InitNeighbors() const;
 
-
     [[nodiscard]] GreedySearchResult GreedySearch(size_t s_idx,
                                                   const Point<T>& query,
                                                   const std::vector<HashSet<size_t>>& n_out,

src/main.cpp

@@ -1,5 +1,6 @@
 #include "faiss.hpp"
 #include "in_memory_graph.hpp"
+#include "on_disk_fvecs.hpp"
 #include "on_disk_graph.hpp"
 #include "utils.hpp"
 
@@ -8,17 +9,25 @@
 #include <faiss/IndexPQ.h>
 #include <faiss/index_io.h>
 
-#include <random>
 #include <ranges>
 #include <vector>
 
+constexpr auto FVECS_FILENAME = "/data/deep1M_base.fvecs";
+constexpr auto GRAPH_INDEX_FILENAME = "deep1m.index";
+constexpr auto PQ_INDEX_FILENAME = "deep1m.pq_index";
+constexpr size_t R = 70;
+constexpr size_t L = 75;
+constexpr size_t K = 16;
+constexpr size_t NITER = 50;
+
 void Search(std::span<const urukrama::Point<float>> queries)
 {
-    urukrama::OnDiskGraph<float> on_disk_graph("deep1m.index");
-    auto index_pq = dynamic_cast<faiss::IndexPQ*>(faiss::read_index("deep1m.index_pq"));
+    urukrama::OnDiskGraph<float> on_disk_graph(GRAPH_INDEX_FILENAME);
+    auto index_pq = dynamic_cast<faiss::IndexPQ*>(faiss::read_index(PQ_INDEX_FILENAME));
 
     size_t total_search_time = 0;
-    size_t n_search = 100;
+    size_t cnt = 0;
+    size_t n_search = 50000;
 
     for (const auto& [p_idx, p]: queries | std::views::enumerate | std::views::take(n_search)) {
         using namespace std::chrono;
@@ -28,49 +37,82 @@ void Search(std::span<const urukrama::Point<float>> queries)
         auto t1 = high_resolution_clock::now();
 
         total_search_time += duration_cast<microseconds>(t1 - t0).count();
+        cnt += (top.front().first == 0);
 
         for (const auto [dist, n_idx]: top | std::views::take(1)) {
             ksp::log::Info("Top entry: p_idx=[{}], n_idx=[{}], dist=[{}]", p_idx, n_idx, dist);
         }
     }
 
     ksp::log::Info("Mean processing time: {}us", total_search_time / n_search);
+    ksp::log::Info("Recall: {}", float(cnt) / float(n_search));
 }
 
-void BuildIndex(const std::vector<urukrama::Point<float>>& points)
+void BuildIndexOnDisk(std::string_view filename)
 {
-    urukrama::InMemoryGraph graph(std::vector{points}, 70, 75);
-    urukrama::WriteOnDisk(graph, "deep1m.index");
+    // Build Graph index
+    {
+        urukrama::OnDiskFVecs<false> on_disk_fvecs(FVECS_FILENAME);
+        std::vector clusters = on_disk_fvecs.ComputeClusters(K, NITER);
+        size_t medoid_idx = on_disk_fvecs.FindMedoid();
+
+        auto on_disk_graph = urukrama::OnDiskGraph<float, true>::Init(GRAPH_INDEX_FILENAME,
+                                                                      medoid_idx,
+                                                                      on_disk_fvecs.Dimension(),
+                                                                      on_disk_fvecs.Size(),
+                                                                      R,
+                                                                      L);
+
+        for (const auto& [cluster_idx, p_indices]: clusters | std::views::enumerate) {
+            std::vector<urukrama::Point<float>> points;
+
+            for (size_t p_idx: p_indices) {
+                auto point_span = on_disk_fvecs.GetPoint(p_idx);
 
-    // TODO: inplace
-    std::vector<float> flat_points;
-    for (const auto& x: points | std::views::join) {
-        flat_points.emplace_back(x);
+                points.emplace_back(
+                    Eigen::Map<const urukrama::Point<float>>(point_span.data(), Eigen::Index(point_span.size())));
+            }
+
+            ksp::log::Info("Building graph for cluster: cluster_idx=[{}], cluster_size=[{}]",
+                           cluster_idx,
+                           p_indices.size());
+
+            urukrama::InMemoryGraph graph(std::vector{points}, 70, 75);
+
+            on_disk_graph.Merge(graph, p_indices);
+        }
     }
 
-    auto index_pq = BuildIndexPQ(flat_points, points.front().size());
-    faiss::write_index(&index_pq, "deep1m.index_pq");
+    // Build PQ index
+    {
+        auto [flat_points, dimension, num_points] = urukrama::FVecsRead(FVECS_FILENAME);
+        auto index_pq = urukrama::BuildIndexPQ(flat_points, dimension);
+        faiss::write_index(&index_pq, PQ_INDEX_FILENAME);
+    }
 }
 
-int main()
+auto LoadPoints()
 {
-    auto points = [] {
-        auto [flat_points, dimension, num_points] = urukrama::FVecsRead("/data/deep1M_base.fvecs");
-        std::vector<urukrama::Point<float>> points;
+    auto [flat_points, dimension, num_points] = urukrama::FVecsRead(FVECS_FILENAME);
+    std::vector<urukrama::Point<float>> points;
 
-        for (const auto& p: flat_points | std::views::chunk(dimension)) {
-            points.emplace_back(Eigen::Map<const urukrama::Point<float>>(p.data(), Eigen::Index(dimension)));
-        }
+    for (const auto& p: flat_points | std::views::chunk(dimension)) {
+        points.emplace_back(Eigen::Map<const urukrama::Point<float>>(p.data(), Eigen::Index(dimension)));
+    }
 
-        std::mt19937_64 random_engine{std::random_device{}()};
-        std::ranges::shuffle(points, random_engine);
+    return points;
+}
+
+int main()
+{
+    urukrama::OnDiskFVecs<true> on_disk_fvecs(FVECS_FILENAME);
+    on_disk_fvecs.Shuffle();
 
-        return points;
-    }();
+    BuildIndexOnDisk(FVECS_FILENAME);
 
+    auto points = LoadPoints();
     ksp::log::Info("Loaded points: size=[{}]", points.size());
 
-    BuildIndex(points);
     Search(points);
 
     return 0;