57 lines
2.5 KiB
C++
57 lines
2.5 KiB
C++
#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
|