Bann save load one file Test PR

apps/build_disk_index.cpp

@@ -103,13 +103,21 @@ int main(int argc, char **argv)
 
     bool use_filters = (label_file != "") ? true : false;
     diskann::Metric metric;
-    if (dist_fn == std::string("l2"))
+    if (dist_fn == std::string("l2")) 
+    {
         metric = diskann::Metric::L2;
+    }
     else if (dist_fn == std::string("mips"))
+    {
         metric = diskann::Metric::INNER_PRODUCT;
+    }
+    else if (dist_fn == std::string("cosine"))
+    {
+        metric = diskann::Metric::COSINE;
+    }
     else
     {
-        std::cout << "Error. Only l2 and mips distance functions are supported" << std::endl;
+        std::cout << "Error. Only l2, cosine, and mips distance functions are supported" << std::endl;
         return -1;
     }
 

apps/search_disk_index.cpp

@@ -4,6 +4,7 @@
 #include "common_includes.h"
 #include <boost/program_options.hpp>
 
+#include "utils.h"
 #include "index.h"
 #include "disk_utils.h"
 #include "math_utils.h"
@@ -47,6 +48,44 @@ void print_stats(std::string category, std::vector<float> percentiles, std::vect
     diskann::cout << std::endl;
 }
 
+template<typename T, typename LabelT>
+void parse_labels_of_query(const std::string &filters_for_query,
+                                      std::unique_ptr<diskann::PQFlashIndex<T, LabelT>> &pFlashIndex,
+                                      std::vector<LabelT> &label_ids_for_query)
+{
+    std::vector<std::string> label_strs_for_query;
+    diskann::split_string(filters_for_query, FILTER_OR_SEPARATOR, label_strs_for_query);
+    for (auto &label_str_for_query : label_strs_for_query)
+    {
+        label_ids_for_query.push_back(pFlashIndex->get_converted_label(label_str_for_query));
+    }
+}
+
+template<typename T, typename LabelT> 
+void populate_label_ids(const std::vector<std::string> &filters_of_queries,
+                        std::unique_ptr<diskann::PQFlashIndex<T, LabelT>> &pFlashIndex,
+                        std::vector<std::vector<LabelT>> &label_ids_of_queries, bool apply_one_to_all, uint32_t query_count)
+{
+    if (apply_one_to_all)
+    {
+        std::vector<LabelT> label_ids_of_query;
+        parse_labels_of_query(filters_of_queries[0], pFlashIndex, label_ids_of_query);
+        for (uint32_t i = 0; i < query_count; i++)
+        {
+            label_ids_of_queries.push_back(label_ids_of_query);
+        }
+    }
+    else
+    {
+        for (auto &filters_of_query : filters_of_queries)
+        {
+            std::vector<LabelT> label_ids_of_query;
+            parse_labels_of_query(filters_of_query, pFlashIndex, label_ids_of_query);
+            label_ids_of_queries.push_back(label_ids_of_query);
+        }
+    }
+}
+
 template <typename T, typename LabelT = uint32_t>
 int search_disk_index(diskann::Metric &metric, const std::string &index_path_prefix,
                       const std::string &result_output_prefix, const std::string &query_file, std::string &gt_file,
@@ -173,6 +212,14 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
         diskann::cout << "..done" << std::endl;
     }
 
+    std::vector<std::vector<LabelT>> per_query_label_ids;
+    if (filtered_search)
+    {
+        populate_label_ids(query_filters, _pFlashIndex, per_query_label_ids, (query_filters.size() == 1), query_num );
+    }
+
+
+
     diskann::cout.setf(std::ios_base::fixed, std::ios_base::floatfield);
     diskann::cout.precision(2);
 
@@ -236,19 +283,10 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
             }
             else
             {
-                LabelT label_for_search;
-                if (query_filters.size() == 1)
-                { // one label for all queries
-                    label_for_search = _pFlashIndex->get_converted_label(query_filters[0]);
-                }
-                else
-                { // one label for each query
-                    label_for_search = _pFlashIndex->get_converted_label(query_filters[i]);
-                }
                 _pFlashIndex->cached_beam_search(
                     query + (i * query_aligned_dim), recall_at, L, query_result_ids_64.data() + (i * recall_at),
-                    query_result_dists[test_id].data() + (i * recall_at), optimized_beamwidth, true, label_for_search,
-                    use_reorder_data, stats + i);
+                    query_result_dists[test_id].data() + (i * recall_at), optimized_beamwidth, true, per_query_label_ids[i],
+                    search_io_limit, use_reorder_data, stats + i);
             }
         }
         auto e = std::chrono::high_resolution_clock::now();
@@ -270,6 +308,9 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
         auto mean_cpuus = diskann::get_mean_stats<float>(stats, query_num,
                                                          [](const diskann::QueryStats &stats) { return stats.cpu_us; });
 
+        auto mean_hops = diskann::get_mean_stats<uint32_t>(
+            stats, query_num, [](const diskann::QueryStats &stats) { return stats.n_hops; });
+
         double recall = 0;
         if (calc_recall_flag)
         {
@@ -283,10 +324,12 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
                       << std::setw(16) << mean_cpuus;
         if (calc_recall_flag)
         {
-            diskann::cout << std::setw(16) << recall << std::endl;
+            diskann::cout << std::setw(16) << recall << std::endl ; 
         }
         else
+        {
             diskann::cout << std::endl;
+        }
         delete[] stats;
     }
 
@@ -443,7 +486,6 @@ int main(int argc, char **argv)
     {
         query_filters = read_file_to_vector_of_strings(query_filters_file);
     }
-
     try
     {
         if (!query_filters.empty() && label_type == "ushort")

apps/search_memory_index.cpp

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT license.
 
 #include <cstring>
@@ -143,6 +143,11 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
     }
 
     double best_recall = 0.0;
+    std::int64_t value = 3;
+    std::function<bool(const int64_t&, float&, bool&)> callback_func = [value](const int64_t &id, float &reRankScore, bool &earlystop) -> bool {
+        //diskann::cout << "check values for ID: " << id << std::endl;
+        return id % value == 0;
+    };
 
     for (uint32_t test_id = 0; test_id < Lvec.size(); test_id++)
     {
@@ -179,8 +184,12 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
             }
             else if (tags)
             {
-                index->search_with_tags(query + i * query_aligned_dim, recall_at, L,
-                                        query_result_tags.data() + i * recall_at, nullptr, res);
+                if (callback_func){
+                    index->search_with_callback(query + i * query_aligned_dim, recall_at, L, query_result_tags.data() + i * recall_at, nullptr, res, callback_func);
+                }else{
+                    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];
@@ -441,7 +450,7 @@ int main(int argc, char **argv)
             else if (data_type == std::string("uint8"))
             {
                 return search_memory_index<uint8_t>(metric, index_path_prefix, result_path, query_file, gt_file,
-                                                    num_threads, K, print_all_recalls, Lvec, dynamic, tags,
+                                                    num_threads, K, print_all_recalls, Lvec, dynamic, tags, 
                                                     show_qps_per_thread, query_filters, fail_if_recall_below);
             }
             else if (data_type == std::string("float"))

apps/test_streaming_scenario.cpp

@@ -393,7 +393,7 @@ int main(int argc, char **argv)
                                        "with each line corresponding to a graph node");
         optional_configs.add_options()("universal_label", po::value<std::string>(&universal_label)->default_value(""),
                                        "Universal label, if using it, only in conjunction with labels_file");
-        optional_configs.add_options()("FilteredLbuild,Lf", po::value<uint32_t>(&Lf)->default_value(0),
+        optional_configs.add_options()("FilteredLbuild", po::value<uint32_t>(&Lf)->default_value(0),
                                        "Build complexity for filtered points, higher value "
                                        "results in better graphs");
         optional_configs.add_options()("label_type", po::value<std::string>(&label_type)->default_value("uint"),

include/abstract_data_store.h

@@ -9,6 +9,8 @@
 #include "types.h"
 #include "windows_customizations.h"
 #include "distance.h"
+#include "aligned_file_reader.h"
+
 
 namespace diskann
 {
@@ -21,13 +23,17 @@ template <typename data_t> class AbstractDataStore
     virtual ~AbstractDataStore() = default;
 
     // Return number of points returned
-    virtual location_t load(const std::string &filename) = 0;
+    virtual location_t load(const std::string &filename, size_t offset) = 0;
+#ifdef EXEC_ENV_OLS
+    virtual location_t load(AlignedFileReader &reader, size_t offset) = 0;
+#endif
 
     // Why does store take num_pts? Since store only has capacity, but we allow
     // resizing we can end up in a situation where the store has spare capacity.
     // To optimize disk utilization, we pass the number of points that are "true"
     // points, so that the store can discard the empty locations before saving.
     virtual size_t save(const std::string &filename, const location_t num_pts) = 0;
+    virtual size_t save(std::ofstream &writer, const location_t num_pts, size_t offset) = 0;
 
     DISKANN_DLLEXPORT virtual location_t capacity() const;
 

include/abstract_graph_store.h

@@ -7,6 +7,8 @@
 #include <vector>
 #include "types.h"
 
+class AlignedFileReader;
+
 namespace diskann
 {
 
@@ -21,11 +23,20 @@ class AbstractGraphStore
     virtual ~AbstractGraphStore() = default;
 
     // returns tuple of <nodes_read, start, num_frozen_points>
-    virtual std::tuple<uint32_t, uint32_t, size_t> load(const std::string &index_path_prefix,
-                                                        const size_t num_points) = 0;
+#ifdef EXEC_ENV_OLS
+    virtual std::tuple<uint32_t, uint32_t, size_t> load(AlignedFileReader &reader, const size_t num_points,
+                                                        size_t offset) = 0;
+#else
+    virtual std::tuple<uint32_t, uint32_t, size_t> load(const std::string &index_path_prefix, const size_t num_points,
+                                                        size_t offset) = 0;
+#endif
+
     virtual int store(const std::string &index_path_prefix, const size_t num_points, const size_t num_fz_points,
                       const uint32_t start) = 0;
 
+    virtual int store(std::ofstream &writer, const size_t num_points, const size_t num_fz_points, const uint32_t start,
+                      size_t offset) = 0;
+
     // not synchronised, user should use lock when necvessary.
     virtual const std::vector<location_t> &get_neighbours(const location_t i) const = 0;
     virtual void add_neighbour(const location_t i, location_t neighbour_id) = 0;

include/abstract_index.h

@@ -64,6 +64,10 @@ class AbstractIndex
     size_t search_with_tags(const data_type *query, const uint64_t K, const uint32_t L, tag_type *tags,
                             float *distances, std::vector<data_type *> &res_vectors);
 
+    // Initialize space for res_vectors before calling.
+    template <typename data_type, typename tag_type>
+    size_t search_with_callback(const data_type *query, const uint64_t K, const uint32_t L, tag_type *tags, float *distances, std::vector<data_type *> &res_vectors, const std::function<bool(const int64_t&, float&, bool&)> callback);
+
     // Added search overload that takes L as parameter, so that we
     // can customize L on a per-query basis without tampering with "Parameters"
     // IDtype is either uint32_t or uint64_t
@@ -121,6 +125,11 @@ class AbstractIndex
     virtual int _get_vector_by_tag(TagType &tag, DataType &vec) = 0;
     virtual size_t _search_with_tags(const DataType &query, const uint64_t K, const uint32_t L, const TagType &tags,
                                      float *distances, DataVector &res_vectors) = 0;
+    virtual size_t _search_with_callback(const DataType &query, const uint64_t K, const uint32_t L, const TagType &tags,
+                                         float *distances, DataVector &res_vectors, const std::function<bool(const int64_t&, float&, bool&)> callback) = 0;
+
+    //virtual size_t _search_with_callback(const DataType &query, const uint64_t K, const uint32_t L, const TagType &tags,
+    //                                 float *distances, DataVector &res_vectors, const std::function<bool(const uint32_t &, float &)> &callback) = 0;
     virtual void _search_with_optimized_layout(const DataType &query, size_t K, size_t L, uint32_t *indices) = 0;
     virtual void _set_universal_label(const LabelType universal_label) = 0;
 };

include/abstract_scratch.h

@@ -0,0 +1,35 @@
+#pragma once
+namespace diskann
+{
+
+template <typename data_t> class PQScratch;
+
+// By somewhat more than a coincidence, it seems that both InMemQueryScratch
+// and SSDQueryScratch have the aligned query and PQScratch objects. So we
+// can put them in a neat hierarchy and keep PQScratch as a standalone class.
+template <typename data_t> class AbstractScratch
+{
+  public:
+    AbstractScratch() = default;
+    // This class does not take any responsibilty for memory management of
+    // its members. It is the responsibility of the derived classes to do so.
+    virtual ~AbstractScratch() = default;
+
+    // Scratch objects should not be copied
+    AbstractScratch(const AbstractScratch &) = delete;
+    AbstractScratch &operator=(const AbstractScratch &) = delete;
+
+    data_t *aligned_query_T()
+    {
+        return _aligned_query_T;
+    }
+    PQScratch<data_t> *pq_scratch()
+    {
+        return _pq_scratch;
+    }
+
+  protected:
+    data_t *_aligned_query_T = nullptr;
+    PQScratch<data_t> *_pq_scratch = nullptr;
+};
+} // namespace diskann

include/aligned_file_reader.h

@@ -117,4 +117,9 @@ class AlignedFileReader
     // process batch of aligned requests in parallel
     // NOTE :: blocking call
     virtual void read(std::vector<AlignedRead> &read_reqs, IOContext &ctx, bool async = false) = 0;
+
+#ifdef USE_BING_INFRA
+    // wait for completion of one request in a batch of requests
+    virtual void wait(IOContext &ctx, int &completedIndex) = 0;
+#endif
 };

include/defaults.h

@@ -17,7 +17,7 @@ const uint32_t NUM_FROZEN_POINTS_STATIC = 0;
 const uint32_t NUM_FROZEN_POINTS_DYNAMIC = 1;
 
 // In-mem index related limits
-const float GRAPH_SLACK_FACTOR = 1.3;
+const float GRAPH_SLACK_FACTOR = 1.3f;
 
 // SSD Index related limits
 const uint64_t MAX_GRAPH_DEGREE = 512;
@@ -30,5 +30,9 @@ const uint32_t MAX_DEGREE = 64;
 const uint32_t BUILD_LIST_SIZE = 100;
 const uint32_t SATURATE_GRAPH = false;
 const uint32_t SEARCH_LIST_SIZE = 100;
+
+const size_t VISITED_RESERVE = 4096;
+const size_t MAX_FILTERS_PER_QUERY = 4096;
+
 } // namespace defaults
 } // namespace diskann

include/distance.h

@@ -1,5 +1,6 @@
 #pragma once
 #include "windows_customizations.h"
+#include <cstdint>
 #include <cstring>
 
 namespace diskann

include/in_mem_data_store.h

@@ -24,8 +24,12 @@ template <typename data_t> class InMemDataStore : public AbstractDataStore<data_
     InMemDataStore(const location_t capacity, const size_t dim, std::unique_ptr<Distance<data_t>> distance_fn);
     virtual ~InMemDataStore();
 
-    virtual location_t load(const std::string &filename) override;
-    virtual size_t save(const std::string &filename, const location_t num_points) override;
+    virtual location_t load(const std::string &filename, size_t offset = 0) override;
+#ifdef EXEC_ENV_OLS
+    virtual location_t load(AlignedFileReader &reader, size_t offset = 0) override;
+#endif
+    virtual size_t save(const std::string &filename, const location_t num_pts) override;
+    virtual size_t save(std::ofstream &writer, const location_t num_pts, size_t offset) override;
 
     virtual size_t get_aligned_dim() const override;
 
@@ -59,9 +63,9 @@ template <typename data_t> class InMemDataStore : public AbstractDataStore<data_
     virtual location_t expand(const location_t new_size) override;
     virtual location_t shrink(const location_t new_size) override;
 
-    virtual location_t load_impl(const std::string &filename);
+    virtual location_t load_impl(const std::string &filename, size_t offset);
 #ifdef EXEC_ENV_OLS
-    virtual location_t load_impl(AlignedFileReader &reader);
+    virtual location_t load_impl(AlignedFileReader &reader, size_t offset);
 #endif
 
   private: