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;
@@ -343,7 +344,9 @@ int main(int argc, char **argv)
     path labels_file_to_use = final_index_path_prefix + "_label_formatted.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::set<std::string> raw_universal_label_set = {universal_label};
+    diskann::InMemFilterStore<uint32_t>::convert_labels_string_to_int(label_data_path, labels_file_to_use,
+                                                                      labels_map_file, raw_universal_label_set);
 
     // 2. parse label file and create necessary data structures
     std::vector<label_set> point_ids_to_labels;

apps/test_insert_deletes_consolidate.cpp

@@ -189,8 +189,7 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa
 
     if (universal_label != "")
     {
-        LabelT u_label = 0;
-        index->set_universal_label(u_label);
+        index->set_universal_labels({universal_label}, /*dynamic index*/ true);
     }
 
     if (points_to_skip > num_points)
@@ -265,7 +264,8 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa
         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);
+            diskann::InMemFilterStore<LabelT>::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);
             location_to_labels = std::get<0>(parse_result);
         }
@@ -315,7 +315,8 @@ void build_incremental_index(const std::string &data_path, diskann::IndexWritePa
         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);
+            diskann::InMemFilterStore<LabelT>::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);
             location_to_labels = std::get<0>(parse_result);
         }

apps/test_streaming_scenario.cpp

@@ -219,7 +219,8 @@ void build_incremental_index(const std::string &data_path, const uint32_t L, con
     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);
+        diskann::InMemFilterStore<LabelT>::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);
         pts_to_labels = std::get<0>(parse_result);
     }
@@ -253,8 +254,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}, true);
     }
 
     if (max_points_to_insert == 0)

include/abstract_filter_store.h

@@ -0,0 +1,78 @@
+#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<label_type> &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_converted_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;
+    DISKANN_DLLEXPORT virtual void calculate_best_medoids(const size_t num_points_to_load,
+                                                          const uint32_t num_candidates) = 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::vector<std::string> &universal_labels,
+                                                        bool dynamic_index = false) = 0;
+    // DISKANN_DLLEXPORT virtual const label_type get_universal_label() const = 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) = 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

@@ -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);
+    virtual void set_universal_labels(const std::vector<std::string> &raw_universal_labels,
+                                      bool dynamic_index = false) = 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/in_mem_filter_store.h

@@ -0,0 +1,85 @@
+#pragma once
+#include <abstract_filter_store.h>
+
+namespace diskann
+{
+
+// This class is responsible for filter actions in index, and should not be used outside.
+template <typename label_type> class InMemFilterStore : public AbstractFilterStore<label_type>
+{
+  public:
+    InMemFilterStore(const size_t num_points);
+    ~InMemFilterStore() = default;
+
+    // needs some internal lock
+    bool detect_common_filters(uint32_t point_id, bool search_invocation,
+                               const std::vector<label_type> &incoming_labels,
+                               const FilterMatchStrategy filter_match_strategy) override;
+
+    const std::vector<label_type> &get_labels_by_location(const location_t point_id) override;
+    void set_labels_to_location(const location_t location, const std::vector<label_type> &labels);
+    void swap_labels(const location_t location_first, const location_t location_second) override;
+    const tsl::robin_set<label_type> &get_all_label_set() override;
+    void add_to_label_set(label_type &label) override;
+    // Throws: out of range exception
+    void add_label_to_location(const location_t point_id, label_type label) override;
+    // returns internal mapping for given raw_label
+    label_type get_converted_label(const std::string &raw_label) override;
+
+    void update_medoid_by_label(const label_type &label, const uint32_t new_medoid) override;
+    const uint32_t &get_medoid_by_label(const label_type &label) override;
+    const std::unordered_map<label_type, uint32_t> &get_labels_to_medoids() override;
+    bool label_has_medoid(const label_type &label) override;
+    void calculate_best_medoids(const size_t num_points_to_load, const uint32_t num_candidates) override;
+
+    // takes raw universal labels and map them internally.
+    void set_universal_labels(const std::vector<std::string> &raw_universal_labels,
+                              bool dyanmic_index = false) override;
+    // const label_type get_universal_label() const;
+
+    // ideally takes raw label file and then genrate internal mapping file and keep the info of mapping
+    size_t load_raw_labels(const std::string &raw_labels_file) override;
+
+    void save_labels(const std::string &save_path, const size_t total_points) override;
+    // 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.
+    void save_raw_labels(const std::string &save_path, const size_t total_points) override;
+    void save_medoids(const std::string &save_path) override;
+    void save_label_map(const std::string &save_path) override;
+    void save_universal_label(const std::string &save_path) override;
+
+    // The function is static so it remains the source of truth across the code. Returns label map
+    DISKANN_DLLEXPORT static std::unordered_map<std::string, label_type> convert_labels_string_to_int(
+        const std::string &inFileName, const std::string &outFileName, const std::string &mapFileName,
+        const std::set<std::string> &raw_universal_labels);
+
+  protected:
+    // This is for internal use and only loads already parsed file, used by index in during load().
+    size_t load_labels(const std::string &labels_file) override;
+    size_t load_medoids(const std::string &labels_to_medoid_file) override;
+    void load_label_map(const std::string &labels_map_file) override;
+    void load_universal_labels(const std::string &universal_labels_file) override;
+
+  private:
+    size_t _num_points;
+    std::vector<std::vector<label_type>> _location_to_labels;
+    tsl::robin_set<label_type> _labels;
+    std::unordered_map<std::string, label_type> _label_map;
+
+    // medoids
+    std::unordered_map<label_type, uint32_t> _label_to_medoid_id;
+    std::unordered_map<uint32_t, uint32_t> _medoid_counts; // medoids only happen for filtered index
+
+    // universal label
+    bool _use_universal_label = false;
+    tsl::robin_set<label_type> _mapped_universal_labels;
+    std::set<std::string> _raw_universal_labels;
+
+    // populates _loaction_to labels and _labels from given label file
+    size_t parse_label_file(const std::string &label_file);
+
+    bool detect_common_filters_by_set_intersection(uint32_t point_id, bool search_invocation,
+                                                   const std::vector<label_type> &incoming_labels);
+};
+
+} // namespace diskann

include/index.h

@@ -21,6 +21,7 @@
 #include "in_mem_data_store.h"
 #include "in_mem_graph_store.h"
 #include "abstract_index.h"
+#include "in_mem_filter_store.h"
 
 #define OVERHEAD_FACTOR 1.1
 #define EXPAND_IF_FULL 0
@@ -103,7 +104,8 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
                                                 const size_t num_points_to_load,
                                                 const std::vector<TagT> &tags = std::vector<TagT>());
 
-    DISKANN_DLLEXPORT void set_universal_label(const LabelT &label);
+    // DISKANN_DLLEXPORT void set_universal_label(const LabelT &label);
+    DISKANN_DLLEXPORT void set_universal_labels(const std::vector<std::string> &raw_labels, bool dynamic_index = false);
 
     // Get converted integer label from string to int map (_label_map)
     DISKANN_DLLEXPORT LabelT get_converted_label(const std::string &raw_label);
@@ -222,8 +224,6 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     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) override;
 
-    virtual void _set_universal_label(const LabelType universal_label) override;
-
     // No copy/assign.
     Index(const Index<T, TagT, LabelT> &) = delete;
     Index<T, TagT, LabelT> &operator=(const Index<T, TagT, LabelT> &) = delete;
@@ -334,6 +334,9 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     // Graph related data structures
     std::unique_ptr<AbstractGraphStore> _graph_store;
 
+    // Filter Store
+    std::unique_ptr<AbstractFilterStore<LabelT>> _filter_store;
+
     char *_opt_graph = nullptr;
 
     T *_data = nullptr; // coordinates of all base points
@@ -369,18 +372,14 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     // Filter Support
 
     bool _filtered_index = false;
-    // Location to label is only updated during insert_point(), all other reads are protected by
-    // default as a location can only be released at end of consolidate deletes
-    std::vector<std::vector<LabelT>> _location_to_labels;
-    tsl::robin_set<LabelT> _labels;
-    std::string _labels_file;
-    std::unordered_map<LabelT, uint32_t> _label_to_start_id;
-    std::unordered_map<uint32_t, uint32_t> _medoid_counts;
-
-    bool _use_universal_label = false;
-    LabelT _universal_label = 0;
+    /*  std::vector<std::vector<LabelT>> _pts_to_labels;
+      tsl::robin_set<LabelT> _labels;
+      std::unordered_map<LabelT, uint32_t> _label_to_medoid_id;
+      std::unordered_map<uint32_t, uint32_t> _medoid_counts;
+      bool _use_universal_label = false;
+      LabelT _universal_label = 0;
+      std::unordered_map<std::string, LabelT> _label_map;*/
     uint32_t _filterIndexingQueueSize;
-    std::unordered_map<std::string, LabelT> _label_map;
 
     // Indexing parameters
     uint32_t _indexingQueueSize;

include/index_config.h

@@ -45,7 +45,7 @@ struct IndexConfig
     IndexConfig(DataStoreStrategy data_strategy, GraphStoreStrategy graph_strategy, Metric metric, size_t dimension,
                 size_t max_points, size_t num_pq_chunks, size_t num_frozen_points, bool dynamic_index, bool enable_tags,
                 bool pq_dist_build, bool concurrent_consolidate, bool use_opq, bool filtered_index,
-                std::string &data_type, const std::string &tag_type, const std::string &label_type,
+                const std::string &data_type, const std::string &tag_type, const std::string &label_type,
                 std::shared_ptr<IndexWriteParameters> index_write_params,
                 std::shared_ptr<IndexSearchParams> index_search_params)
         : data_strategy(data_strategy), graph_strategy(graph_strategy), metric(metric), dimension(dimension),
@@ -124,12 +124,6 @@ class IndexConfigBuilder
         return *this;
     }
 
-    IndexConfigBuilder &is_filtered(bool is_filtered)
-    {
-        this->_filtered_index = is_filtered;
-        return *this;
-    }
-
     IndexConfigBuilder &with_num_pq_chunks(size_t num_pq_chunks)
     {
         this->_num_pq_chunks = num_pq_chunks;
@@ -194,6 +188,12 @@ class IndexConfigBuilder
         return *this;
     }
 
+    IndexConfigBuilder &is_filtered(bool is_filtered)
+    {
+        this->_filtered_index = is_filtered;
+        return *this;
+    }
+
     IndexConfig build()
     {
         if (_data_type == "" || _data_type.empty())