From aabc44da230b36587dbb17e395b8d2ae628d3a28 Mon Sep 17 00:00:00 2001 From: Chakshu Dhannawat Date: Fri, 24 Jul 2026 14:19:42 +0900 Subject: [PATCH] Fix: treat NumPy uint8 vectors as u8, not bit-packed binary A NumPy uint8 buffer is ambiguous: the same bytes back both bit-packed binary vectors (b1x8) and unsigned byte vectors (u8). numpy_string_to_kind maps uint8 to b1x8 for backwards compatibility, so when no explicit dtype was passed to add/search, a u8 index reinterpreted its data as packed bits and the vectors silently collapsed to (near-)zero. int8 worked because it maps to i8 unambiguously. Disambiguate against the index's own scalar kind: when the caller gave no dtype and the buffer resolves to b1x8 but the index is u8, use u8. Binary indexes are unaffected, so bit-packed uint8 vectors keep working. Applied at the add, search, and cluster dispatch sites through a shared helper, and added a scalar_kind() accessor to the multi-shard wrapper so it compiles for both index types. Adds a regression test that stores sparse uint8 counts, checks the round trip and search ranking, and confirms the binary path is untouched. The test fails on the current code (vector stored as all zeros) and passes with the fix. Closes #595 --- python/lib.cpp | 33 +++++++++++++++++++++-------- python/scripts/test_index.py | 41 ++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 9 deletions(-) diff --git a/python/lib.cpp b/python/lib.cpp index 601297d4..f8d897a8 100644 --- a/python/lib.cpp +++ b/python/lib.cpp @@ -97,6 +97,9 @@ struct dense_indexes_py_t { void merge(std::shared_ptr 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(); } + scalar_kind_t scalar_kind() const noexcept { + return shards_.empty() ? scalar_kind_t::unknown_k : shards_[0]->scalar_kind(); + } index_limits_t limits() const noexcept { return {size(), std::numeric_limits::max()}; } void merge_paths(std::vector const& paths, bool view = true, std::size_t threads = 0) { @@ -180,6 +183,24 @@ scalar_kind_t numpy_string_to_kind(std::string const& name) { return scalar_kind_t::unknown_k; } +/// @brief Resolves the scalar kind of a NumPy buffer for a specific index. +/// +/// NumPy `uint8` buffers are ambiguous: the same bytes back both bit-packed +/// binary vectors (`b1x8`) and unsigned byte vectors (`u8`), and +/// `numpy_string_to_kind` maps `uint8` to `b1x8` for backwards compatibility. +/// When the caller passed no explicit `scalar_kind`, disambiguate against the +/// index's own scalar kind, so a `u8` index isn't fed its data as packed bits +/// (which silently zeroed out the vectors). See issue #595. +template +scalar_kind_t resolve_buffer_kind(scalar_kind_t requested, py::buffer_info const& buffer_info, index_at const& index) { + if (requested != scalar_kind_t::unknown_k) + return requested; + scalar_kind_t detected = numpy_string_to_kind(buffer_info.format); + if (detected == scalar_kind_t::b1x8_k && index.scalar_kind() == scalar_kind_t::u8_k) + return scalar_kind_t::u8_k; + return detected; +} + template void forward_error(result_at&& result) { if (!result) @@ -286,9 +307,7 @@ static void add_many_to_index( // // 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); + scalar_kind_t kind = resolve_buffer_kind(scalar_kind, vectors_info, index); switch (kind) { case scalar_kind_t::f64_k: add_typed_to_index(index, keys_info, vectors_info, force_copy, threads, progress); break; case scalar_kind_t::f32_k: add_typed_to_index(index, keys_info, vectors_info, force_copy, threads, progress); break; @@ -501,9 +520,7 @@ static py::tuple search_many_in_index( // std::atomic 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); + scalar_kind_t kind = resolve_buffer_kind(scalar_kind, vectors_info, index); switch (kind) { case scalar_kind_t::f64_k: search_typed(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(index, vectors_info, wanted, exact, threads, keys_py, distances_py, counts_py, stats_visited_members, stats_computed_distances, progress); break; @@ -780,9 +797,7 @@ static py::tuple cluster_vectors( // rows_lookup_gt 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); + scalar_kind_t kind = resolve_buffer_kind(scalar_kind, queries_info, index); { py::gil_scoped_release release; std::unique_lock lock(*index.mutex_ptr_); diff --git a/python/scripts/test_index.py b/python/scripts/test_index.py index 4a339914..4b33c6af 100644 --- a/python/scripts/test_index.py +++ b/python/scripts/test_index.py @@ -165,6 +165,47 @@ def test_index_get_missing_keys(multi): assert index.get(1) is None +def test_u8_vectors_not_misread_as_binary(): + """`uint8` vectors must be stored as bytes, not bit-packed binary (#595). + + A NumPy `uint8` buffer is ambiguous: the same bytes back both bit-packed + binary vectors (`b1x8`) and unsigned byte vectors (`u8`). When no explicit + `dtype` is passed to `add`/`search`, the buffer used to resolve to `b1x8`, + so a `u8` index silently reinterpreted its data as packed bits and collapsed + the vectors to (near-)zero. The index's own scalar kind must disambiguate. + """ + reset_randomness() + ndim = 8 + # Sparse small counts, like the reporter's feature-count vectors. + vector = np.zeros(ndim, dtype=np.uint8) + vector[2] = 3 + vector[5] = 7 + + index = Index(ndim=ndim, metric=MetricKind.L2sq, dtype=ScalarKind.U8) + index.add(0, vector) # no explicit dtype: the previously-broken path + + # Round-trip: the stored bytes must match the input, not zeros. + stored = index.get(0, ScalarKind.U8) + assert np.array_equal(stored, vector), f"u8 vector corrupted on add: {stored}" + + # Search must be coherent: the exact vector is its own nearest neighbor at + # distance 0, and a far vector ranks strictly behind it. + far = np.zeros(ndim, dtype=np.uint8) + far[0] = 200 + far[7] = 200 + index.add(1, far) + matches = index.search(vector, 2) + assert matches.keys[0] == 0 + assert float(matches.distances[0]) == 0.0 + assert float(matches.distances[1]) > 0.0 + + # The fix must not touch binary indexes: uint8 still means bit-packed there. + binary = Index(ndim=64, metric=MetricKind.Hamming, dtype=ScalarKind.B1) + packed = np.array([0b10101010] * 8, dtype=np.uint8) + binary.add(0, packed) + assert float(binary.search(packed, 1).distances[0]) == 0.0 + + @pytest.mark.parametrize("ndim", [3, 97, 256]) @pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq]) @pytest.mark.parametrize("batch_size", [1, 7, 1024])