multithread support - fixing Q Planning stats

Author: Ananya Sutradhar

apps/search_memory_index.cpp

@@ -197,6 +197,10 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
     //query_num = 1;
     double best_recall = 0.0;
 
+    // Declare per-test method counters
+    std::vector<uint32_t> test_num_brutes(Lvec.size(), 0);
+    std::vector<uint32_t> test_num_graphs(Lvec.size(), 0);
+
     // std::cout << "[PERF] About to start search_with_filters. Press Enter to continue..." << std::endl;
     // std::cin.get(); // Wait for user input
 
@@ -227,7 +231,7 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
         //    std::cin.get(); // Wait for user input
         //}
 
-        // Variables for OpenMP reductions (latency, distance comparisons, recall, and timing)
+        // Variables for OpenMP reductions (latency, distance comparisons, recall, timing, and method counts)
         float local_brute_lat = 0.0f;
         float local_graph_lat = 0.0f;
         float local_brute_dist_cmp = 0.0f;
@@ -236,6 +240,8 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
         float local_graph_recall = 0.0f;
         double local_intersection_time = 0.0;
         double local_jaccard_time = 0.0;
+        uint32_t local_num_brutes = 0;
+        uint32_t local_num_graphs = 0;
 
         auto s = std::chrono::high_resolution_clock::now();
         omp_set_num_threads(num_threads);
@@ -244,7 +250,7 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
         std::vector<std::vector<uint32_t>> temp_result_ids(query_num, std::vector<uint32_t>(K_or_L));
         std::vector<std::vector<float>> temp_result_dists(query_num, std::vector<float>(K_or_L));
 
-#pragma omp parallel for schedule(dynamic, 1) reduction(+:local_brute_lat,local_graph_lat,local_brute_dist_cmp,local_graph_dist_cmp,local_brute_recall,local_graph_recall,local_intersection_time,local_jaccard_time)
+#pragma omp parallel for schedule(dynamic, 1) reduction(+:local_brute_lat,local_graph_lat,local_brute_dist_cmp,local_graph_dist_cmp,local_brute_recall,local_graph_recall,local_intersection_time,local_jaccard_time,local_num_brutes,local_num_graphs)
         for (int64_t i = 0; i < (int64_t)query_num; i++)
         {
             int method_used = 0;  // 0 = brute force, 1 = graph search
@@ -266,42 +272,24 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
 
             if (filtered_search && !tags)
             {
-                uint32_t old_b, old_g, old_c;
-                old_b = diskann::num_brutes;
-                old_g = diskann::num_graphs;
-                method_used = 0;
                 std::vector<std::vector<std::string>> raw_filter = query_filters.size() == 1 ? query_filters[0] : query_filters[i];
 
                 // Reset global timing variables before search (thread-unsafe but better than accumulation)
                 // Note: This is not fully thread-safe but provides more accurate per-query measurements
                 diskann::curr_intersection_time = 0.0;
                 diskann::curr_jaccard_time = 0.0;
 
-                auto retval = index->search_with_filters(query + i * query_aligned_dim, raw_filter, K_or_L, L,
+                auto [num_candidates, distance_comparisons, method_type] = index->search_with_filters(query + i * query_aligned_dim, raw_filter, K_or_L, L,
                                                          temp_result_ids[i].data(),
                                                          temp_result_dists[i].data());
 
                 // Capture timing for this query (now represents only this query's timing)
                 delta_intersection = diskann::curr_intersection_time;
                 delta_jaccard = diskann::curr_jaccard_time;
 
-                //// Debug: print num_graphs after search
-                //if (num_threads == 1 && i % 100 == 0) {
-                //    std::cout << "[DEBUG] Query " << i << ": search_with_filters completed, num_graphs = "<< diskann::num_graphs << std::endl;
-                //}
-                ////std::cout << "[DEBUG] After search_with_filters: num_graphs = " << diskann::num_graphs << std::endl;
-
-                if (diskann::num_graphs > old_g) {
-                    method_used = 1;  // Graph search
-                    // Debug the classifications
-                    // if (i < 5) std::cout << "[DEBUG] Query " << i << ": classified as GRAPH search" << std::endl;
-                }
-                else {
-                    method_used = 0;  // Brute force
-                    // Debug the classifications
-                    // if (i < 5) std::cout << "[DEBUG] Query " << i << ": classified as BRUTE FORCE search" << std::endl;
-                }
-                cmp_stats[i] = retval.second;
+                // Method information is now directly available from the search function - no race condition!
+                method_used = method_type;
+                cmp_stats[i] = distance_comparisons;
 //                                filter_match_time[i] = time_to_get_valid*1000000;
                 //                dist_cmp_time[i] = time_to_compare*1000000;
             }
@@ -385,12 +373,20 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
                 local_brute_lat += latency_stats[i];
                 local_brute_dist_cmp += (float)cmp_stats[i];
                 local_brute_recall += cur_recall;
+                local_num_brutes++;
+                if (i < 10) { // Only debug first 10 queries to avoid spam
+
+                }
                 break;
             case 1:
                 query_result_class[test_id][i] = 1;
                 local_graph_lat += latency_stats[i];
                 local_graph_dist_cmp += (float)cmp_stats[i];
                 local_graph_recall += cur_recall;
+                local_num_graphs++;
+                if (i < 10) { // Only debug first 10 queries to avoid spam
+
+                }
                 break;
             }
         }
@@ -412,6 +408,18 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
         brute_recalls[test_id] = local_brute_recall;
         graph_recalls[test_id] = local_graph_recall;
 
+        // DEBUG: Print final values after OpenMP loop
+
+
+
+
+
+
+
+        
+        // Store method counts for this specific test
+        test_num_brutes[test_id] = local_num_brutes;
+        test_num_graphs[test_id] = local_num_graphs;
 
         std::chrono::duration<double> diff = std::chrono::high_resolution_clock::now() - s;
 
@@ -504,19 +512,29 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
             float avg_intersection_time_us = (float)((local_intersection_time * 1000000.0) / (double)query_num);
             float avg_jaccard_time_us = (float)((local_jaccard_time * 1000000.0) / (double)query_num);
 
+            // DEBUG: Print recall calculation values
+
+
+
+
+
+
+
+
+
             std::cout << std::setw(4) << L << std::setw(4) << recall_at << std::setw(8) << displayed_qps << std::setw(18) << avg_cmps
                       << std::setw(20) << (float)mean_latency 
                       << std::setw(20) << (float)latency_999
                       << std::setw(20) << (float)latency_99
                       << std::setw(20) << (float)latency_95
                       << std::setw(10) << (float)recalls[0] 
-                      << std::setw(22) << (float)(brute_dist_cmp[test_id] * 1.0) / (diskann::num_brutes.load() * 1.0) 
-                      << std::setw(22) << (float)(brute_lat[test_id] * 1.0) / (diskann::num_brutes.load() * 1.0) 
-                      << std::setw(20) << (float)(brute_recalls[test_id] * 100.0) / (diskann::num_brutes.load() * recall_at * 1.0) 
+                      << std::setw(22) << (float)(brute_dist_cmp[test_id] * 1.0) / (test_num_brutes[test_id] * 1.0) 
+                      << std::setw(22) << (float)(brute_lat[test_id] * 1.0) / (test_num_brutes[test_id] * 1.0) 
+                      << std::setw(20) << (float)(brute_recalls[test_id] * 100.0) / (test_num_brutes[test_id] * recall_at * 1.0) 
                       << std::setw(22)
-                      << (float)(graph_dist_cmp[test_id] * 1.0) / (diskann::num_graphs.load() * 1.0)
-                      << std::setw(22) << (float)(graph_lat[test_id] * 1.0) / (diskann::num_graphs.load() * 1.0) 
-                      << std::setw(20) << (float)(graph_recalls[test_id] * 100.0) / (diskann::num_graphs.load() * recall_at * 1.0)
+                      << (float)(graph_dist_cmp[test_id] * 1.0) / (test_num_graphs[test_id] * 1.0)
+                      << std::setw(22) << (float)(graph_lat[test_id] * 1.0) / (test_num_graphs[test_id] * 1.0) 
+                      << std::setw(20) << (float)(graph_recalls[test_id] * 100.0) / (test_num_graphs[test_id] * recall_at * 1.0)
                     //   << std::setw(18) << filter_eval_time_us
                     //   << std::setw(18) << penalty_detection_time_us
                     //   << std::setw(18) << core_algo_time_us
@@ -525,8 +543,9 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
                       << std::endl;
         }
     }
-    std::cout << "num_graphs " << diskann::num_graphs.load() << std::endl;
-    std::cout << "num_brutes " << diskann::num_brutes.load() << std::endl;
+    // Show per-L method counts (should be same for each L)
+    std::cout << "num_graphs " << test_num_graphs[0] << std::endl;
+    std::cout << "num_brutes " << test_num_brutes[0] << std::endl;
 
     // Print detailed timing breakdown summary
     if (filtered_search) {
@@ -569,8 +588,8 @@ int search_memory_index(diskann::Metric &metric, const std::string &index_path,
         // std::cout << "Total filter overhead: " << total_filter_overhead << " μs/query" << std::endl;
 
         std::cout << "Breakdown percentage:" << std::endl;
-        std::cout << "  Graph searches: " << (100.0 * diskann::num_graphs.load()) / query_num << "%" << std::endl;
-        std::cout << "  Brute force searches: " << (100.0 * diskann::num_brutes.load()) / query_num << "%" << std::endl;
+        std::cout << "  Graph searches: " << (100.0 * test_num_graphs[0]) / query_num << "%" << std::endl;
+        std::cout << "  Brute force searches: " << (100.0 * test_num_brutes[0]) / query_num << "%" << std::endl;
         std::cout << "=================================" << std::endl;
     }
 

include/abstract_index.h

@@ -75,11 +75,11 @@ class AbstractIndex
     // Filter support search
     // IndexType is either uint32_t or uint64_t
     template <typename IndexType>
-    std::pair<uint32_t, uint32_t> search_with_filters(const DataType &query, const std::string &raw_label,
+    std::tuple<uint32_t, uint32_t, uint32_t> search_with_filters(const DataType &query, const std::string &raw_label,
                                                       const size_t K, const uint32_t L, IndexType *indices,
                                                       float *distances);
     template <typename IndexType>
-    std::pair<uint32_t, uint32_t> search_with_filters(const DataType &query, const std::vector<std::vector<std::string>> &raw_label,
+    std::tuple<uint32_t, uint32_t, uint32_t> search_with_filters(const DataType &query, const std::vector<std::vector<std::string>> &raw_label,
                                                       const size_t K, const uint32_t L, IndexType *indices,
                                                       float *distances);
 
@@ -114,7 +114,7 @@ class AbstractIndex
     virtual void _build(const DataType &data, const size_t num_points_to_load, TagVector &tags) = 0;
     virtual std::pair<uint32_t, uint32_t> _search(const DataType &query, const size_t K, const uint32_t L,
                                                   std::any &indices, float *distances = nullptr) = 0;
-    virtual std::pair<uint32_t, uint32_t> _search_with_filters(const DataType &query,
+    virtual std::tuple<uint32_t, uint32_t, uint32_t> _search_with_filters(const DataType &query,
                                                                const std::vector<std::vector<std::string>> &filter_label,
                                                                const size_t K, const uint32_t L, std::any &indices,
                                                                float *distances) = 0;

include/index.h

@@ -190,7 +190,7 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
                                               const std::string filter_label = "");
 
     template <typename IndexType>
-    DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> search_with_filters(const T *query,
+    DISKANN_DLLEXPORT std::tuple<uint32_t, uint32_t, uint32_t> search_with_filters(const T *query,
                                                                         const std::vector<std::vector<LabelT>> &filter_label,
                                                                         const size_t K, const uint32_t L,
                                                                         IndexType *indices, float *distances);
@@ -256,7 +256,7 @@ template <typename T, typename TagT = uint32_t, typename LabelT = uint32_t> clas
 
     virtual std::pair<uint32_t, uint32_t> _search(const DataType &query, const size_t K, const uint32_t L,
                                                   std::any &indices, float *distances = nullptr) override;
-    virtual std::pair<uint32_t, uint32_t> _search_with_filters(const DataType &query,
+    virtual std::tuple<uint32_t, uint32_t, uint32_t> _search_with_filters(const DataType &query,
                                                                const std::vector<std::vector<std::string>> &filter_label_raw,
                                                                const size_t K, const uint32_t L, std::any &indices,
                                                                float *distances) override;

src/abstract_index.cpp

@@ -34,7 +34,7 @@ size_t AbstractIndex::search_with_tags(const data_type *query, const uint64_t K,
 }
 
 template <typename IndexType>
-std::pair<uint32_t, uint32_t> AbstractIndex::search_with_filters(const DataType &query, const std::string &raw_label,
+std::tuple<uint32_t, uint32_t, uint32_t> AbstractIndex::search_with_filters(const DataType &query, const std::string &raw_label,
                                                                  const size_t K, const uint32_t L, IndexType *indices,
                                                                  float *distances)
 {
@@ -47,7 +47,7 @@ std::pair<uint32_t, uint32_t> AbstractIndex::search_with_filters(const DataType
 }
 
 template <typename IndexType>
-std::pair<uint32_t, uint32_t> AbstractIndex::search_with_filters(const DataType &query,
+std::tuple<uint32_t, uint32_t, uint32_t> AbstractIndex::search_with_filters(const DataType &query,
                                                                  const std::vector<std::vector<std::string>> &raw_label,
                                                                  const size_t K, const uint32_t L, IndexType *indices,
                                                                  float *distances)
@@ -169,18 +169,18 @@ template DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> AbstractIndex::search<u
 template DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> AbstractIndex::search<int8_t, uint64_t>(
     const int8_t *query, const size_t K, const uint32_t L, uint64_t *indices, float *distances);
 
-template DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> AbstractIndex::search_with_filters<uint32_t>(
+template DISKANN_DLLEXPORT std::tuple<uint32_t, uint32_t, uint32_t> AbstractIndex::search_with_filters<uint32_t>(
     const DataType &query, const std::string &raw_label, const size_t K, const uint32_t L, uint32_t *indices,
     float *distances);
 
-template DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> AbstractIndex::search_with_filters<uint64_t>(
+template DISKANN_DLLEXPORT std::tuple<uint32_t, uint32_t, uint32_t> AbstractIndex::search_with_filters<uint64_t>(
     const DataType &query, const std::string &raw_label, const size_t K, const uint32_t L, uint64_t *indices,
     float *distances);
-template DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> AbstractIndex::search_with_filters<uint32_t>(
+template DISKANN_DLLEXPORT std::tuple<uint32_t, uint32_t, uint32_t> AbstractIndex::search_with_filters<uint32_t>(
     const DataType &query, const std::vector<std::vector<std::string>> &raw_label, const size_t K, const uint32_t L,
     uint32_t *indices, float *distances);
 
-template DISKANN_DLLEXPORT std::pair<uint32_t, uint32_t> AbstractIndex::search_with_filters<uint64_t>(
+template DISKANN_DLLEXPORT std::tuple<uint32_t, uint32_t, uint32_t> AbstractIndex::search_with_filters<uint64_t>(
     const DataType &query, const std::vector<std::vector<std::string>> &raw_label, const size_t K, const uint32_t L,
     uint64_t *indices, float *distances);