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

Author: chakshu-dhannawat

python/lib.cpp

@@ -97,6 +97,9 @@ struct dense_indexes_py_t {
     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(); }
+    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<std::size_t>::max()}; }
 
     void merge_paths(std::vector<std::string> 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 <typename index_at>
+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 <typename result_at> 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<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;
@@ -501,9 +520,7 @@ static py::tuple search_many_in_index( //
     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);
+    scalar_kind_t kind = resolve_buffer_kind(scalar_kind, vectors_info, index);
     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;
@@ -780,9 +797,7 @@ static py::tuple cluster_vectors(        //
     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);
+    scalar_kind_t kind = resolve_buffer_kind(scalar_kind, queries_info, index);
     {
         py::gil_scoped_release release;
         std::unique_lock<std::mutex> lock(*index.mutex_ptr_);

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