Create a data store abstraction

.clang-format

@@ -1,5 +1,5 @@
 ---
-BasedOnStyle:  Microsoft
+BasedOnStyle:  Google
 ---
 Language: Cpp
 SortIncludes: false

.github/workflows/pr-test.yml

@@ -172,7 +172,7 @@ jobs:
     - name:  test a streaming index
       if: success() || failure()
       run: |
-        ${{ env.diskann_built_tests }}/test_streaming_scenario --data_type int8 --dist_fn l2 --data_path rand_int8_10D_10K_norm50.0.bin --index_path_prefix index_stream -R 64 -L 600 --alpha 1.2 --insert_threads 4 --consolidate_threads 4 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 200
+        ${{ env.diskann_built_tests }}/test_streaming_scenario --data_type int8 --dist_fn l2 --data_path ./rand_int8_10D_10K_norm50.0.bin --index_path_prefix index_stream -R 64 -L 600 --alpha 1.2 --insert_threads 4 --consolidate_threads 4 --max_points_to_insert 10000 --active_window 4000 --consolidate_interval 2000 --start_point_norm 200
         ${{ env.diskann_built_utils }}/compute_groundtruth --data_type int8 --dist_fn l2 --base_file index_stream.after-streaming-act4000-cons2000-max10000.data --query_file rand_int8_10D_1K_norm50.0.bin --K 100 --gt_file gt100_base-act4000-cons2000-max10000 --tags_file index_stream.after-streaming-act4000-cons2000-max10000.tags
         ${{ env.diskann_built_tests }}/search_memory_index --data_type int8 --dist_fn l2 --fail_if_recall_below 70 --index_path_prefix index_stream.after-streaming-act4000-cons2000-max10000 --result_path res_stream --query_file ./rand_int8_10D_1K_norm50.0.bin --gt_file gt100_base-act4000-cons2000-max10000 -K 10 -L 20 40 60 80 100 -T 64 --dynamic true --tags 1
 

include/abstract_data_store.h

@@ -0,0 +1,105 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <vector>
+#include <string>
+
+#include "types.h"
+#include "windows_customizations.h"
+
+namespace diskann {
+
+template <typename data_t>
+class AbstractDataStore {
+ public:
+  AbstractDataStore(const location_t capacity, const size_t dim);
+
+  // Return number of points returned
+  virtual location_t load(const std::string &filename) = 0;
+
+  // Why does store take num_pts? Since store only has capacity, but we allow
+  // resizing we can end up in a situation where the store has spare capacity.
+  // To optimize disk utilization, we pass the number of points that are "true"
+  // points, so that the store can discard the empty locations before saving.
+  virtual size_t save(const std::string &filename,
+                      const location_t num_pts) = 0;
+
+  DISKANN_DLLEXPORT virtual location_t capacity() const;
+
+  DISKANN_DLLEXPORT virtual size_t get_dims() const;
+
+  // Implementers can choose to return _dim if they are not
+  // concerned about memory alignment.
+  // Returns _dim aligned to a 8-byte value. Used for allocating
+  // aligned memory for efficiency/simplicity of code.
+  virtual size_t get_aligned_dim() const = 0;
+
+  // populate the store with vectors (either from a pointer or bin file),
+  // potentially after normalizing the vectors if the metric deems so
+  virtual void populate_data(const data_t *vectors,
+                             const location_t num_pts) = 0;
+  virtual void populate_data(const std::string &filename,
+                             const size_t offset) = 0;
+
+  // reverse of populate, save the first num_pts many points back to bin file
+  virtual void save_data_to_bin(const std::string &filename,
+                                const location_t num_pts) = 0;
+
+  // Returns the updated capacity of the datastore. Clients should check
+  // if resize actually changed the capacity to new_num_points before
+  // proceeding with operations. See the code below:
+  //  auto new_capcity = data_store->resize(new_num_points);
+  //  if ( new_capacity >= new_num_points) {
+  //   //PROCEED
+  //  else
+  //    //ERROR.
+  virtual location_t resize(const location_t new_num_points);
+
+  // operations on vectors
+  virtual void get_vector(const location_t i, data_t *dest) const = 0;
+  virtual void set_vector(const location_t i, const data_t *const vector) = 0;
+  virtual void prefetch_vector(const location_t loc) = 0;
+
+  // internal shuffle operations to move around vectors
+  virtual void reposition_points(const location_t start_loc,
+                                 const location_t end_loc,
+                                 const location_t num_points) = 0;
+  virtual void copy_points(const location_t from_loc, const location_t to_loc,
+                           const location_t num_points) = 0;
+
+  // metric specific operations
+
+  virtual float get_distance(const data_t *query,
+                             const location_t loc) const = 0;
+  virtual void get_distance(const data_t *query, const location_t *locations,
+                            const uint32_t location_count,
+                            float *distances) const = 0;
+  virtual float get_distance(const location_t loc1,
+                             const location_t loc2) const = 0;
+
+  // stats of the data stored in store
+  // Returns the point in the dataset that is closest to the mean of all points
+  // in the dataset
+  virtual location_t calculate_medoid() const = 0;
+
+  // search helpers
+  virtual size_t get_alignment_factor() const = 0;
+
+ protected:
+  // Expand the datastore to new_num_points. Returns the new capacity created,
+  // which should be == new_num_points in the normal case. Implementers can also
+  // return _capacity to indicate that there are not implementing this method.
+  virtual location_t expand(const location_t new_num_points) = 0;
+
+  // Shrink the datastore to new_num_points. It is NOT an error if shrink
+  // doesn't reduce the capacity so callers need to check this correctly. See
+  // also for "default" implementation
+  virtual location_t shrink(const location_t new_num_points) = 0;
+
+  location_t _capacity;
+  size_t _dim;
+};
+
+}  // namespace diskann

include/abstract_graph_store.h

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "types.h"
+
+namespace diskann {
+
+class AbstractGraphStore {
+ public:
+  AbstractGraphStore(const size_t max_pts) : _capacity(max_pts) {}
+
+  virtual int load(const std::string &index_path_prefix) = 0;
+  virtual int store(const std::string &index_path_prefix) = 0;
+
+  virtual void get_adj_list(const location_t i,
+                            std::vector<location_t> &neighbors) = 0;
+  virtual void set_adj_list(const location_t i,
+                            std::vector<location_t> &neighbors) = 0;
+
+ private:
+  size_t _capacity;
+};
+
+}  // namespace diskann

include/aligned_file_reader.h

@@ -18,11 +18,10 @@ typedef io_context_t IOContext;
 #include <minwinbase.h>
 
 #ifndef USE_BING_INFRA
-struct IOContext
-{
-    HANDLE fhandle = NULL;
-    HANDLE iocp = NULL;
-    std::vector<OVERLAPPED> reqs;
+struct IOContext {
+  HANDLE fhandle = NULL;
+  HANDLE iocp = NULL;
+  std::vector<OVERLAPPED> reqs;
 };
 #else
 #include "IDiskPriorityIO.h"
@@ -32,32 +31,25 @@ struct IOContext
 // errors.
 // Because of such callous copying, we have to use ptr->atomic instead
 // of atomic, as atomic is not copyable.
-struct IOContext
-{
-    enum Status
-    {
-        READ_WAIT = 0,
-        READ_SUCCESS,
-        READ_FAILED,
-        PROCESS_COMPLETE
-    };
-
-    std::shared_ptr<ANNIndex::IDiskPriorityIO> m_pDiskIO = nullptr;
-    std::shared_ptr<std::vector<ANNIndex::AsyncReadRequest>> m_pRequests;
-    std::shared_ptr<std::vector<Status>> m_pRequestsStatus;
-
-    // waitonaddress on this memory to wait for IO completion signal
-    // reader should signal this memory after IO completion
-    // TODO: WindowsAlignedFileReader can be modified to take advantage of this
-    //   and can largely share code with the file reader for Bing.
-    mutable volatile long m_completeCount = 0;
-
-    IOContext()
-        : m_pRequestsStatus(new std::vector<Status>()), m_pRequests(new std::vector<ANNIndex::AsyncReadRequest>())
-    {
-        (*m_pRequestsStatus).reserve(MAX_IO_DEPTH);
-        (*m_pRequests).reserve(MAX_IO_DEPTH);
-    }
+struct IOContext {
+  enum Status { READ_WAIT = 0, READ_SUCCESS, READ_FAILED, PROCESS_COMPLETE };
+
+  std::shared_ptr<ANNIndex::IDiskPriorityIO> m_pDiskIO = nullptr;
+  std::shared_ptr<std::vector<ANNIndex::AsyncReadRequest> > m_pRequests;
+  std::shared_ptr<std::vector<Status> > m_pRequestsStatus;
+
+  // waitonaddress on this memory to wait for IO completion signal
+  // reader should signal this memory after IO completion
+  // TODO: WindowsAlignedFileReader can be modified to take advantage of this
+  //   and can largely share code with the file reader for Bing.
+  mutable volatile long m_completeCount = 0;
+
+  IOContext()
+      : m_pRequestsStatus(new std::vector<Status>()),
+        m_pRequests(new std::vector<ANNIndex::AsyncReadRequest>()) {
+    (*m_pRequestsStatus).reserve(MAX_IO_DEPTH);
+    (*m_pRequests).reserve(MAX_IO_DEPTH);
+  }
 };
 #endif
 
@@ -71,50 +63,47 @@ struct IOContext
 #include "utils.h"
 
 // NOTE :: all 3 fields must be 512-aligned
-struct AlignedRead
-{
-    uint64_t offset; // where to read from
-    uint64_t len;    // how much to read
-    void *buf;       // where to read into
-
-    AlignedRead() : offset(0), len(0), buf(nullptr)
-    {
-    }
-
-    AlignedRead(uint64_t offset, uint64_t len, void *buf) : offset(offset), len(len), buf(buf)
-    {
-        assert(IS_512_ALIGNED(offset));
-        assert(IS_512_ALIGNED(len));
-        assert(IS_512_ALIGNED(buf));
-        // assert(malloc_usable_size(buf) >= len);
-    }
+struct AlignedRead {
+  uint64_t offset;  // where to read from
+  uint64_t len;     // how much to read
+  void *buf;        // where to read into
+
+  AlignedRead() : offset(0), len(0), buf(nullptr) {}
+
+  AlignedRead(uint64_t offset, uint64_t len, void *buf)
+      : offset(offset), len(len), buf(buf) {
+    assert(IS_512_ALIGNED(offset));
+    assert(IS_512_ALIGNED(len));
+    assert(IS_512_ALIGNED(buf));
+    // assert(malloc_usable_size(buf) >= len);
+  }
 };
 
-class AlignedFileReader
-{
-  protected:
-    tsl::robin_map<std::thread::id, IOContext> ctx_map;
-    std::mutex ctx_mut;
-
-  public:
-    // returns the thread-specific context
-    // returns (io_context_t)(-1) if thread is not registered
-    virtual IOContext &get_ctx() = 0;
-
-    virtual ~AlignedFileReader(){};
-
-    // register thread-id for a context
-    virtual void register_thread() = 0;
-    // de-register thread-id for a context
-    virtual void deregister_thread() = 0;
-    virtual void deregister_all_threads() = 0;
-
-    // Open & close ops
-    // Blocking calls
-    virtual void open(const std::string &fname) = 0;
-    virtual void close() = 0;
-
-    // process batch of aligned requests in parallel
-    // NOTE :: blocking call
-    virtual void read(std::vector<AlignedRead> &read_reqs, IOContext &ctx, bool async = false) = 0;
+class AlignedFileReader {
+ protected:
+  tsl::robin_map<std::thread::id, IOContext> ctx_map;
+  std::mutex ctx_mut;
+
+ public:
+  // returns the thread-specific context
+  // returns (io_context_t)(-1) if thread is not registered
+  virtual IOContext &get_ctx() = 0;
+
+  virtual ~AlignedFileReader(){};
+
+  // register thread-id for a context
+  virtual void register_thread() = 0;
+  // de-register thread-id for a context
+  virtual void deregister_thread() = 0;
+  virtual void deregister_all_threads() = 0;
+
+  // Open & close ops
+  // Blocking calls
+  virtual void open(const std::string &fname) = 0;
+  virtual void close() = 0;
+
+  // process batch of aligned requests in parallel
+  // NOTE :: blocking call
+  virtual void read(std::vector<AlignedRead> &read_reqs, IOContext &ctx,
+                    bool async = false) = 0;
 };

include/ann_exception.h

@@ -11,24 +11,25 @@
 #define __FUNCSIG__ __PRETTY_FUNCTION__
 #endif
 
-namespace diskann
-{
+namespace diskann {
 
-class ANNException : public std::runtime_error
-{
-  public:
-    DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode);
-    DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode, const std::string &funcSig,
-                                   const std::string &fileName, uint32_t lineNum);
+class ANNException : public std::runtime_error {
+ public:
+  DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode);
+  DISKANN_DLLEXPORT ANNException(const std::string &message, int errorCode,
+                                 const std::string &funcSig,
+                                 const std::string &fileName, uint32_t lineNum);
 
-  private:
-    int _errorCode;
+ private:
+  int _errorCode;
 };
 
-class FileException : public ANNException
-{
-  public:
-    DISKANN_DLLEXPORT FileException(const std::string &filename, std::system_error &e, const std::string &funcSig,
-                                    const std::string &fileName, uint32_t lineNum);
+class FileException : public ANNException {
+ public:
+  DISKANN_DLLEXPORT FileException(const std::string &filename,
+                                  std::system_error &e,
+                                  const std::string &funcSig,
+                                  const std::string &fileName,
+                                  uint32_t lineNum);
 };
-} // namespace diskann
+}  // namespace diskann

include/boost_dynamic_bitset_fwd.h

@@ -3,9 +3,10 @@
 
 #pragma once
 
-namespace boost
-{
+namespace boost {
 #ifndef BOOST_DYNAMIC_BITSET_FWD_HPP
-template <typename Block = unsigned long, typename Allocator = std::allocator<Block>> class dynamic_bitset;
+template <typename Block = unsigned long,
+          typename Allocator = std::allocator<Block>>
+class dynamic_bitset;
 #endif
-} // namespace boost
+}  // namespace boost