Remove restriction of caching only 10 % nodes in disk search

apps/search_disk_index.cpp

@@ -179,7 +179,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
     std::string recall_string = "Recall@" + std::to_string(recall_at);
     diskann::cout << std::setw(6) << "L" << std::setw(12) << "Beamwidth" << std::setw(16) << "QPS" << std::setw(16)
                   << "Mean Latency" << std::setw(16) << "99.9 Latency" << std::setw(16) << "Mean IOs" << std::setw(16)
-                  << "Mean IO (us)" << std::setw(16) << "CPU (s)";
+                  << "Mean IO (us)" << std::setw(16) << "CPU (s)" << std::setw(16)<< "PQ Training(s)";
     if (calc_recall_flag)
     {
         diskann::cout << std::setw(16) << recall_string << std::endl;
@@ -272,6 +272,8 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
 
         auto mean_io_us = diskann::get_mean_stats<float>(stats, query_num,
                                                          [](const diskann::QueryStats &stats) { return stats.io_us; });
+        auto mean_pq_training_us = diskann::get_mean_stats<float>(stats, query_num,
+                                                         [](const diskann::QueryStats &stats) { return stats.pq_training_us; });
 
         double recall = 0;
         if (calc_recall_flag)
@@ -283,7 +285,7 @@ int search_disk_index(diskann::Metric &metric, const std::string &index_path_pre
 
         diskann::cout << std::setw(6) << L << std::setw(12) << optimized_beamwidth << std::setw(16) << qps
                       << std::setw(16) << mean_latency << std::setw(16) << latency_999 << std::setw(16) << mean_ios
-                      << std::setw(16) << mean_io_us << std::setw(16) << mean_cpuus;
+                      << std::setw(16) << mean_io_us << std::setw(16) << mean_cpuus << std::setw(16) << mean_pq_training_us;
         if (calc_recall_flag)
         {
             diskann::cout << std::setw(16) << recall << std::endl;

include/percentile_stats.h

@@ -23,6 +23,7 @@ struct QueryStats
     float total_us = 0; // total time to process query in micros
     float io_us = 0;    // total time spent in IO
     float cpu_us = 0;   // total time spent in CPU
+    float pq_training_us = 0; // total time spent in PQ training
 
     unsigned n_4k = 0;         // # of 4kB reads
     unsigned n_8k = 0;         // # of 8kB reads

python/include/builder.h

@@ -18,10 +18,9 @@ void build_disk_index(diskann::Metric metric, const std::string &data_file_path,
 
 template <typename DT, typename TagT = DynamicIdType, typename LabelT = filterT>
 void build_memory_index(diskann::Metric metric, const std::string &vector_bin_path,
-                           const std::string &index_output_path, uint32_t graph_degree, uint32_t complexity,
-                           float alpha, uint32_t num_threads, bool use_pq_build,
-                           size_t num_pq_bytes, bool use_opq, bool use_tags = false,
-                           const std::string& filter_labels_file = "", const std::string& universal_label = "",
-                           uint32_t filter_complexity = 0);
+                        const std::string &index_output_path, uint32_t graph_degree, uint32_t complexity, float alpha,
+                        uint32_t num_threads, bool use_pq_build, size_t num_pq_bytes, bool use_opq,
+                        bool use_tags = false, const std::string &filter_labels_file = "",
+                        const std::string &universal_label = "", uint32_t filter_complexity = 0);
 
-}
+} // namespace diskannpy

python/include/dynamic_memory_index.h

@@ -18,8 +18,7 @@ namespace py = pybind11;
 namespace diskannpy
 {
 
-template <typename DT>
-class DynamicMemoryIndex
+template <typename DT> class DynamicMemoryIndex
 {
   public:
     DynamicMemoryIndex(diskann::Metric m, size_t dimensions, size_t max_vectors, uint32_t complexity,
@@ -31,19 +30,18 @@ class DynamicMemoryIndex
     void load(const std::string &index_path);
     int insert(const py::array_t<DT, py::array::c_style | py::array::forcecast> &vector, DynamicIdType id);
     py::array_t<int> batch_insert(py::array_t<DT, py::array::c_style | py::array::forcecast> &vectors,
-                                  py::array_t<DynamicIdType, py::array::c_style | py::array::forcecast> &ids, int32_t num_inserts,
-                                  int num_threads = 0);
+                                  py::array_t<DynamicIdType, py::array::c_style | py::array::forcecast> &ids,
+                                  int32_t num_inserts, int num_threads = 0);
     int mark_deleted(DynamicIdType id);
     void save(const std::string &save_path, bool compact_before_save = false);
-    NeighborsAndDistances<DynamicIdType> search(py::array_t<DT, py::array::c_style | py::array::forcecast> &query, uint64_t knn,
-                                      uint64_t complexity);
-    NeighborsAndDistances<DynamicIdType> batch_search(py::array_t<DT, py::array::c_style | py::array::forcecast> &queries,
-                                            uint64_t num_queries, uint64_t knn, uint64_t complexity,
-                                            uint32_t num_threads);
+    NeighborsAndDistances<DynamicIdType> search(py::array_t<DT, py::array::c_style | py::array::forcecast> &query,
+                                                uint64_t knn, uint64_t complexity);
+    NeighborsAndDistances<DynamicIdType> batch_search(
+        py::array_t<DT, py::array::c_style | py::array::forcecast> &queries, uint64_t num_queries, uint64_t knn,
+        uint64_t complexity, uint32_t num_threads);
     void consolidate_delete();
     size_t num_points();
 
-
   private:
     const uint32_t _initial_search_complexity;
     const diskann::IndexWriteParameters _write_parameters;

src/pq_flash_index.cpp

@@ -355,13 +355,12 @@ void PQFlashIndex<T, LabelT>::cache_bfs_levels(uint64_t num_nodes_to_cache, std:
 
     tsl::robin_set<uint32_t> node_set;
 
-    // Do not cache more than 10% of the nodes in the index
-    uint64_t tenp_nodes = (uint64_t)(std::round(this->_num_points * 0.1));
-    if (num_nodes_to_cache > tenp_nodes)
+    // If num_nodes_to_cache is more than total_nodes then reduce num_nodes_to_cache to total_nodes.
+    if (num_nodes_to_cache > this->_num_points)
     {
-        diskann::cout << "Reducing nodes to cache from: " << num_nodes_to_cache << " to: " << tenp_nodes
-                      << "(10 percent of total nodes:" << this->_num_points << ")" << std::endl;
-        num_nodes_to_cache = tenp_nodes == 0 ? 1 : tenp_nodes;
+        diskann::cout << "Reducing nodes to cache from: " << num_nodes_to_cache
+                      << " to total nodes:" << this->_num_points << std::endl;
+        num_nodes_to_cache = this->_num_points;
     }
     diskann::cout << "Caching " << num_nodes_to_cache << "..." << std::endl;
 
@@ -1276,7 +1275,7 @@ void PQFlashIndex<T, LabelT>::cached_beam_search(const T *query1, const uint64_t
     if (beam_width > num_sector_per_nodes * defaults::MAX_N_SECTOR_READS)
         throw ANNException("Beamwidth can not be higher than defaults::MAX_N_SECTOR_READS", -1, __FUNCSIG__, __FILE__,
                            __LINE__);
-
+    Timer query_timer, io_timer, cpu_timer, pq_training_timer;
     ScratchStoreManager<SSDThreadData<T>> manager(this->_thread_data);
     auto data = manager.scratch_space();
     IOContext &ctx = data->ctx;
@@ -1334,11 +1333,18 @@ void PQFlashIndex<T, LabelT>::cached_beam_search(const T *query1, const uint64_t
     const uint64_t num_sectors_per_node =
         _nnodes_per_sector > 0 ? 1 : DIV_ROUND_UP(_max_node_len, defaults::SECTOR_LEN);
 
+    cpu_timer.reset();
+    pq_training_timer.reset();
     // query <-> PQ chunk centers distances
     _pq_table.preprocess_query(query_rotated); // center the query and rotate if
                                                // we have a rotation matrix
     float *pq_dists = pq_query_scratch->aligned_pqtable_dist_scratch;
     _pq_table.populate_chunk_distances(query_rotated, pq_dists);
+    if (stats != nullptr)
+    {
+        stats->pq_training_us = (float)pq_training_timer.elapsed();
+        stats->cpu_us += (float)cpu_timer.elapsed();
+    }
 
     // query <-> neighbor list
     float *dist_scratch = pq_query_scratch->aligned_dist_scratch;
@@ -1350,7 +1356,6 @@ void PQFlashIndex<T, LabelT>::cached_beam_search(const T *query1, const uint64_t
         diskann::aggregate_coords(ids, n_ids, this->data, this->_n_chunks, pq_coord_scratch);
         diskann::pq_dist_lookup(pq_coord_scratch, n_ids, this->_n_chunks, pq_dists, dists_out);
     };
-    Timer query_timer, io_timer, cpu_timer;
 
     tsl::robin_set<uint64_t> &visited = query_scratch->visited;
     NeighborPriorityQueue &retset = query_scratch->retset;