add a new pre-filter search method and improve serialization performa…

…nce (#8)

Signed-off-by: jinjiabao.jjb <jinjiabao.jjb@antgroup.com>
Co-authored-by: jinjiabao.jjb <jinjiabao.jjb@antgroup.com>

extern/diskann/DiskANN/include/pq.h

@@ -11,6 +11,8 @@
 #define NUM_KMEANS_REPS_PQ 12
 #define MAX_PQ_TRAINING_SET_SIZE 256000
 #define MAX_PQ_CHUNKS 512
+#define NUM_CENTROID 256
+#define ALIGNED_SIZE 256
 
 namespace vsag {
 
@@ -22,7 +24,7 @@ namespace diskann
 {
 class FixedChunkPQTable
 {
-    float *tables = nullptr; // pq_tables = float array of size [256 * ndims]
+    float *tables = nullptr; // pq_tables = float array of size [NUM_CENTROID * ndims]
     uint64_t ndims = 0;      // ndims = true dimension of vectors
     uint64_t n_chunks = 0;
     bool use_rotation = false;
@@ -67,19 +69,20 @@ class FixedChunkPQTable
 
 template <typename T> struct PQScratch
 {
-    float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [256 * NCHUNKS]
+    float *aligned_pqtable_dist_scratch = nullptr; // MUST BE AT LEAST [NUM_CENTROID * NCHUNKS]
     float *aligned_dist_scratch = nullptr;         // MUST BE AT LEAST diskann MAX_DEGREE
     uint8_t *aligned_pq_coord_scratch = nullptr;   // MUST BE AT LEAST  [N_CHUNKS * MAX_DEGREE]
     float *rotated_query = nullptr;
     float *aligned_query_float = nullptr;
 
-    PQScratch(size_t graph_degree, size_t aligned_dim)
+    PQScratch(size_t graph_degree, size_t aligned_dim, size_t pq_chunks)
     {
+
         diskann::alloc_aligned((void **)&aligned_pq_coord_scratch,
-                               (size_t)graph_degree * (size_t)MAX_PQ_CHUNKS * sizeof(uint8_t), 256);
-        diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, 256 * (size_t)MAX_PQ_CHUNKS * sizeof(float),
-                               256);
-        diskann::alloc_aligned((void **)&aligned_dist_scratch, (size_t)graph_degree * sizeof(float), 256);
+                               ROUND_UP(((size_t)graph_degree * (size_t)pq_chunks * sizeof(uint8_t)), ALIGNED_SIZE), ALIGNED_SIZE);
+        diskann::alloc_aligned((void **)&aligned_pqtable_dist_scratch, NUM_CENTROID * (size_t)pq_chunks * sizeof(float),
+                               ALIGNED_SIZE);
+        diskann::alloc_aligned((void **)&aligned_dist_scratch, ROUND_UP((size_t)graph_degree * sizeof(float), ALIGNED_SIZE), ALIGNED_SIZE);
         diskann::alloc_aligned((void **)&aligned_query_float, aligned_dim * sizeof(float), 8 * sizeof(float));
         diskann::alloc_aligned((void **)&rotated_query, aligned_dim * sizeof(float), 8 * sizeof(float));
 

extern/diskann/DiskANN/include/pq_flash_index.h

@@ -24,7 +24,7 @@ namespace diskann
 template <typename T, typename LabelT = uint32_t> class PQFlashIndex
 {
   public:
-    DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr<LocalFileReader> &fileReader, diskann::Metric m, size_t len, bool use_bsa = false);
+    DISKANN_DLLEXPORT PQFlashIndex(std::shared_ptr<LocalFileReader> &fileReader, diskann::Metric m, size_t len, size_t dim, bool use_bsa = false);
     DISKANN_DLLEXPORT ~PQFlashIndex();
 
 #ifdef EXEC_ENV_OLS
@@ -45,7 +45,7 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
 
     DISKANN_DLLEXPORT int load_from_separate_paths(uint32_t num_threads, const char *index_filepath,
                                                    std::stringstream &pivots_stream, std::stringstream &compressed_stream);
-    DISKANN_DLLEXPORT int load_from_separate_paths(uint32_t num_threads, std::stringstream &pivots_stream, std::stringstream &compressed_stream,
+    DISKANN_DLLEXPORT int load_from_separate_paths(std::stringstream &pivots_stream, std::stringstream &compressed_stream,
                                                    std::stringstream &tag_stream);
 
     DISKANN_DLLEXPORT size_t load_graph(std::stringstream &in);
@@ -77,7 +77,7 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
                                               uint64_t *indices, float *distances, const uint64_t beam_width,
                                               std::function<bool(int64_t)> filter,
                                               const uint32_t io_limit, const bool reorder = false,
-                                              QueryStats *stats = nullptr);
+                                              QueryStats *stats = nullptr, bool use_for_range = false);
 
     DISKANN_DLLEXPORT int64_t cached_beam_search_async(const T *query, const uint64_t k_search, const uint64_t l_search,
                                                         uint64_t *indices, float *distances, const uint64_t beam_width,
@@ -103,8 +103,6 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
 
     DISKANN_DLLEXPORT int64_t get_memory_usage();
 
-    void set_sector_size(size_t sector_size);
-
   protected:
     DISKANN_DLLEXPORT void use_medoids_data_as_centroids();
     DISKANN_DLLEXPORT void setup_thread_data(uint64_t nthreads, uint64_t visited_reserve = 4096);
@@ -194,7 +192,6 @@ template <typename T, typename LabelT = uint32_t> class PQFlashIndex
 
     // thread-specific scratch
     ConcurrentQueue<SSDThreadData<T> *> thread_data;
-    size_t sector_size = MAX_N_SECTOR_READS;
     uint64_t max_nthreads;
     bool load_flag = false;
     bool count_visited_nodes = false;

extern/diskann/DiskANN/include/scratch.h

@@ -35,7 +35,7 @@ template <typename T> class InMemQueryScratch
     ~InMemQueryScratch();
     // REFACTOR TODO: move all parameters to a new class.
     InMemQueryScratch(uint32_t search_l, uint32_t indexing_l, uint32_t r, uint32_t maxc, size_t dim, size_t aligned_dim,
-                      size_t alignment_factor, bool init_pq_scratch = false);
+                      size_t alignment_factor, bool init_pq_scratch = false, size_t max_degree = MAX_GRAPH_DEGREE, size_t pq_chunk = MAX_PQ_CHUNKS);
     void resize_for_new_L(uint32_t new_search_l);
     void clear();
 
@@ -153,19 +153,14 @@ template <typename T> class SSDQueryScratch
   public:
     T *coord_scratch = nullptr; // MUST BE AT LEAST [sizeof(T) * data_dim]
 
-    char *sector_scratch = nullptr; // MUST BE AT LEAST [MAX_N_SECTOR_READS * SECTOR_LEN]
     size_t sector_idx = 0;          // index of next [SECTOR_LEN] scratch to use
 
     T *aligned_query_T = nullptr;
 
     PQScratch<T> *_pq_scratch;
 
-    tsl::robin_set<size_t> visited;
-    NeighborPriorityQueue retset;
-    std::vector<Neighbor> full_retset;
 
-    SSDQueryScratch(size_t aligned_dim, size_t visited_reserve);
-    SSDQueryScratch(size_t aligned_dim, size_t visited_reserve, size_t sector_size, size_t sector_len);
+    SSDQueryScratch(size_t max_degree, size_t aligned_dim, size_t pq_chunk);
     ~SSDQueryScratch();
 
     void reset();
@@ -175,9 +170,7 @@ template <typename T> class SSDThreadData
 {
   public:
     SSDQueryScratch<T> scratch;
-
-    SSDThreadData(size_t aligned_dim, size_t visited_reserve);
-    SSDThreadData(size_t aligned_dim, size_t visited_reserve, size_t sector_size, size_t sector_len);
+    SSDThreadData(size_t max_degree, size_t aligned_dim, size_t pq_chunk);
     void clear();
 };
 

extern/diskann/DiskANN/src/pq.cpp

@@ -1920,7 +1920,7 @@ void generate_disk_quantized_data(const T* train_data, size_t train_size, size_t
 
     if (compare_metric == diskann::Metric::INNER_PRODUCT)
         generate_pq_data_from_pivots<float>((const float*)train_data, train_size, train_dim, skip_locs, 256, (uint32_t)disk_pq_dims, disk_pq_pivots,
-                                            disk_pq_compressed_vectors, use_opq, rotate);
+                                            disk_pq_compressed_vectors, use_opq, rotate, use_bsa);
     else
         generate_pq_data_from_pivots<T>(train_data, train_size, train_dim, skip_locs, 256, (uint32_t)disk_pq_dims, disk_pq_pivots,
                                         disk_pq_compressed_vectors, use_opq, rotate, use_bsa);