diff --git a/include/usearch/index_dense.hpp b/include/usearch/index_dense.hpp index e638d00c..e151b929 100644 --- a/include/usearch/index_dense.hpp +++ b/include/usearch/index_dense.hpp @@ -1309,11 +1309,10 @@ class index_dense_gt { shared_lock_t lock(slot_lookup_mutex_); offset = (std::min)(offset, slot_lookup_.size()); slot_lookup_.for_each([&](key_and_slot_t const& key_and_slot) { - if (offset) { + if (offset) + // Skip the first `offset` entries --offset; - return; - } - if (limit) { + else if (limit) { *keys = key_and_slot.key; ++keys; --limit; diff --git a/python/lib.cpp b/python/lib.cpp index e1461f41..dbb66935 100644 --- a/python/lib.cpp +++ b/python/lib.cpp @@ -798,7 +798,7 @@ template std::vector compute_lev } template -static py::tuple get_typed_vectors_for_keys(index_at const& index, py::buffer keys) { +static py::object get_typed_vectors_for_keys(index_at const& index, py::buffer keys) { py::buffer_info keys_info = keys.request(); if (keys_info.ndim != 1) @@ -806,27 +806,37 @@ static py::tuple get_typed_vectors_for_keys(index_at const& index, py::buffer ke Py_ssize_t keys_count = keys_info.shape[0]; byte_t const* keys_data = reinterpret_cast(keys_info.ptr); - py::tuple results(keys_count); - for (Py_ssize_t task_idx = 0; task_idx != keys_count; ++task_idx) { - dense_key_t key = *reinterpret_cast(keys_data + task_idx * keys_info.strides[0]); - std::size_t vectors_count = index.count(key); - if (!vectors_count) { - results[task_idx] = py::none(); - continue; - } + if (index.multi()) { + py::tuple results(keys_count); + + for (Py_ssize_t task_idx = 0; task_idx != keys_count; ++task_idx) { + dense_key_t key = *reinterpret_cast(keys_data + task_idx * keys_info.strides[0]); + std::size_t vectors_count = index.count(key); + if (!vectors_count) { + results[task_idx] = py::none(); + continue; + } - py::array_t result_py({static_cast(vectors_count), // - static_cast(index.scalar_words())}); + py::array_t result_py({static_cast(vectors_count), // + static_cast(index.scalar_words())}); + auto result_py2d = result_py.template mutable_unchecked<2>(); + index.get(key, (internal_at*)&result_py2d(0, 0), vectors_count); + results[task_idx] = result_py; + } + return results; + } else { + py::array_t result_py({keys_count, static_cast(index.scalar_words())}); auto result_py2d = result_py.template mutable_unchecked<2>(); - index.get(key, (internal_at*)&result_py2d(0, 0), vectors_count); - results[task_idx] = result_py; + for (Py_ssize_t task_idx = 0; task_idx != keys_count; ++task_idx) { + dense_key_t key = *reinterpret_cast(keys_data + task_idx * keys_info.strides[0]); + index.get(key, (internal_at*)&result_py2d(task_idx, 0), 1); + } + return result_py; } - - return results; } -template py::tuple get_many(index_at const& index, py::buffer keys, scalar_kind_t scalar_kind) { +template py::object get_many(index_at const& index, py::buffer keys, scalar_kind_t scalar_kind) { if (scalar_kind == scalar_kind_t::f32_k) return get_typed_vectors_for_keys(index, keys); else if (scalar_kind == scalar_kind_t::f64_k)