Method search

python_bindings/symqg_bindings.cpp:40–74

    }

    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);
    }

    void save(const std::string& path) const {
        if (!built_) {