attribute diversity

Author: Sanhaoji2

AnyBuildLogs/latest.txt

@@ -1 +1 @@
-20260116-140907-2cca51ce
+20260127-171934-6eb71ae2

include/defaults.h

@@ -37,6 +37,8 @@ const bool NUM_DIVERSE_BUILD = 1;
 
 const bool REORDER_INDEX = false;
 const uint32_t REORDER_DIM = 0;
+const bool ATTRIBUTE_DIVERSITY = false;
+const float ATTR_DIST_THRESHOLD = 0.2f;
 
 } // namespace defaults
 } // namespace diskann

include/index.h

@@ -267,8 +267,12 @@ 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, size_t& total_labels);
-    void parse_seller_file(const std::string& label_file, size_t& num_pts_labels);
+    template <typename ValueT>
+    void parse_integer_string_file(const std::string &file_path, size_t &num_points, size_t& total_values,
+                                   std::vector<std::vector<ValueT>>& location_to_values, 
+                                   tsl::robin_set<ValueT>* unique_values = nullptr);
+    void parse_seller_file(const std::string& label_file, size_t& num_pts_labels, 
+                           std::vector<uint32_t>& location_to_seller, uint32_t& num_unique_sellers);
 
     void convert_pts_label_to_bitmask(std::vector<std::vector<LabelT>>& pts_to_labels, simple_bitmask_buf& bitmask_buf, size_t num_labels);
 
@@ -357,6 +361,8 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     void initialize_query_scratch(uint32_t num_threads, uint32_t search_l, uint32_t indexing_l, uint32_t r,
                                   uint32_t maxc, size_t dim, size_t bitmask_size = 0);
 
+    double attribute_distance(const std::vector<uint32_t> &a, const std::vector<uint32_t> &b);
+
     // Do not call without acquiring appropriate locks
     // call public member functions save and load to invoke these.
     DISKANN_DLLEXPORT size_t save_graph(std::string filename);
@@ -426,6 +432,10 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     std::vector<uint32_t> _location_to_seller;
     uint32_t _num_unique_sellers = 0;
     std::string _seller_file;
+    bool _attribute_diversity = false;
+    float _attr_dist_threshold = 0.2f;
+    std::string _attribute_file;
+    std::vector<std::vector<std::uint32_t>> _location_to_attributes;
 
     bool _use_universal_label = false;
     LabelT _universal_label = 0;

include/parameters.h

@@ -32,13 +32,18 @@ class IndexWriteParameters
     const bool diverse_index;
     const std::string seller_file;
     const uint32_t num_diverse_build;
+    const bool attribute_diversity;
+    const std::string attribute_file;
+    const float attr_dist_threshold;
 
     IndexWriteParameters(const uint32_t search_list_size, const uint32_t max_degree, const bool saturate_graph,
                          const uint32_t max_occlusion_size, const float alpha, const uint32_t num_threads,
-                         const uint32_t filter_list_size, bool diverse_index, const std::string& seller_file, uint32_t num_diverse_build)
+                         const uint32_t filter_list_size, bool diverse_index, const std::string& seller_file, uint32_t num_diverse_build,
+                         bool attribute_diversity, const std::string& attribute_file, float attr_dist_threshold)
         : search_list_size(search_list_size), max_degree(max_degree), saturate_graph(saturate_graph),
           max_occlusion_size(max_occlusion_size), alpha(alpha), num_threads(num_threads),
-          filter_list_size(filter_list_size), diverse_index(diverse_index), seller_file(seller_file), num_diverse_build(num_diverse_build)
+          filter_list_size(filter_list_size), diverse_index(diverse_index), seller_file(seller_file), num_diverse_build(num_diverse_build),
+          attribute_diversity(attribute_diversity), attribute_file(attribute_file), attr_dist_threshold(attr_dist_threshold)
     {
     }
 
@@ -100,6 +105,24 @@ class IndexWriteParametersBuilder
         return *this;
     }
 
+    IndexWriteParametersBuilder& with_attribute_diversity(const bool attribute_diversity)
+    {
+        _attribute_diversity = attribute_diversity;
+        return *this;
+    }
+
+    IndexWriteParametersBuilder& with_attr_dist_threshold(const float attr_dist_threshold)
+    {
+        _attr_dist_threshold = attr_dist_threshold;
+        return *this;
+    }
+
+    IndexWriteParametersBuilder& with_attribute_file(const std::string attribute_file)
+    {
+        _attribute_file = attribute_file;
+        return *this;
+    }
+
     IndexWriteParametersBuilder &with_alpha(const float alpha)
     {
         _alpha = alpha;
@@ -121,13 +144,16 @@ class IndexWriteParametersBuilder
     IndexWriteParameters build() const
     {
         return IndexWriteParameters(_search_list_size, _max_degree, _saturate_graph, _max_occlusion_size, _alpha,
-                                    _num_threads, _filter_list_size, _diverse_index, _seller_file, _num_diverse_build);
+                                    _num_threads, _filter_list_size, _diverse_index, _seller_file, _num_diverse_build,
+                                    _attribute_diversity, _attribute_file, _attr_dist_threshold);
     }
 
     IndexWriteParametersBuilder(const IndexWriteParameters &wp)
         : _search_list_size(wp.search_list_size), _max_degree(wp.max_degree),
           _max_occlusion_size(wp.max_occlusion_size), _saturate_graph(wp.saturate_graph), _alpha(wp.alpha),
-          _filter_list_size(wp.filter_list_size)
+          _num_threads(wp.num_threads), _filter_list_size(wp.filter_list_size), _diverse_index(wp.diverse_index),
+          _seller_file(wp.seller_file), _num_diverse_build(wp.num_diverse_build), _attribute_diversity(wp.attribute_diversity),
+          _attribute_file(wp.attribute_file), _attr_dist_threshold(wp.attr_dist_threshold)
     {
     }
     IndexWriteParametersBuilder(const IndexWriteParametersBuilder &) = delete;
@@ -143,7 +169,10 @@ class IndexWriteParametersBuilder
     uint32_t _filter_list_size{defaults::FILTER_LIST_SIZE};
     bool _diverse_index{ defaults::DIVERSE_INDEX };
     std::string _seller_file{ defaults::EMPTY_STRING };
-    uint32_t _num_diverse_build{ defaults::NUM_DIVERSE_BUILD };
+    uint32_t _num_diverse_build{ defaults::NUM_DIVERSE_BUILD }; 
+    bool _attribute_diversity{ defaults::ATTRIBUTE_DIVERSITY };
+    std::string _attribute_file{ defaults::EMPTY_STRING };
+    float _attr_dist_threshold{ defaults::ATTR_DIST_THRESHOLD };
 };
 
 struct IndexLoadParams