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])