Method search

python_bindings/ivf_bindings.cpp:63–107

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

        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 (!built_) {
            throw std::runtime_error("IvfIndex 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");
        }
        if (k == 0 || k > max_elements_) {
            throw std::invalid_argument("k must be between 1 and max_elements");
        }
        if (nprobe == 0) {
            throw std::invalid_argument("nprobe must be positive");
        }

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