inline progress_t(std::nullptr_t = nullptr) : func_(&dummy_progress) {}

inline progress_t(progress_func_t const& func) : func_(func ? func : &dummy_progress) {}

// The callback may be invoked from a worker thread while the main thread

// has released the GIL around a long-running C++ call. `gil_scoped_acquire`

// is reference-counted and a no-op when the GIL is already held, so this

// is safe in every context.

inline bool operator()(std::size_t processed, std::size_t total) const noexcept {

py::gil_scoped_acquire acquire;

return func_(processed, total);

}

private:

static inline bool dummy_progress(std::size_t /*processed*/, std::size_t /*total*/) { return true; }

progress_func_t func_;

using native_t = index_dense_t;

using native_t::add;

using native_t::capacity;

using native_t::reserve;

using native_t::search;

using native_t::size;

// Serializes Python-thread access to the underlying C++ index. The native

// `index_dense_t` assumes a single owning Python thread (it carries

// per-worker `cast_buffer_` slots indexed by an executor-local `thread_idx`

// that the binding picks for each call). Multiple Python threads calling

// a heavy op on the same index would otherwise collide on those slots.

// Locked around every binding entry point that releases the GIL.

//

// Held via `unique_ptr` so that the wrapper remains move-constructible -

// `std::mutex` itself is neither copyable nor movable, and pybind11's

// return-by-value factories require movability.

mutable std::unique_ptr<std::mutex> mutex_ptr_ = std::make_unique<std::mutex>();

dense_index_py_t(native_t&& base) : index_dense_t(std::move(base)) {}

std::vector<std::shared_ptr<dense_index_py_t>> shards_;

mutable std::unique_ptr<std::mutex> mutex_ptr_ = std::make_unique<std::mutex>();

void merge(std::shared_ptr<dense_index_py_t> shard) { shards_.push_back(shard); }

std::size_t bytes_per_vector() const noexcept { return shards_.empty() ? 0 : shards_[0]->bytes_per_vector(); }

std::size_t scalar_words() const noexcept { return shards_.empty() ? 0 : shards_[0]->scalar_words(); }

index_limits_t limits() const noexcept { return {size(), std::numeric_limits<std::size_t>::max()}; }

void merge_paths(std::vector<std::string> const& paths, bool view = true, std::size_t threads = 0) {

if (!threads)

threads = std::thread::hardware_concurrency();

shards_.reserve(shards_.size() + paths.size());

std::mutex shards_mutex;

// Release the GIL *before* taking the per-index mutex so a Python

// thread waiting on the mutex doesn't hold the GIL - otherwise a

// worker thread in the current owner would block forever in

// `gil_scoped_acquire`.

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*mutex_ptr_);

executor_default_t{threads}.dynamic(paths.size(), [&](std::size_t, std::size_t task_idx) {

index_dense_t index = index_dense_t::make(paths[task_idx].c_str(), view);

if (!index)

return false;

auto shared_index = std::make_shared<dense_index_py_t>(std::move(index));

std::unique_lock<std::mutex> lock(shards_mutex);

shards_.push_back(shared_index);

py::gil_scoped_acquire acquire;

if (PyErr_CheckSignals() != 0)

throw py::error_already_set();

return true;

});

}

std::size_t size() const noexcept {

std::size_t result = 0;

for (auto const& shard : shards_)

result += shard->size();

return result;

}

std::size_t dimensions, //

scalar_kind_t scalar_kind, //

std::size_t connectivity, //

std::size_t expansion_add, //

std::size_t expansion_search, //

metric_kind_t metric_kind, //

metric_punned_signature_t metric_signature, //

std::uintptr_t metric_uintptr, //

bool multi, //

bool enable_key_lookups) {

index_dense_config_t config(connectivity, expansion_add, expansion_search);

config.multi = multi;

config.enable_key_lookups = enable_key_lookups;

metric_t metric = //

metric_uintptr //

? metric_t::stateless(dimensions, metric_uintptr, metric_signature, metric_kind, scalar_kind)

: metric_t::builtin(dimensions, metric_kind, scalar_kind);

if (metric.missing())

throw std::invalid_argument("Unsupported metric!");

using index_state_t = typename index_dense_t::state_result_t;

index_state_t state = index_dense_t::make(metric, config);

if (!state)

throw std::invalid_argument(state.error.release());

return std::move(state.index);

// https://docs.python.org/3/library/struct.html#format-characters

if (name == "B" || name == "<B" || name == "u1" || name == "|u1")

return scalar_kind_t::b1x8_k;

else if (name == "b" || name == "<b" || name == "i1" || name == "|i1")

return scalar_kind_t::i8_k;

else if (name == "e" || name == "<e" || name == "f2" || name == "<f2")

return scalar_kind_t::f16_k;

else if (name == "f" || name == "<f" || name == "f4" || name == "<f4")

return scalar_kind_t::f32_k;

else if (name == "d" || name == "<d" || name == "i8" || name == "<i8")

return scalar_kind_t::f64_k;

else

return scalar_kind_t::unknown_k;

if (!result)

throw std::invalid_argument(result.error.release());

int signals = PyErr_CheckSignals();

if (signals != 0)

throw py::error_already_set();

dense_index_py_t& index, //

py::buffer_info const& keys_info, py::buffer_info const& vectors_info, //

bool force_copy, std::size_t threads, //

progress_func_t const& progress) {

Py_ssize_t vectors_count = vectors_info.shape[0];

byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);

byte_t const* keys_data = reinterpret_cast<byte_t const*>(keys_info.ptr);

atomic_error_t atomic_error{nullptr};

// Progress status

progress_t progress_{progress};

std::atomic<std::size_t> processed{0};

{

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*index.mutex_ptr_);

if (!index.try_reserve(index_limits_t(ceil2(index.size() + vectors_count), threads)))

throw std::invalid_argument("Out of memory!");

executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {

dense_key_t key = *reinterpret_cast<dense_key_t const*>(keys_data + task_idx * keys_info.strides[0]);

scalar_at const* vector =

reinterpret_cast<scalar_at const*>(vectors_data + task_idx * vectors_info.strides[0]);

dense_add_result_t result = index.add(key, vector, thread_idx, force_copy);

if (!result) {

atomic_error = result.error.release();

return false;

}

// We don't want to check for signals from multiple threads

++processed;

if (thread_idx == 0) {

py::gil_scoped_acquire acquire;

if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), vectors_count)) {

atomic_error.store("Operation has been terminated");

return false;

}

}

return true;

});

}

// At the end report the latest numbers, because the reporter thread may be finished earlier

progress_(processed.load(), vectors_count);

// Raise the error from a single thread

auto error = atomic_error.load();

if (error) {

PyErr_SetString(PyExc_RuntimeError, error);

throw py::error_already_set();

}

index_at& index, py::buffer keys, py::buffer vectors, //

bool force_copy, std::size_t threads, //

progress_func_t const& progress, //

scalar_kind_t scalar_kind = scalar_kind_t::unknown_k) {

py::buffer_info keys_info = keys.request();

py::buffer_info vectors_info = vectors.request();

if (keys_info.itemsize != sizeof(dense_key_t))

throw std::invalid_argument("Incompatible key type!");

if (keys_info.strides[0] != static_cast<Py_ssize_t>(keys_info.itemsize))

throw std::invalid_argument("Keys array must be C-contiguous.");

if (keys_info.ndim != 1)

throw std::invalid_argument("Keys must be placed in a single-dimensional array!");

if (vectors_info.ndim != 2)

throw std::invalid_argument("Expects a matrix of vectors to add!");

if (vectors_info.strides[1] != static_cast<Py_ssize_t>(vectors_info.itemsize))

throw std::invalid_argument("Matrix rows must be contiguous, try `ascontiguousarray`.");

Py_ssize_t keys_count = keys_info.shape[0];

Py_ssize_t vectors_count = vectors_info.shape[0];

Py_ssize_t vectors_dimensions = vectors_info.shape[1];

if (vectors_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))

throw std::invalid_argument("The number of vector dimensions doesn't match!");

if (keys_count != vectors_count)

throw std::invalid_argument("Number of keys and vectors must match!");

if (!threads)

threads = std::thread::hardware_concurrency();

// `add_typed_to_index` does the `try_reserve` + executor work inside its

// own GIL-released, mutex-locked region; we just dispatch on the scalar

// kind here.

// clang-format off

scalar_kind_t kind = (scalar_kind != scalar_kind_t::unknown_k)

? scalar_kind

: numpy_string_to_kind(vectors_info.format);

switch (kind) {

case scalar_kind_t::f64_k: add_typed_to_index<f64_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::f32_k: add_typed_to_index<f32_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::bf16_k: add_typed_to_index<bf16_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::f16_k: add_typed_to_index<f16_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::e5m2_k: add_typed_to_index<e5m2_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::e4m3_k: add_typed_to_index<e4m3_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::e3m2_k: add_typed_to_index<e3m2_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::e2m3_k: add_typed_to_index<e2m3_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::i8_k: add_typed_to_index<i8_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::u8_k: add_typed_to_index<u8_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

case scalar_kind_t::b1x8_k: add_typed_to_index<b1x8_t>(index, keys_info, vectors_info, force_copy, threads, progress); break;

default: throw std::invalid_argument("Incompatible scalars in the vectors matrix: " + vectors_info.format);

}

// clang-format on

dense_index_py_t& index, py::buffer_info& vectors_info, //

std::size_t wanted, bool exact, std::size_t threads, //

py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,

std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances,

progress_func_t const& progress) {

auto keys_py2d = keys_py.template mutable_unchecked<2>();

auto distances_py2d = distances_py.template mutable_unchecked<2>();

auto counts_py1d = counts_py.template mutable_unchecked<1>();

Py_ssize_t vectors_count = vectors_info.shape[0];

byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);

if (!threads)

threads = std::thread::hardware_concurrency();

// Progress status

progress_t progress_{progress};

std::atomic<std::size_t> processed{0};

atomic_error_t atomic_error{nullptr};

{

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*index.mutex_ptr_);

if (!index.try_reserve(index_limits_t(index.size(), threads)))

throw std::invalid_argument("Out of memory!");

executor_default_t{threads}.dynamic(vectors_count, [&](std::size_t thread_idx, std::size_t task_idx) {

scalar_at const* vector = (scalar_at const*)(vectors_data + task_idx * vectors_info.strides[0]);

dense_search_result_t result = index.search(vector, wanted, thread_idx, exact);

if (!result) {

atomic_error = result.error.release();

return false;

}

counts_py1d(task_idx) =

static_cast<Py_ssize_t>(result.dump_to(&keys_py2d(task_idx, 0), &distances_py2d(task_idx, 0), wanted));

stats_visited_members += result.visited_members;

stats_computed_distances += result.computed_distances;

// We don't want to check for signals from multiple threads

++processed;

if (thread_idx == 0) {

py::gil_scoped_acquire acquire;

if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), vectors_count)) {

atomic_error.store("Operation has been terminated");

return false;

}

}

return true;

});

}

// At the end report the latest numbers, because the reporter thread may be finished earlier

progress_(processed.load(), vectors_count);

// Raise the error from a single thread

auto error = atomic_error.load();

if (error) {

PyErr_SetString(PyExc_RuntimeError, error);

throw py::error_already_set();

}

dense_indexes_py_t& indexes, py::buffer_info& vectors_info, //

std::size_t wanted, bool exact, std::size_t threads, //

py::array_t<dense_key_t>& keys_py, py::array_t<distance_t>& distances_py, py::array_t<Py_ssize_t>& counts_py,

std::atomic<std::size_t>& stats_visited_members, std::atomic<std::size_t>& stats_computed_distances,

progress_func_t const& progress) {

auto keys_py2d = keys_py.template mutable_unchecked<2>();

auto distances_py2d = distances_py.template mutable_unchecked<2>();

auto counts_py1d = counts_py.template mutable_unchecked<1>();

Py_ssize_t vectors_count = vectors_info.shape[0];

byte_t const* vectors_data = reinterpret_cast<byte_t const*>(vectors_info.ptr);

for (std::size_t vector_idx = 0; vector_idx != static_cast<std::size_t>(vectors_count); ++vector_idx)

counts_py1d(vector_idx) = 0;

if (!threads)

threads = std::thread::hardware_concurrency();

bitset_t query_mutexes(static_cast<std::size_t>(vectors_count));

if (!query_mutexes)

throw std::bad_alloc();

// Progress status

progress_t progress_{progress};

std::atomic<std::size_t> processed{0};

atomic_error_t atomic_error{nullptr};

{

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*indexes.mutex_ptr_);

executor_default_t{threads}.dynamic(indexes.shards_.size(), [&](std::size_t thread_idx, std::size_t task_idx) {

dense_index_py_t& index = *indexes.shards_[task_idx].get();

index_limits_t limits;

limits.members = index.size();

limits.threads_add = 0;

limits.threads_search = 1;

if (!index.try_reserve(limits)) {

atomic_error = "Out of memory!";

return false;

}

for (std::size_t vector_idx = 0; vector_idx != static_cast<std::size_t>(vectors_count); ++vector_idx) {

scalar_at const* vector = (scalar_at const*)(vectors_data + vector_idx * vectors_info.strides[0]);

dense_search_result_t result = index.search(vector, wanted, 0, exact);

if (!result) {

atomic_error = result.error.release();

return false;

}

{

auto lock = query_mutexes.lock(vector_idx);

counts_py1d(vector_idx) = static_cast<Py_ssize_t>(result.merge_into( //

&keys_py2d(vector_idx, 0), //

&distances_py2d(vector_idx, 0), //

static_cast<std::size_t>(counts_py1d(vector_idx)), //

wanted));

}

stats_visited_members += result.visited_members;

stats_computed_distances += result.computed_distances;

// We don't want to check for signals from multiple threads

++processed;

if (thread_idx == 0) {

py::gil_scoped_acquire acquire;

if (PyErr_CheckSignals() != 0 || !progress_(processed.load(), indexes.shards_.size())) {

atomic_error.store("Operation has been terminated");

return false;

}

}

}

return true;

});

}

// At the end report the latest numbers, because the reporter thread may be finished earlier

progress_(processed.load(), indexes.shards_.size());

// Raise the error from a single thread

auto error = atomic_error.load();

if (error) {

PyErr_SetString(PyExc_RuntimeError, error);

throw py::error_already_set();

}

index_at& index, py::buffer vectors, std::size_t wanted, bool exact, std::size_t threads,

progress_func_t const& progress, scalar_kind_t scalar_kind = scalar_kind_t::unknown_k) {

if (wanted == 0)

return py::tuple(5);

// Clamp threads to hardware limit instead of throwing

threads = std::min<std::size_t>(threads, std::thread::hardware_concurrency());

py::buffer_info vectors_info = vectors.request();

if (vectors_info.ndim != 2)

throw std::invalid_argument("Expects a matrix of vectors to add!");

Py_ssize_t vectors_count = vectors_info.shape[0];

Py_ssize_t vectors_dimensions = vectors_info.shape[1];

if (vectors_dimensions != static_cast<Py_ssize_t>(index.scalar_words()))

throw std::invalid_argument("The number of vector dimensions doesn't match!");

if (vectors_info.strides[1] != static_cast<Py_ssize_t>(vectors_info.itemsize))

throw std::invalid_argument("Matrix rows must be contiguous, try `ascontiguousarray`.");

py::array_t<dense_key_t> keys_py({vectors_count, static_cast<Py_ssize_t>(wanted)});

py::array_t<distance_t> distances_py({vectors_count, static_cast<Py_ssize_t>(wanted)});

py::array_t<Py_ssize_t> counts_py(vectors_count);

std::atomic<std::size_t> stats_visited_members(0);

std::atomic<std::size_t> stats_computed_distances(0);

// clang-format off

scalar_kind_t kind = (scalar_kind != scalar_kind_t::unknown_k)

? scalar_kind

: numpy_string_to_kind(vectors_info.format);

switch (kind) {

case scalar_kind_t::f64_k: search_typed<f64_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::f32_k: search_typed<f32_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::bf16_k: search_typed<bf16_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::f16_k: search_typed<f16_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::e5m2_k: search_typed<e5m2_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::e4m3_k: search_typed<e4m3_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::e3m2_k: search_typed<e3m2_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::e2m3_k: search_typed<e2m3_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::i8_k: search_typed<i8_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::u8_k: search_typed<u8_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

case scalar_kind_t::b1x8_k: search_typed<b1x8_t>(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break;

default: throw std::invalid_argument("Incompatible scalars in the query matrix: " + vectors_info.format);

}

// clang-format on

py::tuple results(5);

results[0] = keys_py;

results[1] = distances_py;

results[2] = counts_py;

results[3] = stats_visited_members.load();

results[4] = stats_computed_distances.load();

return results;

py::buffer dataset, py::buffer queries, //

std::size_t wanted, std::size_t threads, //

metric_kind_t metric_kind, //

metric_punned_signature_t metric_signature, //

std::uintptr_t metric_uintptr, //

progress_func_t const& progress_func, //

scalar_kind_t scalar_kind = scalar_kind_t::unknown_k) {

if (wanted == 0)

return py::tuple(5);

py::buffer_info dataset_info = dataset.request();

py::buffer_info queries_info = queries.request();

if (dataset_info.ndim != 2 || queries_info.ndim != 2)

throw std::invalid_argument("Expects a matrix of dataset to add!");

if (dataset_info.strides[1] != static_cast<Py_ssize_t>(dataset_info.itemsize))

throw std::invalid_argument("Dataset rows must be contiguous, try `ascontiguousarray`.");

if (queries_info.strides[1] != static_cast<Py_ssize_t>(queries_info.itemsize))

throw std::invalid_argument("Queries rows must be contiguous, try `ascontiguousarray`.");

std::size_t dataset_count = static_cast<std::size_t>(dataset_info.shape[0]);

std::size_t dataset_dimensions = static_cast<std::size_t>(dataset_info.shape[1]);

std::size_t dataset_stride = static_cast<std::size_t>(dataset_info.strides[0]);

std::size_t queries_stride = static_cast<std::size_t>(queries_info.strides[0]);

std::size_t queries_count = static_cast<std::size_t>(queries_info.shape[0]);

std::size_t queries_dimensions = static_cast<std::size_t>(queries_info.shape[1]);

if (dataset_dimensions != queries_dimensions)

throw std::invalid_argument("The number of vector dimensions doesn't match!");

if (wanted > dataset_count)

throw std::invalid_argument("You can't request more matches than in the dataset!");

scalar_kind_t dataset_kind =

(scalar_kind != scalar_kind_t::unknown_k) ? scalar_kind : numpy_string_to_kind(dataset_info.format);

scalar_kind_t queries_kind =

(scalar_kind != scalar_kind_t::unknown_k) ? scalar_kind : numpy_string_to_kind(queries_info.format);

if (dataset_kind != queries_kind)

throw std::invalid_argument("The types of vectors don't match!");

std::size_t dimensions = static_cast<std::size_t>(queries_dimensions);

metric_t metric = //

metric_uintptr //

? metric_t::stateless(dimensions, metric_uintptr, metric_signature, metric_kind, queries_kind)

: metric_t::builtin(dimensions, metric_kind, queries_kind);

if (!metric)

throw std::invalid_argument("Unsupported metric!");

py::array_t<dense_key_t> keys_py({static_cast<Py_ssize_t>(queries_count), static_cast<Py_ssize_t>(wanted)});

py::array_t<distance_t> distances_py({static_cast<Py_ssize_t>(queries_count), static_cast<Py_ssize_t>(wanted)});

py::array_t<Py_ssize_t> counts_py(static_cast<Py_ssize_t>(queries_count));

auto keys_py2d = keys_py.template mutable_unchecked<2>();

auto distances_py2d = distances_py.template mutable_unchecked<2>();

auto counts_py1d = counts_py.template mutable_unchecked<1>();

byte_t const* dataset_data = reinterpret_cast<byte_t const*>(dataset_info.ptr);

byte_t const* queries_data = reinterpret_cast<byte_t const*>(queries_info.ptr);

for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)

counts_py1d(query_idx) = wanted;

if (!threads)

threads = std::thread::hardware_concurrency();

// Dispatch brute-force search

progress_t progress{progress_func};

executor_default_t executor{threads};

exact_search_t search;

exact_search_results_t offsets_and_distances;

{

py::gil_scoped_release release;

offsets_and_distances = search( //

dataset_data, dataset_count, dataset_stride, //

queries_data, queries_count, queries_stride, //

wanted, metric, executor, [&](std::size_t passed, std::size_t total) {

py::gil_scoped_acquire acquire;

return PyErr_CheckSignals() == 0 && progress(passed, total);

});

}

if (!offsets_and_distances)

throw std::bad_alloc();

// Export the results

for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx) {

dense_key_t* query_keys = &keys_py2d(query_idx, 0);

distance_t* query_distances = &distances_py2d(query_idx, 0);

auto query_result = offsets_and_distances.at(query_idx);

for (std::size_t i = 0; i != wanted; ++i)

query_keys[i] = static_cast<dense_key_t>(query_result[i].offset),

query_distances[i] = query_result[i].distance;

}

py::tuple results(5);

results[0] = keys_py;

results[1] = distances_py;

results[2] = counts_py;

results[3] = 0;

results[4] = static_cast<std::size_t>(dataset_count * queries_count);

return results;

py::buffer dataset, //

std::size_t wanted, //

std::size_t max_iterations, //

double inertia_threshold, //

double max_seconds, //

double min_shifts, //

std::uint64_t seed, //

std::size_t threads, //

scalar_kind_t scalar_kind, //

metric_kind_t metric_kind, //

progress_func_t const& progress_func) {

using distance_t = typename kmeans_clustering_t::distance_t;

py::buffer_info dataset_info = dataset.request();

if (dataset_info.ndim != 2)

throw std::invalid_argument("Expects a matrix (rank-2 tensor) of dataset to cluster!");

if (dataset_info.strides[1] != static_cast<Py_ssize_t>(dataset_info.itemsize))

throw std::invalid_argument("Dataset rows must be contiguous, try `ascontiguousarray`.");

std::size_t dataset_count = static_cast<std::size_t>(dataset_info.shape[0]);

std::size_t dataset_dimensions = static_cast<std::size_t>(dataset_info.shape[1]);

std::size_t dataset_stride = static_cast<std::size_t>(dataset_info.strides[0]);

scalar_kind_t dataset_kind = numpy_string_to_kind(dataset_info.format);

std::size_t bytes_per_scalar = bits_per_scalar_word(dataset_kind) / CHAR_BIT;

std::vector<std::size_t> point_to_centroid_index(dataset_count, 0);

std::vector<distance_t> point_to_centroid_distance(dataset_count, 0);

std::vector<byte_t> centroids(wanted * dataset_dimensions * bytes_per_scalar, 0);

if (!threads)

threads = std::thread::hardware_concurrency();

// Dispatch brute-force search

progress_t progress{progress_func};

executor_default_t executor{threads};

kmeans_clustering_t engine;

engine.metric_kind = metric_kind;

engine.quantization_kind = scalar_kind;

engine.max_iterations = max_iterations;

engine.min_shifts = min_shifts;

engine.max_seconds = max_seconds;

engine.inertia_threshold = inertia_threshold;

kmeans_clustering_result_t result;

{

py::gil_scoped_release release;

result = engine( //

reinterpret_cast<byte_t const*>(dataset_info.ptr), dataset_count, dataset_stride, //

centroids.data(), wanted, dataset_dimensions * bytes_per_scalar, //

point_to_centroid_index.data(), point_to_centroid_distance.data(), dataset_kind, dataset_dimensions,

executor, [&](std::size_t passed, std::size_t total) {

py::gil_scoped_acquire acquire;

return PyErr_CheckSignals() == 0 && progress(passed, total);

});

}

if (!result)

throw std::runtime_error(result.error.release());

// Following constructor doesn't seem to be documented, but it's used in the source code of `pybind11`

// https://github.com/pybind/pybind11/blob/aeda49ed0b4e6e8abba7abc265ace86a6c26ba66/include/pybind11/numpy.h#L918-L919

// https://github.com/pybind/pybind11/blob/aeda49ed0b4e6e8abba7abc265ace86a6c26ba66/include/pybind11/buffer_info.h#L60-L75

py::buffer_info centroids_info;

centroids_info.ptr = reinterpret_cast<void*>(centroids.data());

centroids_info.itemsize = dataset_info.itemsize;

centroids_info.size = wanted * dataset_dimensions;

centroids_info.format = dataset_info.format;

centroids_info.ndim = 2;

centroids_info.shape = {static_cast<ssize_t>(wanted), static_cast<ssize_t>(dataset_dimensions)};

centroids_info.strides = {static_cast<ssize_t>(dataset_dimensions * bytes_per_scalar),

static_cast<ssize_t>(bytes_per_scalar)};

py::tuple results(3);

results[0] = py::array_t<std::size_t>({static_cast<ssize_t>(dataset_count)}, point_to_centroid_index.data());

results[1] = py::array_t<distance_t>({static_cast<ssize_t>(dataset_count)}, point_to_centroid_distance.data());

results[2] = py::array(centroids_info);

return results;

byte_t* data_;

std::size_t stride_;

rows_lookup_gt(void* data, std::size_t stride) noexcept : data_((byte_t*)data), stride_(stride) {}

scalar_at* operator[](std::size_t i) const noexcept { return reinterpret_cast<scalar_at*>(data_ + i * stride_); }

std::ptrdiff_t operator-(rows_lookup_gt const& other) const noexcept { return (data_ - other.data_) / stride_; }

rows_lookup_gt operator+(std::size_t n) const noexcept { return {data_ + stride_ * n, stride_}; }

template <typename other_scalar_at> rows_lookup_gt<other_scalar_at> as() const noexcept { return {data_, stride_}; }

index_at& index, py::buffer queries, //

std::size_t min_count, std::size_t max_count, std::size_t threads, progress_func_t const& progress,

scalar_kind_t scalar_kind = scalar_kind_t::unknown_k) {

// Clamp threads to hardware limit instead of throwing

threads = std::min<std::size_t>(threads, std::thread::hardware_concurrency());

py::buffer_info queries_info = queries.request();

if (queries_info.ndim != 2)

throw std::invalid_argument("Expects a matrix of queries to add!");

std::size_t queries_count = static_cast<std::size_t>(queries_info.shape[0]);

std::size_t queries_stride = static_cast<std::size_t>(queries_info.strides[0]);

std::size_t queries_dimensions = static_cast<std::size_t>(queries_info.shape[1]);

if (queries_dimensions != index.scalar_words())

throw std::invalid_argument("The number of vector dimensions doesn't match!");

py::array_t<dense_key_t> keys_py({Py_ssize_t(queries_count), Py_ssize_t(1)});

py::array_t<distance_t> distances_py({Py_ssize_t(queries_count), Py_ssize_t(1)});

dense_clustering_result_t cluster_result;

executor_default_t executor{threads};

auto keys_py2d = keys_py.template mutable_unchecked<2>();

auto distances_py2d = distances_py.template mutable_unchecked<2>();

dense_key_t* keys_ptr = reinterpret_cast<dense_key_t*>(&keys_py2d(0, 0));

distance_t* distances_ptr = reinterpret_cast<distance_t*>(&distances_py2d(0, 0));

index_dense_clustering_config_t config;

config.min_clusters = min_count;

config.max_clusters = max_count;

rows_lookup_gt<byte_t const> queries_begin(queries_info.ptr, queries_stride);

rows_lookup_gt<byte_t const> queries_end = queries_begin + queries_count;

// clang-format off

scalar_kind_t kind = (scalar_kind != scalar_kind_t::unknown_k)

? scalar_kind

: numpy_string_to_kind(queries_info.format);

{

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*index.mutex_ptr_);

switch (kind) {

case scalar_kind_t::f64_k: cluster_result = index.cluster(queries_begin.as<f64_t const>(), queries_end.as<f64_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::f32_k: cluster_result = index.cluster(queries_begin.as<f32_t const>(), queries_end.as<f32_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::bf16_k: cluster_result = index.cluster(queries_begin.as<bf16_t const>(), queries_end.as<bf16_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::f16_k: cluster_result = index.cluster(queries_begin.as<f16_t const>(), queries_end.as<f16_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::e5m2_k: cluster_result = index.cluster(queries_begin.as<e5m2_t const>(), queries_end.as<e5m2_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::e4m3_k: cluster_result = index.cluster(queries_begin.as<e4m3_t const>(), queries_end.as<e4m3_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::e3m2_k: cluster_result = index.cluster(queries_begin.as<e3m2_t const>(), queries_end.as<e3m2_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::e2m3_k: cluster_result = index.cluster(queries_begin.as<e2m3_t const>(), queries_end.as<e2m3_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::i8_k: cluster_result = index.cluster(queries_begin.as<i8_t const>(), queries_end.as<i8_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::u8_k: cluster_result = index.cluster(queries_begin.as<u8_t const>(), queries_end.as<u8_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

case scalar_kind_t::b1x8_k: cluster_result = index.cluster(queries_begin.as<b1x8_t const>(), queries_end.as<b1x8_t const>(), config, keys_ptr, distances_ptr, executor, progress_t{progress}); break;

default: throw std::invalid_argument("Incompatible scalars in the query matrix: " + queries_info.format);

}

}

// clang-format on

cluster_result.error.raise();

// Those would be set to 1 for all entries, in case of success

py::array_t<Py_ssize_t> counts_py(queries_count);

auto counts_py1d = counts_py.template mutable_unchecked<1>();

for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)

counts_py1d(static_cast<Py_ssize_t>(query_idx)) = 1;

py::tuple results(5);

results[0] = keys_py;

results[1] = distances_py;

results[2] = counts_py;

results[3] = cluster_result.visited_members;

results[4] = cluster_result.computed_distances;

return results;

index_at& index, py::array_t<dense_key_t> queries_py, //

std::size_t min_count, std::size_t max_count, std::size_t threads, progress_func_t const& progress) {

// Clamp threads to hardware limit instead of throwing

threads = std::min<std::size_t>(threads, std::thread::hardware_concurrency());

std::size_t queries_count = static_cast<std::size_t>(queries_py.size());

auto queries_py1d = queries_py.template unchecked<1>();

dense_key_t const* queries_begin = &queries_py1d(0);

dense_key_t const* queries_end = queries_begin + queries_count;

py::array_t<dense_key_t> keys_py({Py_ssize_t(queries_count), Py_ssize_t(1)});

py::array_t<distance_t> distances_py({Py_ssize_t(queries_count), Py_ssize_t(1)});

executor_default_t executor{threads};

auto keys_py2d = keys_py.template mutable_unchecked<2>();

auto distances_py2d = distances_py.template mutable_unchecked<2>();

dense_key_t* keys_ptr = reinterpret_cast<dense_key_t*>(&keys_py2d(0, 0));

distance_t* distances_ptr = reinterpret_cast<distance_t*>(&distances_py2d(0, 0));

index_dense_clustering_config_t config;

config.min_clusters = min_count;

config.max_clusters = max_count;

dense_clustering_result_t cluster_result;

{

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*index.mutex_ptr_);

cluster_result =

index.cluster(queries_begin, queries_end, config, keys_ptr, distances_ptr, executor, progress_t{progress});

}

cluster_result.error.raise();

// Those would be set to 1 for all entries, in case of success

py::array_t<Py_ssize_t> counts_py(queries_count);

auto counts_py1d = counts_py.template mutable_unchecked<1>();

for (std::size_t query_idx = 0; query_idx != queries_count; ++query_idx)

counts_py1d(static_cast<Py_ssize_t>(query_idx)) = 1;

py::tuple results(5);

results[0] = keys_py;

results[1] = distances_py;

results[2] = counts_py;

results[3] = cluster_result.visited_members;

results[4] = cluster_result.computed_distances;

return results;

dense_index_py_t const& a, dense_index_py_t const& b, //

std::size_t max_proposals, bool exact, //

progress_func_t const& progress) {

std::unordered_map<dense_key_t, dense_key_t> a_to_b;

dummy_key_to_key_mapping_t b_to_a;

a_to_b.reserve((std::min)(a.size(), b.size()));

index_join_config_t config;

config.max_proposals = max_proposals;

config.exact = exact;

config.expansion = (std::max)(a.expansion_search(), b.expansion_search());

std::size_t threads = (std::min)(a.limits().threads(), b.limits().threads());

executor_default_t executor{threads};

join_result_t result;

{

// Lock the receiver `a`; `b` is read-only from this side. Concurrent

// bidirectional `join(a, b)` and `join(b, a)` from two Python threads

// is unsupported.

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*a.mutex_ptr_);

result = a.join(b, config, a_to_b, b_to_a, executor, progress_t{progress});

}

forward_error(result);

return a_to_b;

using copy_result_t = typename dense_index_py_t::copy_result_t;

index_dense_copy_config_t config;

config.force_vector_copy = force_copy;

copy_result_t result = index.copy(config);

forward_error(result);

return std::move(result.index);

if (!threads)

threads = std::thread::hardware_concurrency();

py::gil_scoped_release release;

std::unique_lock<std::mutex> lock(*index.mutex_ptr_);

if (!index.try_reserve(index_limits_t(index.size(), threads)))

throw std::invalid_argument("Out of memory!");

index.compact(executor_default_t{threads}, progress_t{progress});

py::dict result;

result["matrix_included"] = !meta.config.exclude_vectors;

result["matrix_uses_64_bit_dimensions"] = meta.config.use_64_bit_dimensions;

index_dense_head_t const& head = meta.head;

result["version"] = std::to_string(head.version_major) + "." + //

std::to_string(head.version_minor) + "." + //

std::to_string(head.version_patch);

result["kind_metric"] = metric_kind_t(head.kind_metric);

result["kind_scalar"] = scalar_kind_t(head.kind_scalar);

result["kind_key"] = scalar_kind_t(head.kind_key);

result["kind_compressed_slot"] = scalar_kind_t(head.kind_compressed_slot);

result["count_present"] = std::uint64_t(head.count_present);

result["count_deleted"] = std::uint64_t(head.count_deleted);

result["dimensions"] = std::uint64_t(head.dimensions);

return result;

index.load(path.c_str(), {}, progress_t{progress}).error.raise();

// Reserve memory and threads for restored index.

std::size_t threads = std::thread::hardware_concurrency();

if (!index.try_reserve(index_limits_t(index.size(), threads)))

throw std::invalid_argument("Out of memory!");

index.view(path.c_str(), 0, {}, progress_t{progress}).error.raise();

// Reserve memory and threads for restored index.

std::size_t threads = std::thread::hardware_concurrency();

if (!index.try_reserve(index_limits_t(index.size(), threads)))

throw std::invalid_argument("Out of memory!");

std::size_t serialized_length = index.serialized_length();

// Create an empty bytearray object using CPython API

PyObject* byte_array = PyByteArray_FromStringAndSize(nullptr, 0);

if (!byte_array)

throw std::runtime_error("Could not allocate bytearray object");

// Resize the bytearray object to the desired length