| 82 | } |
| 83 | |
| 84 | py::tuple search(py::handle queries, size_t k, size_t ef = 0, size_t num_threads = 1) { |
| 85 | auto query_array = ensure_2d_array<float>(queries, "queries"); |
| 86 | if (dim_ != 0 && static_cast<size_t>(query_array.shape(1)) != dim_) { |
| 87 | throw std::invalid_argument("query dimension does not match index dim"); |
| 88 | } |
| 89 | if (ef == 0) { |
| 90 | ef = std::max<size_t>(k, 10); |
| 91 | } |
| 92 | |
| 93 | const auto shape = std::vector<ssize_t>{ |
| 94 | static_cast<ssize_t>(query_array.shape(0)), static_cast<ssize_t>(k)}; |
| 95 | auto ids = py::array_t<rabitqlib::PID>(shape); |
| 96 | auto dists = py::array_t<float>(shape); |
| 97 | auto ids_buf = ids.mutable_unchecked<2>(); |
| 98 | auto dists_buf = dists.mutable_unchecked<2>(); |
| 99 | |
| 100 | std::vector<std::vector<std::pair<float, rabitqlib::PID>>> results = index_->search( |
| 101 | query_array.data(), |
| 102 | static_cast<size_t>(query_array.shape(0)), |
| 103 | k, |
| 104 | ef, |
| 105 | num_threads |
| 106 | ); |
| 107 | |
| 108 | for (ssize_t i = 0; i < static_cast<ssize_t>(results.size()); ++i) { |
| 109 | for ( |
| 110 | ssize_t j = 0; |
| 111 | j < static_cast<ssize_t>(std::min<size_t>(k, results[static_cast<size_t>(i)].size())); |
| 112 | ++j |
| 113 | ) { |
| 114 | ids_buf(i, j) = results[static_cast<size_t>(i)][static_cast<size_t>(j)].second; |
| 115 | dists_buf(i, j) = results[static_cast<size_t>(i)][static_cast<size_t>(j)].first; |
| 116 | } |
| 117 | } |
| 118 | return py::make_tuple(ids, dists); |
| 119 | } |
| 120 | |
| 121 | void save(const std::string& path) const { |
| 122 | if (!built_) { |