patelyash/index factory

apps/build_memory_index.cpp

@@ -17,6 +17,7 @@
 
 #include "memory_mapper.h"
 #include "ann_exception.h"
+#include "index_factory.h"
 
 namespace po = boost::program_options;
 
@@ -155,46 +156,42 @@ int main(int argc, char **argv)
     {
         diskann::cout << "Starting index build with R: " << R << "  Lbuild: " << L << "  alpha: " << alpha
                       << "  #threads: " << num_threads << std::endl;
-        if (label_file != "" && label_type == "ushort")
-        {
-            if (data_type == std::string("int8"))
-                return build_in_memory_index<int8_t, uint32_t, uint16_t>(
-                    metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes,
-                    use_opq, label_file, universal_label, Lf);
-            else if (data_type == std::string("uint8"))
-                return build_in_memory_index<uint8_t, uint32_t, uint16_t>(
-                    metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes,
-                    use_opq, label_file, universal_label, Lf);
-            else if (data_type == std::string("float"))
-                return build_in_memory_index<float, uint32_t, uint16_t>(
-                    metric, data_path, R, L, alpha, index_path_prefix, num_threads, use_pq_build, build_PQ_bytes,
-                    use_opq, label_file, universal_label, Lf);
-            else
-            {
-                std::cout << "Unsupported type. Use one of int8, uint8 or float." << std::endl;
-                return -1;
-            }
-        }
-        else
-        {
-            if (data_type == std::string("int8"))
-                return build_in_memory_index<int8_t>(metric, data_path, R, L, alpha, index_path_prefix, num_threads,
-                                                     use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label,
-                                                     Lf);
-            else if (data_type == std::string("uint8"))
-                return build_in_memory_index<uint8_t>(metric, data_path, R, L, alpha, index_path_prefix, num_threads,
-                                                      use_pq_build, build_PQ_bytes, use_opq, label_file,
-                                                      universal_label, Lf);
-            else if (data_type == std::string("float"))
-                return build_in_memory_index<float>(metric, data_path, R, L, alpha, index_path_prefix, num_threads,
-                                                    use_pq_build, build_PQ_bytes, use_opq, label_file, universal_label,
-                                                    Lf);
-            else
-            {
-                std::cout << "Unsupported type. Use one of int8, uint8 or float." << std::endl;
-                return -1;
-            }
-        }
+
+        size_t data_num, data_dim;
+        diskann::get_bin_metadata(data_path, data_num, data_dim);
+
+        auto config = diskann::IndexConfigBuilder()
+                          .with_metric(metric)
+                          .with_dimension(data_dim)
+                          .with_max_points(data_num)
+                          .with_data_load_store_strategy(diskann::MEMORY)
+                          .with_data_type(data_type)
+                          .with_label_type(label_type)
+                          .is_dynamic_index(false)
+                          .is_enable_tags(false)
+                          .is_use_opq(use_opq)
+                          .is_pq_dist_build(use_pq_build)
+                          .with_num_pq_chunks(build_PQ_bytes)
+                          .build();
+
+        auto index_build_params = diskann::IndexWriteParametersBuilder(L, R)
+                                      .with_filter_list_size(Lf)
+                                      .with_alpha(alpha)
+                                      .with_saturate_graph(false)
+                                      .with_num_threads(num_threads)
+                                      .build();
+
+        auto build_params = diskann::IndexBuildParamsBuilder(index_build_params)
+                                .with_universal_label(universal_label)
+                                .with_label_file(label_file)
+                                .with_save_path_prefix(index_path_prefix)
+                                .build();
+        auto index_factory = diskann::IndexFactory(config);
+        auto index = index_factory.create_instance();
+        index->build(data_path, data_num, build_params);
+        index->save(index_path_prefix.c_str());
+        index.reset();
+        return 0;
     }
     catch (const std::exception &e)
     {

apps/search_memory_index.cpp

@@ -20,6 +20,7 @@
 #include "index.h"
 #include "memory_mapper.h"
 #include "utils.h"
+#include "index_factory.h"
 
 namespace po = boost::program_options;
 
@@ -30,14 +31,14 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
                         const bool dynamic, const bool tags, const bool show_qps_per_thread,
                         const std::vector<std::string> &query_filters, const float fail_if_recall_below)
 {
+    using TagT = uint32_t;
     // Load the query file
     T *query = nullptr;
     uint32_t *gt_ids = nullptr;
     float *gt_dists = nullptr;
     size_t query_num, query_dim, query_aligned_dim, gt_num, gt_dim;
     diskann::load_aligned_bin<T>(query_file, query, query_num, query_dim, query_aligned_dim);
 
-    // Check for ground truth
     bool calc_recall_flag = false;
     if (truthset_file != std::string("null") && file_exists(truthset_file))
     {
@@ -66,18 +67,32 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
         }
     }
 
-    using TagT = uint32_t;
-    const bool concurrent = false, pq_dist_build = false, use_opq = false;
-    const size_t num_pq_chunks = 0;
-    using IndexType = diskann::Index<T, TagT, LabelT>;
-    const size_t num_frozen_pts = IndexType::get_graph_num_frozen_points(index_path);
-    IndexType index(metric, query_dim, 0, dynamic, tags, concurrent, pq_dist_build, num_pq_chunks, use_opq,
-                    num_frozen_pts);
-    std::cout << "Index class instantiated" << std::endl;
-    index.load(index_path.c_str(), num_threads, *(std::max_element(Lvec.begin(), Lvec.end())));
+    const size_t num_frozen_pts = diskann::get_graph_num_frozen_points(index_path);
+
+    auto config = diskann::IndexConfigBuilder()
+                      .with_metric(metric)
+                      .with_dimension(query_dim)
+                      .with_max_points(0)
+                      .with_data_load_store_strategy(diskann::MEMORY)
+                      .with_data_type(diskann_type_to_name<T>())
+                      .with_label_type(diskann_type_to_name<LabelT>())
+                      .with_tag_type(diskann_type_to_name<TagT>())
+                      .is_dynamic_index(dynamic)
+                      .is_enable_tags(tags)
+                      .is_concurrent_consolidate(false)
+                      .is_pq_dist_build(false)
+                      .is_use_opq(false)
+                      .with_num_pq_chunks(0)
+                      .with_num_frozen_pts(num_frozen_pts)
+                      .build();
+
+    auto index_factory = diskann::IndexFactory(config);
+    auto index = index_factory.create_instance();
+    index->load(index_path.c_str(), num_threads, *(std::max_element(Lvec.begin(), Lvec.end())));
     std::cout << "Index loaded" << std::endl;
+
     if (metric == diskann::FAST_L2)
-        index.optimize_index_layout();
+        index->optimize_index_layout();
 
     std::cout << "Using " << num_threads << " threads to search" << std::endl;
     std::cout.setf(std::ios_base::fixed, std::ios_base::floatfield);
@@ -148,29 +163,22 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
             auto qs = std::chrono::high_resolution_clock::now();
             if (filtered_search)
             {
-                LabelT filter_label_as_num;
-                if (query_filters.size() == 1)
-                {
-                    filter_label_as_num = index.get_converted_label(query_filters[0]);
-                }
-                else
-                {
-                    filter_label_as_num = index.get_converted_label(query_filters[i]);
-                }
-                auto retval = index.search_with_filters(query + i * query_aligned_dim, filter_label_as_num, recall_at,
-                                                        L, query_result_ids[test_id].data() + i * recall_at,
-                                                        query_result_dists[test_id].data() + i * recall_at);
+                std::string raw_filter = query_filters.size() == 1 ? query_filters[0] : query_filters[i];
+
+                auto retval = index->search_with_filters(query + i * query_aligned_dim, raw_filter, recall_at, L,
+                                                         query_result_ids[test_id].data() + i * recall_at,
+                                                         query_result_dists[test_id].data() + i * recall_at);
                 cmp_stats[i] = retval.second;
             }
             else if (metric == diskann::FAST_L2)
             {
-                index.search_with_optimized_layout(query + i * query_aligned_dim, recall_at, L,
-                                                   query_result_ids[test_id].data() + i * recall_at);
+                index->search_with_optimized_layout(query + i * query_aligned_dim, recall_at, L,
+                                                    query_result_ids[test_id].data() + i * recall_at);
             }
             else if (tags)
             {
-                index.search_with_tags(query + i * query_aligned_dim, recall_at, L,
-                                       query_result_tags.data() + i * recall_at, nullptr, res);
+                index->search_with_tags(query + i * query_aligned_dim, recall_at, L,
+                                        query_result_tags.data() + i * recall_at, nullptr, res);
                 for (int64_t r = 0; r < (int64_t)recall_at; r++)
                 {
                     query_result_ids[test_id][recall_at * i + r] = query_result_tags[recall_at * i + r];
@@ -179,8 +187,8 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
             else
             {
                 cmp_stats[i] = index
-                                   .search(query + i * query_aligned_dim, recall_at, L,
-                                           query_result_ids[test_id].data() + i * recall_at)
+                                   ->search(query + i * query_aligned_dim, recall_at, L,
+                                            query_result_ids[test_id].data() + i * recall_at)
                                    .second;
             }
             auto qe = std::chrono::high_resolution_clock::now();
@@ -245,7 +253,6 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
     }
 
     diskann::aligned_free(query);
-
     return best_recall >= fail_if_recall_below ? 0 : -1;
 }