Class IvfIndex

python_bindings/ivf_bindings.cpp:15–146

namespace rabitqlib::python_bindings {

class IvfIndex {
   public:
    IvfIndex(
        size_t dim,
        size_t max_elements,
        size_t num_clusters,
        size_t nbits,
        const std::string& metric = "l2"
    )
        : dim_(dim)
        , max_elements_(max_elements)
        , num_clusters_(num_clusters)
        , nbits_(nbits)
        , metric_(metric_from_string(metric))
        , index_(std::make_unique<rabitqlib::ivf::IVF>(
              max_elements,
              dim,
              num_clusters,
              nbits,
              metric_,
              rabitqlib::RotatorType::FhtKacRotator
          )) {}

    void build(
        py::handle data,
        py::handle centroids,
        py::handle cluster_ids,
        size_t num_threads = 1,
        bool fast_quantization = false
    ) {
        auto data_array = ensure_2d_array<float>(data, "data");
        auto centroids_array = ensure_2d_array<float>(centroids, "centroids");
        auto cluster_ids_array = ensure_1d_array<rabitqlib::PID>(cluster_ids, "cluster_ids");

        if (static_cast<size_t>(data_array.shape(1)) != dim_) {
            throw std::invalid_argument("data dimension does not match index dim");
        }
        if (static_cast<size_t>(centroids_array.shape(1)) != dim_) {
            throw std::invalid_argument("centroid dimension does not match index dim");
        }
        if (static_cast<size_t>(cluster_ids_array.shape(0)) != static_cast<size_t>(data_array.shape(0))) {
            throw std::invalid_argument("cluster_ids length must match number of rows in data");
        }

        index_->construct(data_array.data(), centroids_array.data(), cluster_ids_array.data(), fast_quantization, num_threads);
        built_ = true;
    }

    py::tuple search(
        py::handle queries,
        size_t k,
        size_t nprobe,
        bool high_accuracy = true,
        size_t num_threads = 1
    ) {
        auto query_array = ensure_2d_array<float>(queries, "queries");
        if (static_cast<size_t>(query_array.shape(1)) != dim_) {
            throw std::invalid_argument("query dimension does not match index dim");