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>(data_array.shape(0)) != max_elements_) {
            throw std::invalid_argument("number of data rows must match index max_elements"
            );
        }
        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>(centroids_array.shape(0)) != num_clusters_) {
            throw std::invalid_argument(
                "number of centroid rows must match index num_clusters"
            );
        }
        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"
            );
        }