FilterStore: unifying filter specific logic

apps/build_memory_index.cpp

@@ -144,6 +144,7 @@ int main(int argc, char **argv)
                           .with_index_write_params(index_build_params)
                           .is_enable_tags(false)
                           .is_use_opq(use_opq)
+                          .is_filtered(label_file != "")
                           .is_pq_dist_build(use_pq_build)
                           .with_num_pq_chunks(build_PQ_bytes)
                           .build();

apps/build_stitched_index.cpp

@@ -18,6 +18,7 @@
 #include "memory_mapper.h"
 #include "parameters.h"
 #include "utils.h"
+#include "in_mem_filter_store.h"
 #include "program_options_utils.hpp"
 
 namespace po = boost::program_options;
@@ -341,10 +342,12 @@ int main(int argc, char **argv)
     handle_args(argc, argv, data_type, input_data_path, final_index_path_prefix, label_data_path, universal_label,
                 num_threads, R, L, stitched_R, alpha);
 
-    path labels_file_to_use = final_index_path_prefix + "_label_formatted.txt";
+    path labels_file_to_use = final_index_path_prefix + "_label_numeric.txt";
     path labels_map_file = final_index_path_prefix + "_labels_map.txt";
 
-    convert_labels_string_to_int(label_data_path, labels_file_to_use, labels_map_file, universal_label);
+    std::string raw_universal_label = universal_label;
+    diskann::InMemFilterStore<uint32_t>::convert_label_to_numeric(label_data_path, labels_file_to_use, labels_map_file,
+                                                                  raw_universal_label);
 
     // 2. parse label file and create necessary data structures
     std::vector<label_set> point_ids_to_labels;

apps/search_disk_index.cpp

@@ -239,11 +239,11 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
                 LabelT label_for_search;
                 if (query_filters.size() == 1)
                 { // one label for all queries
-                    label_for_search = _pFlashIndex->get_converted_label(query_filters[0]);
+                    label_for_search = _pFlashIndex->get_numeric_label(query_filters[0]);
                 }
                 else
                 { // one label for each query
-                    label_for_search = _pFlashIndex->get_converted_label(query_filters[i]);
+                    label_for_search = _pFlashIndex->get_numeric_label(query_filters[i]);
                 }
                 _pFlashIndex->cached_beam_search(
                     query + (i * query_aligned_dim), recall_at, L, query_result_ids_64.data() + (i * recall_at),

apps/test_insert_deletes_consolidate.cpp

@@ -94,7 +94,7 @@ std::string get_save_filename(const std::string &save_path, size_t points_to_ski
 
 template <typename T, typename TagT, typename LabelT>
 void insert_till_next_checkpoint(diskann::AbstractIndex &index, size_t start, size_t end, int32_t thread_count, T *data,
-                                 size_t aligned_dim, std::vector<std::vector<LabelT>> &location_to_labels)
+                                 size_t aligned_dim, std::vector<std::vector<std::string>> &location_to_labels)
 {
     diskann::Timer insert_timer;
 #pragma omp parallel for num_threads(thread_count) schedule(dynamic)
@@ -187,10 +187,10 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa
     diskann::IndexFactory index_factory = diskann::IndexFactory(index_config);
     auto index = index_factory.create_instance();
 
+    /* remove set_universal_label from here and set it through filter store only*/
     if (universal_label != "")
     {
-        LabelT u_label = 0;
-        index->set_universal_label(u_label);
+        index->set_universal_labels(universal_label);
     }
 
     if (points_to_skip > num_points)
@@ -255,18 +255,15 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa
                   << " points since the data file has only that many" << std::endl;
     }
 
-    std::vector<std::vector<LabelT>> location_to_labels;
+    std::vector<std::vector<std::string>> location_to_labels;
     if (concurrent)
     {
         // handle labels
         const auto save_path_inc = get_save_filename(save_path + ".after-concurrent-delete-", points_to_skip,
                                                      points_to_delete_from_beginning, last_point_threshold);
-        std::string labels_file_to_use = save_path_inc + "_label_formatted.txt";
-        std::string mem_labels_int_map_file = save_path_inc + "_labels_map.txt";
         if (has_labels)
         {
-            convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label);
-            auto parse_result = diskann::parse_formatted_label_file<LabelT>(labels_file_to_use);
+            auto parse_result = diskann::parse_raw_label_file(label_file);
             location_to_labels = std::get<0>(parse_result);
         }
 
@@ -311,12 +308,11 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa
     {
         const auto save_path_inc = get_save_filename(save_path + ".after-delete-", points_to_skip,
                                                      points_to_delete_from_beginning, last_point_threshold);
-        std::string labels_file_to_use = save_path_inc + "_label_formatted.txt";
+        std::string labels_file_to_use = save_path_inc + "_label_numeric.txt";
         std::string mem_labels_int_map_file = save_path_inc + "_labels_map.txt";
         if (has_labels)
         {
-            convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label);
-            auto parse_result = diskann::parse_formatted_label_file<LabelT>(labels_file_to_use);
+            auto parse_result = diskann::parse_raw_label_file(labels_file_to_use);
             location_to_labels = std::get<0>(parse_result);
         }
 

apps/test_streaming_scenario.cpp

@@ -87,7 +87,7 @@ std::string get_save_filename(const std::string &save_path, size_t active_window
 
 template <typename T, typename TagT, typename LabelT>
 void insert_next_batch(diskann::AbstractIndex &index, size_t start, size_t end, size_t insert_threads, T *data,
-                       size_t aligned_dim, std::vector<std::vector<LabelT>> &pts_to_labels)
+                       size_t aligned_dim, std::vector<std::vector<std::string>> &pts_to_labels)
 {
     try
     {
@@ -211,16 +211,14 @@ void build_incremental_index(const std::string &data_path, const uint32_t L, con
     size_t dim, aligned_dim;
     size_t num_points;
 
-    std::vector<std::vector<LabelT>> pts_to_labels;
+    std::vector<std::vector<std::string>> pts_to_labels;
 
     const auto save_path_inc =
         get_save_filename(save_path + ".after-streaming-", active_window, consolidate_interval, max_points_to_insert);
-    std::string labels_file_to_use = save_path_inc + "_label_formatted.txt";
-    std::string mem_labels_int_map_file = save_path_inc + "_labels_map.txt";
+
     if (has_labels)
     {
-        convert_labels_string_to_int(label_file, labels_file_to_use, mem_labels_int_map_file, universal_label);
-        auto parse_result = diskann::parse_formatted_label_file<LabelT>(labels_file_to_use);
+        auto parse_result = diskann::parse_raw_label_file(label_file);
         pts_to_labels = std::get<0>(parse_result);
     }
 
@@ -253,8 +251,7 @@ void build_incremental_index(const std::string &data_path, const uint32_t L, con
 
     if (universal_label != "")
     {
-        LabelT u_label = 0;
-        index->set_universal_label(u_label);
+        index->set_universal_labels(universal_label);
     }
 
     if (max_points_to_insert == 0)

include/abstract_filter_store.h

@@ -0,0 +1,76 @@
+#pragma once
+#include "common_includes.h"
+#include "utils.h"
+#include <any>
+
+namespace diskann
+{
+
+enum class FilterMatchStrategy
+{
+    SET_INTERSECTION
+};
+// This class is responsible for filter actions in index, and should not be used outside.
+template <typename label_type> class AbstractFilterStore
+{
+  public:
+    DISKANN_DLLEXPORT AbstractFilterStore(const size_t num_points);
+    virtual ~AbstractFilterStore() = default;
+
+    // needs some internal lock + abstract implementation
+    DISKANN_DLLEXPORT virtual bool detect_common_filters(
+        uint32_t point_id, bool search_invocation, const std::vector<label_type> &incoming_labels,
+        const FilterMatchStrategy strategy = FilterMatchStrategy::SET_INTERSECTION) = 0;
+
+    DISKANN_DLLEXPORT virtual const std::vector<label_type> &get_labels_by_location(const location_t point_id) = 0;
+    DISKANN_DLLEXPORT virtual void set_labels_to_location(const location_t location,
+                                                          const std::vector<std::string> &labels) = 0;
+    DISKANN_DLLEXPORT virtual void swap_labels(const location_t location_first, const location_t location_second) = 0;
+
+    DISKANN_DLLEXPORT virtual const tsl::robin_set<label_type> &get_all_label_set() = 0;
+    DISKANN_DLLEXPORT virtual void add_to_label_set(label_type &label) = 0;
+    // Throws: out of range exception
+    DISKANN_DLLEXPORT virtual void add_label_to_location(const location_t point_id, label_type label) = 0;
+    // returns internal mapping for given raw_label
+    DISKANN_DLLEXPORT virtual label_type get_numeric_label(const std::string &raw_label) = 0;
+
+    DISKANN_DLLEXPORT virtual void update_medoid_by_label(const label_type &label, const uint32_t new_medoid) = 0;
+    DISKANN_DLLEXPORT virtual const uint32_t &get_medoid_by_label(const label_type &label) = 0;
+    DISKANN_DLLEXPORT virtual const std::unordered_map<label_type, uint32_t> &get_labels_to_medoids() = 0;
+    DISKANN_DLLEXPORT virtual bool label_has_medoid(const label_type &label) = 0;
+
+    // TODO: in future we may accept a set or vector of universal labels
+    // DISKANN_DLLEXPORT virtual void set_universal_label(label_type universal_label) = 0;
+    DISKANN_DLLEXPORT virtual void set_universal_labels(const std::string &universal_labels) = 0;
+    DISKANN_DLLEXPORT virtual std::pair<bool, label_type> get_universal_label() = 0;
+
+    // takes raw label file and then genrate internal mapping file and keep the info of mapping
+    DISKANN_DLLEXPORT virtual size_t load_raw_labels(const std::string &raw_labels_file,
+                                                     const std::string &raw_universal_label) = 0;
+
+    DISKANN_DLLEXPORT virtual void save_labels(const std::string &save_path, const size_t total_points) = 0;
+    // For dynamic filtered build, we compact the data and hence location_to_labels, we need the compacted version of
+    // raw labels to compute GT correctly.
+    DISKANN_DLLEXPORT virtual void save_raw_labels(const std::string &save_path, const size_t total_points) = 0;
+    DISKANN_DLLEXPORT virtual void save_medoids(const std::string &save_path) = 0;
+    DISKANN_DLLEXPORT virtual void save_label_map(const std::string &save_path) = 0;
+    DISKANN_DLLEXPORT virtual void save_universal_label(const std::string &save_path) = 0;
+
+  protected:
+    // This is for internal use and only loads already parsed file
+    DISKANN_DLLEXPORT virtual size_t load_labels(const std::string &labels_file) = 0;
+    DISKANN_DLLEXPORT virtual size_t load_medoids(const std::string &labels_to_medoid_file) = 0;
+    DISKANN_DLLEXPORT virtual void load_label_map(const std::string &labels_map_file) = 0;
+    DISKANN_DLLEXPORT virtual void load_universal_labels(const std::string &universal_labels_file) = 0;
+
+  private:
+    size_t _num_points;
+
+    // populates pts_to labels and _labels from given label file
+    virtual size_t parse_label_file(const std::string &label_file) = 0;
+
+    // mark Index as friend so it can access protected loads
+    template <typename T, typename TagT, typename LabelT> friend class Index;
+};
+
+} // namespace diskann

include/abstract_index.h

@@ -79,8 +79,8 @@ class AbstractIndex
                                                       float *distances);
 
     // insert points with labels, labels should be present for filtered index
-    template <typename data_type, typename tag_type, typename label_type>
-    int insert_point(const data_type *point, const tag_type tag, const std::vector<label_type> &labels);
+    template <typename data_type, typename tag_type>
+    int insert_point(const data_type *point, const tag_type tag, const std::vector<std::string> &labels);
 
     // insert point for unfiltered index build. do not use with filtered index
     template <typename data_type, typename tag_type> int insert_point(const data_type *point, const tag_type tag);
@@ -103,7 +103,8 @@ class AbstractIndex
     // memory should be allocated for vec before calling this function
     template <typename tag_type, typename data_type> int get_vector_by_tag(tag_type &tag, data_type *vec);
 
-    template <typename label_type> void set_universal_label(const label_type universal_label);
+    // required for dynamic index (they dont use filter store / data store yet)
+    virtual void set_universal_labels(const std::string &raw_universal_labels) = 0;
 
   private:
     virtual void _build(const DataType &data, const size_t num_points_to_load, TagVector &tags) = 0;
@@ -122,6 +123,5 @@ class AbstractIndex
     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 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;
 };
 } // namespace diskann

include/disk_utils.h

@@ -65,7 +65,8 @@ DISKANN_DLLEXPORT int merge_shards(const std::string &vamana_prefix, const std::
                                    const std::string &idmaps_prefix, const std::string &idmaps_suffix,
                                    const uint64_t nshards, uint32_t max_degree, const std::string &output_vamana,
                                    const std::string &medoids_file, bool use_filters = false,
-                                   const std::string &labels_to_medoids_file = std::string(""));
+                                   const std::string &labels_to_medoids_file = std::string(""),
+                                   const std::unordered_map<std::string, uint32_t> &disk_labels_map = {});
 
 DISKANN_DLLEXPORT void extract_shard_labels(const std::string &in_label_file, const std::string &shard_ids_bin,
                                             const std::string &shard_label_file);
@@ -81,7 +82,7 @@ DISKANN_DLLEXPORT int build_merged_vamana_index(std::string base_file, diskann::
                                                 std::string centroids_file, size_t build_pq_bytes, bool use_opq,
                                                 uint32_t num_threads, bool use_filters = false,
                                                 const std::string &label_file = std::string(""),
-                                                const std::string &labels_to_medoids_file = std::string(""),
+                                                const std::string &disk_labels_to_medoids_file = std::string(""),
                                                 const std::string &universal_label = "", const uint32_t Lf = 0);
 
 template <typename T, typename LabelT>

include/filter_utils.h

@@ -57,8 +57,7 @@ DISKANN_DLLEXPORT void generate_label_indices(path input_data_path, path final_i
 DISKANN_DLLEXPORT load_label_index_return_values load_label_index(path label_index_path,
                                                                   uint32_t label_number_of_points);
 
-template <typename LabelT>
-DISKANN_DLLEXPORT std::tuple<std::vector<std::vector<LabelT>>, tsl::robin_set<LabelT>> parse_formatted_label_file(
+DISKANN_DLLEXPORT std::tuple<std::vector<std::vector<std::string>>, tsl::robin_set<std::string>> parse_raw_label_file(
     path label_file);
 
 DISKANN_DLLEXPORT parse_label_file_return_values parse_label_file(path label_data_path, std::string universal_label);