Class SymqgIndex

python_bindings/symqg_bindings.cpp:17–110

namespace rabitqlib::python_bindings {

class SymqgIndex {
   public:
    SymqgIndex(size_t dim, size_t max_degree, const std::string& metric = "l2")
        : dim_(dim)
        , max_degree_(max_degree)
        , metric_(metric_from_string(metric)) {}

    void build(py::handle data, size_t ef_construction, size_t num_threads = 1) {
        auto data_array = ensure_2d_array<float>(data, "data");
        if (static_cast<size_t>(data_array.shape(1)) != dim_) {
            throw std::invalid_argument("data dimension does not match index dim");
        }

        num_points_ = static_cast<size_t>(data_array.shape(0));
        index_ = std::make_unique<rabitqlib::symqg::QuantizedGraph<float>>(
            num_points_, dim_, max_degree_, metric_, rabitqlib::RotatorType::FhtKacRotator
        );

        rabitqlib::symqg::QGBuilder builder(*index_, ef_construction, data_array.data(), num_threads);
        builder.build();
        built_ = true;
    }

    py::tuple search(py::handle queries, size_t k, size_t ef, size_t num_threads = 1) {
        auto query_array = ensure_2d_array<float>(queries, "queries");
        if (!built_) {
            throw std::runtime_error("SymqgIndex must be built or loaded before search");
        }
        if (static_cast<size_t>(query_array.shape(1)) != dim_) {
            throw std::invalid_argument("query dimension does not match index dim");
        }

        index_->set_ef(ef);

        const size_t nq = static_cast<size_t>(query_array.shape(0));
        const auto shape = std::vector<ssize_t>{static_cast<ssize_t>(nq), static_cast<ssize_t>(k)};
        auto ids = py::array_t<rabitqlib::PID>(shape);
        auto dists = py::array_t<float>(shape);
        auto ids_buf = ids.mutable_unchecked<2>();
        auto dists_buf = dists.mutable_unchecked<2>();
    
        rabitqlib::ivf::parallel_for(
            0,
            nq,
            num_threads,
            [&](size_t idx, size_t /*threadId*/) {
                std::vector<rabitqlib::PID> row_ids(k, 0);
                std::vector<float> row_dists(k, 0.0F);
                index_->search(query_array.data() + (idx * dim_), static_cast<uint32_t>(k), row_ids.data(), row_dists.data());
                for (size_t j = 0; j < k; ++j) {
                    ids_buf(static_cast<ssize_t>(idx), static_cast<ssize_t>(j)) = row_ids[j];
                    dists_buf(static_cast<ssize_t>(idx), static_cast<ssize_t>(j)) = row_dists[j];
                }
            }
        );

        return py::make_tuple(ids, dists);
    }