support integer filter label

Author: Sanhaoji2

include/abstract_index.h

@@ -53,7 +53,7 @@ class AbstractIndex
 #ifdef EXEC_ENV_OLS
     virtual void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l) = 0;
 #else
-    virtual void load(const char *index_file, uint32_t num_threads, uint32_t search_l, bool loadBitmaskLabelFile = false) = 0;
+    virtual void load(const char *index_file, uint32_t num_threads, uint32_t search_l, LabelFormatType label_format_type = LabelFormatType::String) = 0;
 #endif
 
     // For FastL2 search on optimized layout

include/filter_match_proxy.h

@@ -0,0 +1,66 @@
+#pragma once
+#include "label_bitmask.h"
+#include "integer_label_vector.h"
+
+namespace diskann
+{
+
+    class filter_match_proxy
+    {
+    public:
+        virtual bool contain_filtered_label(uint32_t id) = 0;
+    };
+
+    template <typename LabelT>
+    class bitmask_filter_match : public filter_match_proxy
+    {
+    public:
+        bitmask_filter_match(simple_bitmask_buf& bitmask_filters,
+            std::vector<std::uint64_t>& query_bitmask_buf,
+            const std::vector<LabelT>& filter_labels,
+            LabelT unv_label);
+
+        virtual bool contain_filtered_label(uint32_t id) override;
+
+    private:
+        simple_bitmask_buf& _bitmask_filters;
+        std::vector<std::uint64_t>& _query_bitmask_buf;
+        simple_bitmask_full_val _bitmask_full_val;
+    };
+
+    template <typename LabelT>
+    class integer_label_filter_match : public filter_match_proxy
+    {
+    public:
+        integer_label_filter_match(integer_label_vector& label_vector,
+            const std::vector<LabelT>& filter_labels,
+            LabelT unv_label);
+
+        virtual bool contain_filtered_label(uint32_t id) override;
+
+    private:
+        integer_label_vector& _label_vector;
+        const std::vector<LabelT>& _filter_labels;
+        LabelT _unv_label;
+    };
+
+template <typename LabelT>
+class label_filter_match_holder : public filter_match_proxy
+{
+public:
+    label_filter_match_holder(simple_bitmask_buf& bitmask_filters,
+        std::vector<std::uint64_t>& query_bitmask_buf,
+        integer_label_vector& label_vector,
+        const std::vector<LabelT>& filter_labels,
+        LabelT unv_label,
+        bool use_integer_labels);
+
+    virtual bool contain_filtered_label(uint32_t id) override;
+
+private:
+    bitmask_filter_match<LabelT> _bitmask_filter_match;
+    integer_label_filter_match<LabelT> _integer_label_filter_match;
+    bool _use_integer_labels;
+};
+
+}

include/index.h

@@ -24,6 +24,7 @@
 #include "percentile_stats.h"
 #include <bitset>
 #include "label_bitmask.h"
+#include "integer_label_vector.h"
 
 #include "quantized_distance.h"
 #include "pq_data_store.h"
@@ -80,7 +81,7 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
 #ifdef EXEC_ENV_OLS
     DISKANN_DLLEXPORT void load(AlignedFileReader &reader, uint32_t num_threads, uint32_t search_l);
 #else
-    DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, uint32_t search_l, bool loadBitmaskLabelFile = false);
+    DISKANN_DLLEXPORT void load(const char *index_file, uint32_t num_threads, uint32_t search_l, LabelFormatType label_format_type = LabelFormatType::String);
 #endif
 
     // get some private variables
@@ -118,6 +119,10 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
 
     DISKANN_DLLEXPORT bool is_set_universal_label() const override;
 
+    DISKANN_DLLEXPORT void enable_integer_label();
+    
+    DISKANN_DLLEXPORT bool integer_label_enabled();
+
     // Set starting point of an index before inserting any points incrementally.
     // The data count should be equal to _num_frozen_pts * _aligned_dim.
     DISKANN_DLLEXPORT void set_start_points(const T *data, size_t data_count);
@@ -253,11 +258,18 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     // determines navigating node of the graph by calculating medoid of datafopt
     uint32_t calculate_entry_point();
 
-    void parse_label_file(const std::string &label_file, size_t &num_pts_labels);
+    void parse_label_file(const std::string &label_file, size_t &num_pts_labels, size_t& total_labels);
     void parse_seller_file(const std::string& label_file, size_t& num_pts_labels);
 
     void convert_pts_label_to_bitmask(std::vector<std::vector<LabelT>>& pts_to_labels, simple_bitmask_buf& bitmask_buf, size_t num_labels);
 
+    void convert_pts_label_to_integer_vector(std::vector<std::vector<LabelT>> &pts_to_labels,
+        integer_label_vector &int_label_vector, size_t total_labels);
+    
+    void aggregate_points_by_bitmask_label(std::unordered_map<LabelT, std::vector<uint32_t>>& label_to_points, size_t num_points_to_load);
+
+    void aggregate_points_by_integer_label(std::unordered_map<LabelT, std::vector<uint32_t>>& label_to_points, size_t num_points_to_load);
+
     std::unordered_map<std::string, LabelT> load_label_map(const std::string &map_file);
 
     // Returns the locations of start point and frozen points suitable for use
@@ -463,6 +475,9 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
 
     simple_bitmask_buf _bitmask_buf;
 
+    bool _use_integer_labels = false;
+    integer_label_vector _label_vector;
+
     TableStats _table_stats;
 
     static const float INDEX_GROWTH_FACTOR;

include/integer_label_vector.h

@@ -0,0 +1,43 @@
+#pragma once
+#include <vector>
+#include <string>
+
+namespace diskann
+{
+
+class integer_label_vector
+{
+public:
+    bool initialize(size_t numpoints, size_t total_labels);
+
+    bool initialize_from_file(const std::string &label_file, size_t &numpoints);
+
+    bool write_to_file(const std::string &label_file) const;
+
+    template <typename LabelT>
+    bool add_labels(uint32_t point_id, std::vector<LabelT> &labels);
+    
+    bool check_label_exists(uint32_t point_id, uint32_t label);
+
+    template <typename LabelT>
+    bool check_label_exists(uint32_t point_id, const std::vector<LabelT> &labels);
+
+    bool check_label_full_contain(uint32_t point_id, const std::vector<uint32_t> &labels);
+
+    bool check_label_full_contain(uint32_t source_point, uint32_t target_point);
+
+    const std::vector<size_t> &get_offset_vector() const;
+
+    const std::vector<uint32_t> &get_data_vector() const;
+
+    size_t get_memory_usage() const;
+
+  private:
+    bool binary_search(size_t start, size_t end, uint32_t label, size_t& last_check);
+
+private:
+  std::vector<size_t> _offset;
+  std::vector<uint32_t> _data;
+};
+
+}

include/label_helper.h

@@ -1,6 +1,8 @@
 #pragma once
 #include "label_bitmask.h"
+#include "integer_label_vector.h"
 #include "percentile_stats.h"
+#include "tsl/robin_set.h"
 #include <string>
 
 namespace diskann
@@ -20,6 +22,113 @@ class label_helper
 
     bool read_bitmask_from_file(const std::string &bitmask_label_file, simple_bitmask_buf &bitmask_buf,
                                 size_t& num_points);
+
+    bool parse_label_file_in_integer(
+        const std::string& label_file,
+        size_t& num_points,
+        integer_label_vector& integer_vector,
+        tsl::robin_set<uint32_t>& labels, TableStats &table_stats);
+
+    template <typename LabelT>
+    bool load_label_map(
+        const std::string& label_map_file,
+        std::unordered_map<std::string, LabelT>& label_map)
+    {
+        std::ifstream infile(label_map_file, std::ios::binary);
+        if (infile.fail())
+        {
+            throw diskann::ANNException(std::string("Failed to open file ") + label_map_file, -1);
+        }
+        infile.seekg(0, std::ios::end);
+        size_t file_size = infile.tellg();
+
+        std::string buffer(file_size, ' ');
+
+        infile.seekg(0, std::ios::beg);
+        infile.read(&buffer[0], file_size);
+        infile.close();
+
+        unsigned line_cnt = 0;
+
+        size_t cur_pos = 0;
+        size_t next_pos = 0;
+        size_t lbl_pos = 0;
+        std::string token;
+        std::string labe_str;
+        while (cur_pos < file_size && cur_pos != std::string::npos)
+        {
+            next_pos = buffer.find('\n', cur_pos);
+            if (next_pos == std::string::npos)
+            {
+                break;
+            }
+
+            lbl_pos = search_string_range(buffer, '\t', cur_pos, next_pos);
+            labe_str.assign(buffer.c_str() + cur_pos, lbl_pos - cur_pos);
+
+            token.assign(buffer.c_str() + lbl_pos + 1, next_pos - lbl_pos - 1);
+            LabelT label_num = (LabelT)std::stoul(token);
+
+            label_map[labe_str] = label_num;
+
+            cur_pos = next_pos + 1;
+
+            line_cnt++;
+        }
+
+        return true;
+    }
+
+    template <typename LabelT>
+    bool load_label_medoids(
+        const std::string& label_medoids_file,
+        std::unordered_map<LabelT, uint32_t>& label_to_start_id)
+    {
+        std::ifstream infile(label_medoids_file, std::ios::binary);
+        if (infile.fail())
+        {
+            throw diskann::ANNException(std::string("Failed to open file ") + label_medoids_file, -1);
+        }
+        infile.seekg(0, std::ios::end);
+        size_t file_size = infile.tellg();
+
+        std::string buffer(file_size, ' ');
+
+        infile.seekg(0, std::ios::beg);
+        infile.read(&buffer[0], file_size);
+        infile.close();
+
+        unsigned line_cnt = 0;
+
+        size_t cur_pos = 0;
+        size_t next_pos = 0;
+        size_t lbl_pos = 0;
+        std::string token;
+        while (cur_pos < file_size && cur_pos != std::string::npos)
+        {
+            next_pos = buffer.find('\n', cur_pos);
+            if (next_pos == std::string::npos)
+            {
+                break;
+            }
+
+            lbl_pos = search_string_range(buffer, ',', cur_pos, next_pos);
+            token.assign(buffer.c_str() + cur_pos, lbl_pos - cur_pos);
+            LabelT label_num = (LabelT)std::stoul(token);
+            
+            token.assign(buffer.c_str() + lbl_pos + 1, next_pos - lbl_pos - 1);
+            uint32_t medoid = (uint32_t)std::stoul(token);
+
+            label_to_start_id[label_num] = medoid;
+
+            cur_pos = next_pos + 1;
+
+            line_cnt++;
+        }
+
+        return true;
+    }
+
   private:
     size_t search_string_range(const std::string& str, char ch, size_t start, size_t end);
 };

include/parameters.h

@@ -12,6 +12,13 @@
 namespace diskann
 {
 
+enum class LabelFormatType :uint8_t
+{
+    String = 0,
+    BitMask = 1,
+    Integer = 2
+};
+
 class IndexWriteParameters
 
 {

include/pq_flash_index.h

@@ -16,6 +16,7 @@
 #include "tsl/robin_map.h"
 #include "tsl/robin_set.h"
 #include "label_bitmask.h"
+#include "integer_label_vector.h"
 
 #define FULL_PRECISION_REORDER_MULTIPLIER 3
 
@@ -47,7 +48,7 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
                                                     const char* labels_filepath, const char* labels_to_medoids_filepath,
                                                     const char* labels_map_filepath, const char* unv_label_filepath,
                                                     const char* seller_filepath,
-                                                    bool load_bitmask_label = false);
+                                                    LabelFormatType label_format_type = LabelFormatType::String);
 #endif
 
     DISKANN_DLLEXPORT void load_cache_list(std::vector<uint32_t> &node_list);
@@ -234,9 +235,12 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
     // filter support
     simple_bitmask_buf _bitmask_buf;
 
+    bool _use_integer_labels = false;
+    integer_label_vector _label_vector;
+
     std::unordered_map<LabelT, std::vector<uint32_t>> _filter_to_medoid_ids;
     bool _use_universal_label = false;
-    LabelT _universal_filter_label;
+    LabelT _universal_filter_label = 0;
     tsl::robin_set<uint32_t> _dummy_pts;
     tsl::robin_set<uint32_t> _has_dummy_pts;
     tsl::robin_map<uint32_t, uint32_t> _dummy_to_real_map;