Initial commit

packages/leann-backend-diskann/third_party/DiskANN/include/quantized_distance.h

@@ -0,0 +1,56 @@
+#pragma once
+#include <memory>
+#include <string>
+#include <vector>
+#include "abstract_scratch.h"
+
+namespace diskann
+{
+template <typename data_t> class PQScratch;
+
+template <typename data_t> class QuantizedDistance
+{
+  public:
+    QuantizedDistance() = default;
+    QuantizedDistance(const QuantizedDistance &) = delete;
+    QuantizedDistance &operator=(const QuantizedDistance &) = delete;
+    virtual ~QuantizedDistance() = default;
+
+    virtual bool is_opq() const = 0;
+    virtual std::string get_quantized_vectors_filename(const std::string &prefix) const = 0;
+    virtual std::string get_pivot_data_filename(const std::string &prefix) const = 0;
+    virtual std::string get_rotation_matrix_suffix(const std::string &pq_pivots_filename) const = 0;
+
+    // Loading the PQ centroid table need not be part of the abstract class.
+    // However, we want to indicate that this function will change once we have a
+    // file reader hierarchy, so leave it here as-is.
+#ifdef EXEC_ENV_OLS
+    virtual void load_pivot_data(MemoryMappedFiles &files, const std::String &pq_table_file, size_t num_chunks) = 0;
+#else
+    virtual void load_pivot_data(const std::string &pq_table_file, size_t num_chunks) = 0;
+#endif
+
+    // Number of chunks in the PQ table. Depends on the compression level used.
+    // Has to be < ndim
+    virtual uint32_t get_num_chunks() const = 0;
+
+    // Preprocess the query by computing chunk distances from the query vector to
+    // various centroids. Since we don't want this class to do scratch management,
+    // we will take a PQScratch object which can come either from Index class or
+    // PQFlashIndex class.
+    virtual void preprocess_query(const data_t *query_vec, uint32_t query_dim, PQScratch<data_t> &pq_scratch) = 0;
+
+    // Workhorse
+    // This function must be called after preprocess_query
+    virtual void preprocessed_distance(PQScratch<data_t> &pq_scratch, const uint32_t id_count, float *dists_out) = 0;
+
+    // Same as above, but convenience function for index.cpp.
+    virtual void preprocessed_distance(PQScratch<data_t> &pq_scratch, const uint32_t n_ids,
+                                       std::vector<float> &dists_out) = 0;
+
+    // Currently this function is required for DiskPQ. However, it too can be subsumed
+    // under preprocessed_distance if we add the appropriate scratch variables to
+    // PQScratch and initialize them in pq_flash_index.cpp::disk_iterate_to_fixed_point()
+    virtual float brute_force_distance(const float *query_vec, uint8_t *base_vec) = 0;
+};
+} // namespace diskann