Fix errors

Author: pianand-p1p13

apps/build_memory_index.cpp

@@ -24,10 +24,11 @@ namespace po = boost::program_options;
 
 int main(int argc, char **argv)
 {
-    std::string data_type, dist_fn, data_path, index_path_prefix, label_file, universal_label, label_type;
-    uint32_t num_threads, R, L, Lf, build_PQ_bytes;
+    std::string data_type, dist_fn, data_path, index_path_prefix, label_file, universal_label, label_type, seller_file, attribute_file;
+    uint32_t num_threads, R, L, Lf, build_PQ_bytes, num_diverse_build;
     float alpha;
-    bool use_pq_build, use_opq;
+    bool use_pq_build, use_opq, diverse_index = false, attribute_diversity = false;
+    
 
     po::options_description desc{
         program_options_utils::make_program_description("build_memory_index", "Build a memory-based DiskANN index.")};
@@ -70,6 +71,12 @@ int main(int argc, char **argv)
                                        program_options_utils::FILTERED_LBUILD);
         optional_configs.add_options()("label_type", po::value<std::string>(&label_type)->default_value("uint"),
                                        program_options_utils::LABEL_TYPE_DESCRIPTION);
+        optional_configs.add_options()("seller_file", po::value<std::string>(&seller_file)->default_value(""),
+                                       program_options_utils::DIVERSITY_FILE);
+        optional_configs.add_options()("attribute_file", po::value<std::string>(&attribute_file)->default_value(""),
+                                       program_options_utils::ATTRIBUTE_DIVERSITY_FILE);
+        optional_configs.add_options()("NumDiverse", po::value<uint32_t>(&num_diverse_build)->default_value(1),
+                                       program_options_utils::NUM_DIVERSE);
 
         // Merge required and optional parameters
         desc.add(required_configs).add(optional_configs);
@@ -112,19 +119,33 @@ int main(int argc, char **argv)
         return -1;
     }
 
+    if(seller_file != "")
+        diverse_index = true;
+
+
+    if(attribute_file != ""){
+        diverse_index = true;
+        attribute_diversity = true;
+    }
     try
     {
         diskann::cout << "Starting index build with R: " << R << "  Lbuild: " << L << "  alpha: " << alpha
                       << "  #threads: " << num_threads << std::endl;
 
         size_t data_num, data_dim;
         diskann::get_bin_metadata(data_path, data_num, data_dim);
+        std::cout<<"Num diverse build: " << num_diverse_build << std::endl;
 
         auto index_build_params = diskann::IndexWriteParametersBuilder(L, R)
                                       .with_filter_list_size(Lf)
                                       .with_alpha(alpha)
                                       .with_saturate_graph(false)
                                       .with_num_threads(num_threads)
+                                      .with_diverse_index(diverse_index)
+                                      .with_seller_file(seller_file)
+                                      .with_num_diverse_build(num_diverse_build)
+                                      .with_attribute_diversity(attribute_diversity)
+                                      .with_attribute_file(attribute_file)
                                       .build();
 
         auto filter_params = diskann::IndexFilterParamsBuilder()

apps/search_memory_index.cpp

@@ -153,8 +153,9 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
             continue;
         }
 
-        query_result_ids[test_id].resize(recall_at * query_num);
-        query_result_dists[test_id].resize(recall_at * query_num);
+        query_result_ids[test_id].resize(recall_at * query_num, std::numeric_limits<uint32_t>::max());
+        query_result_dists[test_id].resize(recall_at * query_num, std::numeric_limits<float>::max());
+
         std::vector<T *> res = std::vector<T *>();
 
         auto s = std::chrono::high_resolution_clock::now();
@@ -165,12 +166,7 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
             auto qs = std::chrono::high_resolution_clock::now();
             if (filtered_search && !tags)
             {
-                std::string raw_filter = query_filters.size() == 1 ? query_filters[0] : query_filters[i];
-
-                auto retval = index->search_with_filters(query + i * query_aligned_dim, raw_filter, recall_at, L,
-                                                         query_result_ids[test_id].data() + i * recall_at,
-                                                         query_result_dists[test_id].data() + i * recall_at);
-                cmp_stats[i] = retval.second;
+                ;
             }
             else if (metric == diskann::FAST_L2)
             {
@@ -198,11 +194,17 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
                 }
             }
             else
-            {
+            {   std::vector<uint32_t> results(L,std::numeric_limits<uint32_t>::max());
+                std::vector<float> dists(L,std::numeric_limits<float>::max());
+                
                 cmp_stats[i] = index
                                    ->search(query + i * query_aligned_dim, recall_at, L,
-                                            query_result_ids[test_id].data() + i * recall_at)
+                                            results.data(), dists.data())
                                    .second;
+                for (uint32_t ctr = 0; ctr < std::min(results.size(),(uint64_t)recall_at); ctr++) {
+                        query_result_ids[test_id][recall_at * i + ctr] = results[ctr];
+                        query_result_dists[test_id][recall_at * i + ctr] = dists[ctr];
+                }
             }
             auto qe = std::chrono::high_resolution_clock::now();
             std::chrono::duration<double> diff = qe - qs;

apps/utils/tsv_to_bin.cpp

@@ -64,9 +64,21 @@ void block_convert_uint8(std::ifstream &reader, std::ofstream &writer, size_t np
     delete[] read_buf;
 }
 
+size_t count_lines(const std::string &filename)
+{
+    std::ifstream file(filename);
+    size_t lines = 0;
+    std::string line;
+    while (std::getline(file, line))
+    {
+        ++lines;
+    }
+    return lines;
+}
+
 int main(int argc, char **argv)
 {
-    if (argc != 6)
+    if (argc != 5)
     {
         std::cout << argv[0]
                   << "<float/int8/uint8> input_filename.tsv output_filename.bin "
@@ -82,7 +94,7 @@ int main(int argc, char **argv)
     }
 
     size_t ndims = atoi(argv[4]);
-    size_t npts = atoi(argv[5]);
+    size_t npts = count_lines(argv[2]);
 
     std::ifstream reader(argv[2], std::ios::binary | std::ios::ate);
     //  size_t          fsize = reader.tellg();

include/index.h

@@ -402,7 +402,7 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
     uint32_t _num_diverse_build = 1;
     uint32_t _max_L_per_seller = 0;
     std::vector<uint32_t> _location_to_seller;
-    std::vector<std::vector<uint32_t>> _location_to_attributes;
+    std::vector<std::vector<std::string>> _location_to_attributes;
     uint32_t _num_unique_sellers = 0;
     std::string _seller_file;
     std::string _attribute_file;

include/program_options_utils.hpp

@@ -77,5 +77,11 @@ const char *UNIVERSAL_LABEL =
     "in the labels file instead of listing all labels for a node.  DiskANN will not automatically assign a "
     "universal label to a node.";
 const char *FILTERED_LBUILD = "Build complexity for filtered points, higher value results in better graphs";
+const char *DIVERSE_INDEX = "Build Diverse Index";
+const char *ATTRIBUTE_DIVERSE_INDEX = "Build Attribute Diverse Index";
+const char *NUM_DIVERSE = "Number of diverse edges needed per node in each local region";
+const char *DIVERSITY_FILE = "Seller diversity file for diverse index";
+const char *ATTRIBUTE_DIVERSITY_FILE = "Attribute diversity file for diverse index";
+
 
 } // namespace program_options_utils

src/index.cpp

@@ -866,7 +866,7 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::iterate_to_fixed_point(
 
     std::vector<Neighbor> &expanded_nodes = scratch->pool();
 
-    if (!_diverse_index)
+    if (!_diverse_index || _attribute_diversity)
     {
         NeighborPriorityQueue& best_L_nodes_ref = scratch->best_l_nodes();
         best_L_nodes_ref.reserve(Lsize);
@@ -1233,16 +1233,30 @@ void Index<T, TagT, LabelT>::prune_search_result(int location, std::vector<uint3
 // #include <execution>
 
 
-double attribute_distance(std::vector<uint32_t> &a, std::vector <uint32_t> &b)
-{
-    std::vector<int> result(a.size());
+// double attribute_distance(std::vector<std::string> &a, std::vector <std::string> &b)
+// {
+//     std::vector<int> result(a.size());
 
-    std::transform(a.begin(), a.end(), b.begin(), result.begin(),
-                   [](int x, int y) { return x == y ? 1 : 0; });
+//     std::transform(a.begin(), a.end(), b.begin(), result.begin(),
+//                    [](std::string x, std::string  y) { return x == y ? 1 : 0; });
+
+//     int sum = std::accumulate(result.begin(), result.end(), 0);
+//     double distance = sum * 1.0 / result.size();
+//     return 1 - distance;
+// }
+double attribute_distance(const std::vector<std::string> &a, const std::vector<std::string> &b)
+{
+    // Use the smaller size to avoid out-of-bounds issues
+    size_t n = std::min(a.size(), b.size());
+    if(n == 0)
+        return 0.0; // or maybe 1.0 if you want full "distance" when both are empty
 
-    int sum = std::accumulate(result.begin(), result.end(), 0);
-    double distance = sum * 1.0 / result.size();
-    return 1 - distance;
+    int matches = 0;
+    for(size_t i = 0; i < n; ++i)
+        if(a[i] == b[i]) matches++;
+
+    double similarity = static_cast<double>(matches) / n;
+    return 1.0 - similarity;
 }
 
 template <typename T, typename TagT, typename LabelT>
@@ -1295,10 +1309,13 @@ void Index<T, TagT, LabelT>::occlude_list(const uint32_t location, std::vector<N
                 {
                     continue;
                 }
-                auto iter_seller = _location_to_seller[iter->id];
-                if (blockers[cur_index].find(iter_seller) != blockers[cur_index].end() && !_attribute_diversity)
-                {
-                    continue;
+
+                if(!_attribute_diversity){
+                    auto iter_seller = _location_to_seller[iter->id];
+                    if (blockers[cur_index].find(iter_seller) != blockers[cur_index].end() && !_attribute_diversity)
+                    {
+                        continue;
+                    }
                 }
 
                 if (blockers[cur_index].find((uint32_t)cur_index) != blockers[cur_index].end() && _attribute_diversity)
@@ -1333,11 +1350,13 @@ void Index<T, TagT, LabelT>::occlude_list(const uint32_t location, std::vector<N
                     {
                         continue;
                     }
-
-                    auto iter2_seller = _location_to_seller[iter2->id];
-                    if (blockers[t].find(iter2_seller) != blockers[t].end())
-                    {
-                        continue;
+                    
+                    if(!_attribute_diversity){
+                        auto iter2_seller = _location_to_seller[iter2->id];
+                        if (blockers[t].find(iter2_seller) != blockers[t].end())
+                        {
+                            continue;
+                        }
                     }
                 }
 
@@ -1369,6 +1388,7 @@ void Index<T, TagT, LabelT>::occlude_list(const uint32_t location, std::vector<N
                             blockers[t].insert(iter_seller);
                         }
                         else{
+                            //double attr_dist = 0.6;
                             double attr_dist = attribute_distance(_location_to_attributes[iter2->id], _location_to_attributes[iter->id]);
                             if(attr_dist > 0.5){ // attribute distance threshold
                                 blockers[t].insert(iter->id);
@@ -1388,7 +1408,7 @@ void Index<T, TagT, LabelT>::occlude_list(const uint32_t location, std::vector<N
                     if (y > cur_alpha * x)
                     {
                         occlude_factor[t] = std::max(occlude_factor[t], eps);
-                        if (_diverse_index)
+                        if (_diverse_index && !_attribute_diversity)
                         {
                             auto iter_seller = _location_to_seller[iter->id];
                             blockers[t].insert(iter_seller);
@@ -1553,7 +1573,9 @@ template <typename T, typename TagT, typename LabelT> void Index<T, TagT, LabelT
 
     diskann::Timer link_timer;
 
+try{
 #pragma omp parallel for schedule(dynamic, 2048)
+
     for (int64_t node_ctr = 0; node_ctr < (int64_t)(visit_order.size()); node_ctr++)
     {
         auto node = visit_order[node_ctr];
@@ -1567,8 +1589,12 @@ template <typename T, typename TagT, typename LabelT> void Index<T, TagT, LabelT
             search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch, true, _filterIndexingQueueSize);
         }
         else
-        {
-            search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch);
+        {   try{
+                search_for_point_and_prune(node, _indexingQueueSize, pruned_list, scratch);
+        
+            } catch (const std::exception &e) {
+                diskann::cerr << "Exception caught during link for node " << node << " : " << e.what() << std::endl;
+            }
         }
         assert(pruned_list.size() > 0);
 
@@ -1586,6 +1612,8 @@ template <typename T, typename TagT, typename LabelT> void Index<T, TagT, LabelT
             diskann::cout << "\r" << (100.0 * node_ctr) / (visit_order.size()) << "% of index build completed."
                           << std::flush;
         }
+    }} catch (const std::exception &e) {
+        diskann::cerr << "Exception caught during link: " << e.what() << std::endl;
     }
 
     if (_nd > 0)
@@ -2272,14 +2300,14 @@ void Index<T, TagT, LabelT>::parse_attribute_file(const std::string &label_file,
         std::istringstream iss(line);
         getline(iss, token, '\t');
         std::istringstream new_iss(token);
-        std::vector<uint32_t> attributes;
+        std::vector<std::string> attributes;
 
         while (getline(new_iss, token, ','))
         {
             token.erase(std::remove(token.begin(), token.end(), '\n'), token.end());
             token.erase(std::remove(token.begin(), token.end(), '\r'), token.end());
-            uint32_t token_as_num = (uint32_t)std::stoul(token);
-            attributes.push_back(token_as_num);
+            //uint32_t token_as_num = (uint32_t)std::stoul(token);
+            attributes.push_back(token);
         }
 
         _location_to_attributes.push_back(attributes);
@@ -2465,7 +2493,7 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::search(const T *query, con
     auto retval = iterate_to_fixed_point(scratch, L, init_ids, false, unused_filter_label, true, maxLperSeller);
 
     NeighborPriorityQueueBase* best_L_nodes;
-    if (!_diverse_index || maxLperSeller == 0) {
+    if (!_diverse_index || maxLperSeller == 0 || _attribute_diversity) {
         best_L_nodes = &(scratch->best_l_nodes());
     }
     else {
@@ -2575,7 +2603,7 @@ std::pair<uint32_t, uint32_t> Index<T, TagT, LabelT>::search_with_filters(const
     auto retval = iterate_to_fixed_point(scratch, L, init_ids, true, filter_vec, true, maxLperSeller);
 
     NeighborPriorityQueueBase* best_L_nodes;
-    if (!_diverse_index) {
+    if (!_diverse_index || _attribute_diversity) {
         best_L_nodes = &(scratch->best_l_nodes());
     }
     else {