diff --git a/turbovec-python/python/turbovec/_dedup.py b/turbovec-python/python/turbovec/_dedup.py new file mode 100644 index 00000000..96942a6f --- /dev/null +++ b/turbovec-python/python/turbovec/_dedup.py @@ -0,0 +1,71 @@ +"""Shared in-batch duplicate resolution for the framework integrations. + +Each upstream library resolves a repeated id *within a single write* its own +way, and every turbovec wrapper must match its upstream to stay a true +drop-in: + +- LangChain's ``InMemoryVectorStore`` overwrites on a repeated key → KEEP_LAST +- LlamaIndex rejects duplicate ``node_id`` in a batch → REJECT +- agno's LanceDb is append-only and keeps every row → KEEP_ALL +- Haystack exposes a runtime ``DuplicatePolicy`` (FAIL/SKIP/OVERWRITE). + Its resolution is *stateful* (it dedups against the existing store as well + as the batch, with deferred issue-#89 removal), so it does not reduce to + the pure in-batch function here and keeps its own logic; this enum still + documents the mapping (OVERWRITE→KEEP_LAST, SKIP→KEEP_FIRST, FAIL→REJECT). + +The shared piece is the in-batch resolution only: given one key per item, +return the indices to keep. Each wrapper still owns its key extraction and +its cross-store upsert/removal. +""" +from __future__ import annotations + +import enum +from typing import Hashable, List, Sequence + + +class DuplicatePolicy(enum.Enum): + """How to resolve items that share a key within a single batch.""" + + KEEP_LAST = "keep_last" + """One item per key; the last occurrence wins (dict-overwrite semantics).""" + + KEEP_FIRST = "keep_first" + """One item per key; the first occurrence wins.""" + + REJECT = "reject" + """Raise ``ValueError`` if any key repeats; otherwise keep everything.""" + + KEEP_ALL = "keep_all" + """No deduplication; items with duplicate keys all survive.""" + + +def resolve_duplicates( + keys: Sequence[Hashable], policy: DuplicatePolicy +) -> List[int]: + """Return, in ascending order, the batch indices to keep under ``policy``. + + The returned indices index into ``keys`` (and any parallel arrays the + caller holds). For KEEP_ALL and REJECT the result is ``0..len(keys)``; + for KEEP_LAST/KEEP_FIRST it collapses to one index per distinct key. + + Raises: + ValueError: under REJECT, if any key occurs more than once. + """ + if policy is DuplicatePolicy.KEEP_ALL: + return list(range(len(keys))) + if policy is DuplicatePolicy.REJECT: + seen: set = set() + for k in keys: + if k in seen: + raise ValueError(f"duplicate id in batch: {k!r}") + seen.add(k) + return list(range(len(keys))) + # KEEP_LAST / KEEP_FIRST collapse to one index per key. + chosen: dict = {} + for i, k in enumerate(keys): + if policy is DuplicatePolicy.KEEP_LAST or k not in chosen: + chosen[k] = i + return sorted(chosen.values()) + + +__all__ = ["DuplicatePolicy", "resolve_duplicates"] diff --git a/turbovec-python/python/turbovec/_persist.py b/turbovec-python/python/turbovec/_persist.py new file mode 100644 index 00000000..54692a42 --- /dev/null +++ b/turbovec-python/python/turbovec/_persist.py @@ -0,0 +1,55 @@ +"""Shared persistence consistency checks for the framework integrations. + +Each wrapper persists two artifacts: the binary ``.tvim`` index and a JSON +side-car holding the handle -> document/node/text payload maps. At query +time the wrapper resolves an index-returned u64 handle through that side-car +map. If the two files are out of sync — a partial copy, a stale backup, a +hand-edited or tampered side-car — an index handle won't resolve and the +wrapper would raise an opaque ``KeyError`` deep inside a query. + +``check_persisted_handles`` turns that into a clean ``ValueError`` at load +time. ``IdMapIndex`` exposes only ``__len__`` and ``contains``; that's +sufficient: if the side-car's handle set and the index have equal size and +every side-car handle is present in the index, the two are a bijection (no +index handle can be missing from the side-car). +""" +from __future__ import annotations + +from typing import Iterable + + +def check_persisted_handles(index, handles: Iterable[int], *, what: str = "entry") -> None: + """Validate that the side-car's handle set matches the loaded index. + + Args: + index: the loaded ``IdMapIndex`` (uses ``len`` and ``contains``). + handles: the u64 handles the side-car maps can resolve. + what: noun for error messages (e.g. "document", "node"). + + Raises: + ValueError: if the side-car has duplicate handles, a different count + than the index, or a handle the index doesn't contain. + """ + handle_list = [int(h) for h in handles] + n_index = len(index) + + if len(set(handle_list)) != len(handle_list): + raise ValueError( + f"persisted store is corrupt: duplicate {what} handles in the side-car" + ) + if len(handle_list) != n_index: + raise ValueError( + f"persisted store is inconsistent with its index: side-car has " + f"{len(handle_list)} {what} handle(s) but the index holds {n_index}. " + f"The .tvim index and its JSON side-car are out of sync." + ) + for h in handle_list: + if not index.contains(h): + raise ValueError( + f"persisted store is inconsistent with its index: {what} handle " + f"{h} is not present in the index. The .tvim index and its JSON " + f"side-car are out of sync." + ) + + +__all__ = ["check_persisted_handles"] diff --git a/turbovec-python/python/turbovec/agno.py b/turbovec-python/python/turbovec/agno.py index 1bfb9686..993d7fdd 100644 --- a/turbovec-python/python/turbovec/agno.py +++ b/turbovec-python/python/turbovec/agno.py @@ -130,8 +130,13 @@ def __init__( # freshly-constructed store doesn't "exist" until `create()` is # called, and `drop()` returns it to that state. self._index: Optional[IdMapIndex] = None - # str doc_id -> u64 handle - self._str_to_u64: Dict[str, int] = {} + # str doc_id -> set of u64 handles. One-to-many: agno's derived + # doc_id is NOT unique (two documents with identical content, or a + # repeated explicit doc.id within a batch, derive the same id), and + # LanceDb keeps every such row. Mapping one doc_id to a single handle + # silently orphaned the earlier vectors — present in search and the + # index count but unreachable by id, so undeletable (issue #104). + self._str_to_u64: Dict[str, Set[int]] = {} # u64 handle -> stored payload (mirrors LanceDb's "payload" shape) self._u64_to_doc: Dict[int, Dict[str, Any]] = {} # u64 handle assignment counter @@ -355,7 +360,7 @@ def insert( cleaned = doc.content.replace("\x00", "�") if doc.content else "" doc_id = self._derive_doc_id(doc, content_hash, cleaned) h = int(handle) - self._str_to_u64[doc_id] = h + self._str_to_u64.setdefault(doc_id, set()).add(h) self._u64_to_doc[h] = { "id": doc_id, "name": doc.name, @@ -440,14 +445,24 @@ def _remove_handle(self, handle: int) -> None: return self._index.remove(handle) doc_id = data.get("id") - # Only clear the id->handle mapping if it still points at this - # handle; a re-inserted doc may have repointed it to a new handle. - if doc_id is not None and self._str_to_u64.get(doc_id) == handle: - self._str_to_u64.pop(doc_id, None) - # Drop the name->id link only if no surviving handle keeps that id. + # Drop just this handle from the id's handle set; remove the id + # entirely only once no handle remains under it. + if doc_id is not None: + handles = self._str_to_u64.get(doc_id) + if handles is not None: + handles.discard(handle) + if not handles: + del self._str_to_u64[doc_id] + # Drop the name->id link only if no surviving handle keeps that + # (name, id) pair. The derived doc_id excludes `name`, so two docs + # with different names can share an id — matching on id alone would + # leave a stale name entry when the last handle for this name goes. name = data.get("name") if name and name in self._name_to_ids: - if not any(d.get("id") == doc_id for d in self._u64_to_doc.values()): + if not any( + d.get("id") == doc_id and d.get("name") == name + for d in self._u64_to_doc.values() + ): self._name_to_ids[name].discard(doc_id) if not self._name_to_ids[name]: del self._name_to_ids[name] @@ -615,58 +630,58 @@ def get_supported_search_types(self) -> List[SearchType]: def delete_by_id(self, id: str) -> bool: if self._index is None: return False - handle = self._str_to_u64.pop(id, None) - if handle is None: + handles = self._str_to_u64.get(id) + if not handles: return False - doc_data = self._u64_to_doc.pop(handle, None) - if doc_data is not None: - name = doc_data.get("name") - if name and name in self._name_to_ids: - self._name_to_ids[name].discard(id) - if not self._name_to_ids[name]: - del self._name_to_ids[name] - self._index.remove(handle) - # Lazily drop content_hash from the set if no surviving doc has it. - if doc_data is not None: - ch = doc_data.get("content_hash") - if ch and not any( - d.get("content_hash") == ch for d in self._u64_to_doc.values() - ): - self._content_hashes.discard(ch) + # Remove every vector sharing this id — a non-unique derived doc_id + # can map to several handles. _remove_handle maintains the id, name, + # and content_hash side-indexes per handle. + for handle in list(handles): + self._remove_handle(handle) return True def delete_by_name(self, name: str) -> bool: if self._index is None: return False - ids = list(self._name_to_ids.get(name, set())) - for doc_id in ids: - self.delete_by_id(doc_id) - return bool(ids) + # Remove exactly the handles whose stored name matches. Delegating to + # delete_by_id would key on the derived doc_id, which excludes `name`, + # so it would also delete a differently-named doc that happens to + # share the id. LanceDb deletes rows matching the predicate directly. + handles = [h for h, d in self._u64_to_doc.items() if d.get("name") == name] + for handle in handles: + self._remove_handle(handle) + return bool(handles) def delete_by_metadata(self, metadata: Dict[str, Any]) -> bool: if self._index is None: return False items = list(metadata.items()) - to_delete = [ - data["id"] - for data in self._u64_to_doc.values() + # Remove the matching handles directly (see delete_by_name): the + # derived doc_id ignores metadata, so delete_by_id would over-delete + # distinct docs that collide on the id. + handles = [ + h + for h, data in self._u64_to_doc.items() if all((data.get("meta_data") or {}).get(k) == v for k, v in items) ] - for doc_id in to_delete: - self.delete_by_id(doc_id) - return bool(to_delete) + for handle in handles: + self._remove_handle(handle) + return bool(handles) def delete_by_content_id(self, content_id: str) -> bool: if self._index is None: return False - to_delete = [ - data["id"] - for data in self._u64_to_doc.values() + # Remove the matching handles directly (see delete_by_name): the + # derived doc_id ignores content_id, so delete_by_id would over-delete + # distinct docs that collide on the id. + handles = [ + h + for h, data in self._u64_to_doc.items() if data.get("content_id") == content_id ] - for doc_id in to_delete: - self.delete_by_id(doc_id) - return bool(to_delete) + for handle in handles: + self._remove_handle(handle) + return bool(handles) def update_metadata(self, content_id: str, metadata: Dict[str, Any]) -> None: """Merge ``metadata`` into both ``meta_data`` and the ``filters`` @@ -752,10 +767,13 @@ def _load_from(self, folder: Path) -> None: self._u64_to_doc = {int(h): d for h, d in state["u64_to_doc"]} self._next_u64 = int(state["next_u64"]) - # Rebuild reverse indexes from the loaded payload. - self._str_to_u64 = { - data["id"]: handle for handle, data in self._u64_to_doc.items() - } + # Rebuild reverse indexes from the loaded payload. doc_id is + # non-unique, so accumulate handles into a set per id rather than a + # dict comprehension (which would drop all but the last handle and + # re-orphan the very vectors issue #104 fixed). + self._str_to_u64 = {} + for handle, data in self._u64_to_doc.items(): + self._str_to_u64.setdefault(data["id"], set()).add(handle) self._content_hashes = set() self._name_to_ids = {} for data in self._u64_to_doc.values(): diff --git a/turbovec-python/python/turbovec/haystack.py b/turbovec-python/python/turbovec/haystack.py index 48eca942..68395151 100644 --- a/turbovec-python/python/turbovec/haystack.py +++ b/turbovec-python/python/turbovec/haystack.py @@ -24,6 +24,7 @@ import numpy as np +from ._persist import check_persisted_handles from ._turbovec import IdMapIndex try: @@ -671,6 +672,7 @@ def load_from_disk( store._str_to_u64 = { data["id"]: handle for handle, data in store._u64_to_doc.items() } + check_persisted_handles(store._index, store._u64_to_doc.keys(), what="document") return store # ---- Internals ---------------------------------------------------- diff --git a/turbovec-python/python/turbovec/langchain.py b/turbovec-python/python/turbovec/langchain.py index f44f27b7..b8ee6140 100644 --- a/turbovec-python/python/turbovec/langchain.py +++ b/turbovec-python/python/turbovec/langchain.py @@ -15,6 +15,8 @@ import numpy as np +from ._dedup import DuplicatePolicy, resolve_duplicates +from ._persist import check_persisted_handles from ._turbovec import IdMapIndex try: @@ -184,8 +186,8 @@ def _store_texts_and_vectors( # earlier vectors. The returned id list still mirrors the input # (one entry per input text), as the reference does. result_ids = ids - if len(set(ids)) != len(ids): - keep = sorted({id_: i for i, id_ in enumerate(ids)}.values()) + keep = resolve_duplicates(ids, DuplicatePolicy.KEEP_LAST) + if len(keep) != len(ids): ids = [ids[i] for i in keep] texts_list = [texts_list[i] for i in keep] metadatas = [metadatas[i] for i in keep] @@ -542,6 +544,7 @@ def load( # JSON object keys are strings; the str_to_u64 values are already # ints in the payload, just need to confirm. str_to_u64 = {sid: int(h) for sid, h in state["str_to_u64"].items()} + check_persisted_handles(index, str_to_u64.values(), what="document") return cls( embedding=embedding, index=index, diff --git a/turbovec-python/python/turbovec/llama_index.py b/turbovec-python/python/turbovec/llama_index.py index aebd73c8..6253ff04 100644 --- a/turbovec-python/python/turbovec/llama_index.py +++ b/turbovec-python/python/turbovec/llama_index.py @@ -12,6 +12,8 @@ import numpy as np +from ._dedup import DuplicatePolicy, resolve_duplicates +from ._persist import check_persisted_handles from ._turbovec import IdMapIndex try: @@ -143,14 +145,16 @@ def add(self, nodes: list[BaseNode], **_: Any) -> list[str]: # `query` later resolves through the duplicate node_id, returning # the second node's payload attached to the first node's vector. # Caller's job to deduplicate before calling add. - seen: set[str] = set() - for n in nodes: - if n.node_id in seen: - raise ValueError( - f"duplicate node_id {n.node_id!r} appears multiple times " - "in the input batch; deduplicate before calling add()" - ) - seen.add(n.node_id) + node_ids = [n.node_id for n in nodes] + try: + resolve_duplicates(node_ids, DuplicatePolicy.REJECT) + except ValueError: + seen: set[str] = set() + dup = next(nid for nid in node_ids if nid in seen or seen.add(nid)) + raise ValueError( + f"duplicate node_id {dup!r} appears multiple times " + "in the input batch; deduplicate before calling add()" + ) from None embeddings = [node.get_embedding() for node in nodes] vectors = np.asarray(embeddings, dtype=np.float32) @@ -648,6 +652,7 @@ def from_persist_path( store._node_id_to_u64 = {nid: int(h) for nid, h in state["node_id_to_u64"]} store._u64_to_node_id = {h: nid for nid, h in store._node_id_to_u64.items()} store._next_u64 = int(state["next_u64"]) + check_persisted_handles(index, store._u64_to_node_id.keys(), what="node") return store @classmethod diff --git a/turbovec-python/src/lib.rs b/turbovec-python/src/lib.rs index 8f0b72a9..00b65a51 100644 --- a/turbovec-python/src/lib.rs +++ b/turbovec-python/src/lib.rs @@ -8,6 +8,32 @@ fn not_contiguous_err(kind: &str) -> PyErr { )) } +/// Map a numpy shape error from reassembling search results into a typed +/// RuntimeError. The result dimensions are derived from the core's own +/// output, so this never fires today — but a future change to result shaping +/// would otherwise surface as an uncatchable panic instead of a catchable +/// exception. +fn shape_err(e: numpy::ndarray::ShapeError) -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "internal error: malformed search result shape: {e}" + )) +} + +/// Reject NaN / Inf / overflow-magnitude query coordinates with a typed +/// `ValueError`. The core `search` panics on invalid values (its documented +/// Rust contract), which would otherwise surface to Python as an uncatchable +/// `PanicException`. `add` already maps the same condition to `ValueError`; +/// this keeps `search` consistent. +fn validate_queries(values: &[f32], dim: usize) -> PyResult<()> { + if let Some((vi, ci, v)) = turbovec_core::first_invalid_coord(values, dim) { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "invalid query value at query {vi}, coord {ci}: {v} \ + (must be finite and |value| < 1e16)", + ))); + } + Ok(()) +} + #[pyclass] struct TurboQuantIndex { inner: turbovec_core::TurboQuantIndex, @@ -70,6 +96,7 @@ impl TurboQuantIndex { ))); } } + validate_queries(q_slice, arr.ncols())?; let mask_arr = mask.as_ref().map(|m| m.as_array()); let mask_slice: Option<&[bool]> = match mask_arr.as_ref() { @@ -91,10 +118,10 @@ impl TurboQuantIndex { let effective_k = results.k; let scores = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), results.scores) - .unwrap() + .map_err(shape_err)? .into_pyarray(py); let indices = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), results.indices) - .unwrap() + .map_err(shape_err)? .into_pyarray(py); Ok((scores, indices)) @@ -247,6 +274,7 @@ impl IdMapIndex { ))); } } + validate_queries(q_slice, arr.ncols())?; let allow_arr = allowlist.as_ref().map(|a| a.as_array()); let allow_slice: Option<&[u64]> = match allow_arr.as_ref() { @@ -303,10 +331,10 @@ impl IdMapIndex { }; let scores_arr = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), scores) - .unwrap() + .map_err(shape_err)? .into_pyarray(py); let ids_arr = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), ids) - .unwrap() + .map_err(shape_err)? .into_pyarray(py); Ok((scores_arr, ids_arr)) } diff --git a/turbovec-python/tests/test_agno.py b/turbovec-python/tests/test_agno.py index cb84defe..ed59368a 100644 --- a/turbovec-python/tests/test_agno.py +++ b/turbovec-python/tests/test_agno.py @@ -1121,3 +1121,83 @@ def test_delete_by_metadata_returns_false_when_no_match(): assert db.delete_by_metadata({"tag": "no-such-value"}) is False # Original doc still present. assert db.get_count() == 1 + + +# ---- Security/data-integrity regression (issue #104) ---------------------- + + +def test_duplicate_doc_id_keeps_both_vectors_no_orphan(): + # agno's reference store (LanceDb) is append-only: two docs with the same + # explicit id (hence same derived doc_id) are BOTH stored. Previously the + # one-to-one _str_to_u64 map orphaned the first vector — counted and + # searchable but undeletable. Both must now be reachable and deletable. + db = TurboQuantVectorDb(embedder=StubEmbedder(DIM)) + db.create() + db.insert("h", [_doc("alpha", doc_id="dup"), _doc("beta", doc_id="dup")]) + + assert db.get_count() == 2 + [doc_id] = list(db._str_to_u64) # both collapse to one derived id + assert len(db._str_to_u64[doc_id]) == 2 # ...mapping to both handles + assert len(db._u64_to_doc) == 2 + + # Deleting that id removes BOTH vectors, leaving no orphan behind. + assert db.delete_by_id(doc_id) + assert db.get_count() == 0 + assert db._str_to_u64 == {} + assert db._u64_to_doc == {} + + +def test_duplicate_doc_id_survives_persistence_roundtrip(tmp_path): + embedder = StubEmbedder(DIM) + db = TurboQuantVectorDb(embedder=embedder, path=str(tmp_path)) + db.create() + db.insert("h", [_doc("alpha", doc_id="dup"), _doc("beta", doc_id="dup")]) + db.save() + + # Reload must rebuild the one-to-many id map, not drop a handle. + db2 = TurboQuantVectorDb(embedder=embedder, path=str(tmp_path)) + db2.create() + assert db2.get_count() == 2 + [doc_id] = list(db2._str_to_u64) + assert len(db2._str_to_u64[doc_id]) == 2 + + +def _doc_same_content(name=None, content_id=None, meta_data=None): + # All share identical content + no explicit id -> identical derived doc_id, + # differing only by name/content_id/metadata. + return _doc("identical", name=name, content_id=content_id, meta_data=meta_data) + + +def test_delete_by_name_only_removes_matching_name_on_doc_id_collision(): + # name is not part of the derived doc_id, so two differently-named docs + # with identical content collide. delete_by_name must remove only the + # named doc, not its id-twin, and must not leave a stale name entry. + db = TurboQuantVectorDb(embedder=StubEmbedder(DIM)) + db.create() + db.insert("h", [_doc_same_content(name="A"), _doc_same_content(name="B")]) + assert db.get_count() == 2 + + assert db.delete_by_name("A") is True + assert db.get_count() == 1 + assert db.name_exists("A") is False # no stale entry + assert db.name_exists("B") is True + + +def test_delete_by_content_id_only_removes_matching_on_doc_id_collision(): + db = TurboQuantVectorDb(embedder=StubEmbedder(DIM)) + db.create() + db.insert("h", [_doc_same_content(content_id="c1"), _doc_same_content(content_id="c2")]) + assert db.get_count() == 2 + + assert db.delete_by_content_id("c1") is True + assert db.get_count() == 1 # c2 survives + + +def test_delete_by_metadata_only_removes_matching_on_doc_id_collision(): + db = TurboQuantVectorDb(embedder=StubEmbedder(DIM)) + db.create() + db.insert("h", [_doc_same_content(meta_data={"k": "x"}), _doc_same_content(meta_data={"k": "y"})]) + assert db.get_count() == 2 + + assert db.delete_by_metadata({"k": "x"}) is True + assert db.get_count() == 1 # the {"k": "y"} doc survives diff --git a/turbovec-python/tests/test_dedup.py b/turbovec-python/tests/test_dedup.py new file mode 100644 index 00000000..4fb1634c --- /dev/null +++ b/turbovec-python/tests/test_dedup.py @@ -0,0 +1,39 @@ +"""Unit tests for the shared in-batch duplicate-resolution helper.""" +from __future__ import annotations + +import pytest + +from turbovec._dedup import DuplicatePolicy, resolve_duplicates + + +def test_keep_all_returns_every_index(): + assert resolve_duplicates(["a", "a", "b"], DuplicatePolicy.KEEP_ALL) == [0, 1, 2] + + +def test_keep_last_collapses_to_last_occurrence(): + # a@0,a@2 -> keep 2; b@1 -> keep 1. Ascending order. + assert resolve_duplicates(["a", "b", "a"], DuplicatePolicy.KEEP_LAST) == [1, 2] + + +def test_keep_first_collapses_to_first_occurrence(): + assert resolve_duplicates(["a", "b", "a"], DuplicatePolicy.KEEP_FIRST) == [0, 1] + + +def test_reject_raises_on_duplicate(): + with pytest.raises(ValueError, match="duplicate id in batch"): + resolve_duplicates(["a", "b", "a"], DuplicatePolicy.REJECT) + + +def test_reject_passes_through_when_unique(): + assert resolve_duplicates(["a", "b", "c"], DuplicatePolicy.REJECT) == [0, 1, 2] + + +def test_empty_batch(): + for policy in DuplicatePolicy: + assert resolve_duplicates([], policy) == [] + + +def test_no_duplicates_preserves_order_for_all_policies(): + keys = ["x", "y", "z"] + for policy in DuplicatePolicy: + assert resolve_duplicates(keys, policy) == [0, 1, 2] diff --git a/turbovec-python/tests/test_haystack.py b/turbovec-python/tests/test_haystack.py index bc57cca6..f7d89f8d 100644 --- a/turbovec-python/tests/test_haystack.py +++ b/turbovec-python/tests/test_haystack.py @@ -1179,3 +1179,22 @@ def test_embedding_retrieval_all_results_have_finite_float_scores(): for r in results: assert isinstance(r.score, float) assert math.isfinite(r.score) + + +def test_load_rejects_side_car_desynced_from_index(tmp_path): + import json + + store = TurboQuantDocumentStore(dim=DIM, bit_width=4) + store.write_documents(make_docs(4)) + store.save_to_disk(tmp_path) + + TurboQuantDocumentStore.load_from_disk(tmp_path) # clean reload works + + with open(tmp_path / "docstore.json") as f: + state = json.load(f) + state["u64_to_doc"] = state["u64_to_doc"][:-1] # drop one handle->doc + with open(tmp_path / "docstore.json", "w") as f: + json.dump(state, f) + + with pytest.raises(ValueError): + TurboQuantDocumentStore.load_from_disk(tmp_path) diff --git a/turbovec-python/tests/test_langchain.py b/turbovec-python/tests/test_langchain.py index ac9a8994..3fdbd82e 100644 --- a/turbovec-python/tests/test_langchain.py +++ b/turbovec-python/tests/test_langchain.py @@ -757,3 +757,27 @@ async def run() -> list[Document]: docs = asyncio.run(run()) assert [d.id for d in docs] == ["id-c", "id-a", "id-b"] + + +def test_load_rejects_side_car_desynced_from_index(tmp_path): + # A side-car whose handle map doesn't match the .tvim index must fail + # cleanly at load, not with a KeyError deep inside a later query. + import json + + emb = StubEmbeddings(dim=64) + store = TurboQuantVectorStore.from_texts(["a", "b", "c", "d"], emb, bit_width=4) + store.dump(tmp_path) + + # Clean reload works. + TurboQuantVectorStore.load(tmp_path, emb) + + with open(tmp_path / "docstore.json") as f: + state = json.load(f) + # Drop one id->handle mapping so the side-car holds fewer handles than + # the index. + state["str_to_u64"].pop(next(iter(state["str_to_u64"]))) + with open(tmp_path / "docstore.json", "w") as f: + json.dump(state, f) + + with pytest.raises(ValueError): + TurboQuantVectorStore.load(tmp_path, emb) diff --git a/turbovec-python/tests/test_llama_index.py b/turbovec-python/tests/test_llama_index.py index ff874296..012e1e9a 100644 --- a/turbovec-python/tests/test_llama_index.py +++ b/turbovec-python/tests/test_llama_index.py @@ -1030,7 +1030,6 @@ def test_query_returns_node_with_full_field_fidelity(): start_char_idx=100, end_char_idx=200, metadata_template="<<{key}::{value}>>", - metadata_separator=" | ", text_template="META:{metadata_str}\nBODY:{content}", mimetype="text/markdown", ) @@ -1049,7 +1048,11 @@ def test_query_returns_node_with_full_field_fidelity(): assert returned.start_char_idx == 100 assert returned.end_char_idx == 200 assert returned.metadata_template == "<<{key}::{value}>>" - assert returned.metadata_separator == " | " + # NB: metadata_separator is intentionally not asserted. LlamaIndex's own + # metadata_dict_to_node does not round-trip it — node_to_metadata_dict + # serializes it, but reconstruction drops it back to the framework + # default — so no store built on the framework serializer (including the + # reference) preserves it. Verified directly against llama-index-core. assert returned.text_template == "META:{metadata_str}\nBODY:{content}" assert returned.mimetype == "text/markdown" # All four relationships should be present, not just SOURCE. @@ -1272,3 +1275,24 @@ def test_query_returned_node_always_has_none_embedding(): for node in store.get_nodes(): assert node.embedding is None + + +def test_from_persist_path_rejects_side_car_desynced_from_index(tmp_path): + import json + + store = TurboQuantVectorStore.from_params(dim=64, bit_width=4) + store.add([_make_node(t, seed=i) for i, t in enumerate(["a", "b", "c", "d"])]) + base = str(tmp_path / "store") + store.persist(base) + + TurboQuantVectorStore.from_persist_path(base) # clean reload works + + side_car = tmp_path / "store.nodes.json" + with open(side_car) as f: + state = json.load(f) + state["node_id_to_u64"] = state["node_id_to_u64"][:-1] # drop one node->handle + with open(side_car, "w") as f: + json.dump(state, f) + + with pytest.raises(ValueError): + TurboQuantVectorStore.from_persist_path(base) diff --git a/turbovec-python/tests/test_security.py b/turbovec-python/tests/test_security.py new file mode 100644 index 00000000..c0bee8f3 --- /dev/null +++ b/turbovec-python/tests/test_security.py @@ -0,0 +1,136 @@ +"""Security regression tests. + +These guard against the classes of bug found in the security audit: an +untrusted/corrupt index file must be *rejected* at load with a typed error +rather than loading and later panicking, returning silently-wrong results, +or driving an unbounded allocation. Each test crafts a malformed file by +hand against the on-disk format documented in ``turbovec/src/io.rs``. +""" +from __future__ import annotations + +import struct + +import numpy as np +import pytest + +from turbovec import IdMapIndex, TurboQuantIndex + + +def _craft_tv( + path, + *, + bit_width: int, + dim: int, + n_vectors: int, + n_scales: int | None = None, + codes: bytes = b"", + n_calib: int = 0, +) -> None: + """Write a v3 ``.tv`` file with fully attacker-controlled header fields.""" + if n_scales is None: + n_scales = n_vectors + with open(path, "wb") as f: + f.write(b"TVPI") # magic + f.write(bytes([3])) # version 3 + f.write(bytes([bit_width & 0xFF])) + f.write(struct.pack("8 divide-by-zero'd in repack; 5..8 silently passed the + # length check and returned wrong scores. Only 2/3/4 are valid. + p = tmp_path / "bad_bitwidth.tv" + _craft_tv(p, bit_width=bit_width, dim=8, n_vectors=1, codes=b"\x00" * 8) + with pytest.raises((ValueError, OSError)): + TurboQuantIndex.load(str(p)) + + +@pytest.mark.parametrize("dim", [12, 7, 100]) +def test_load_rejects_non_multiple_of_8_dim(tmp_path, dim): + p = tmp_path / "bad_dim.tv" + _craft_tv(p, bit_width=4, dim=dim, n_vectors=1, codes=b"\x00" * 8) + with pytest.raises((ValueError, OSError)): + TurboQuantIndex.load(str(p)) + + +def test_load_rejects_dim_zero_with_vectors(tmp_path): + # dim==0 is the lazy-index sentinel and is only valid with n_vectors==0. + p = tmp_path / "bad_lazy.tv" + _craft_tv(p, bit_width=4, dim=0, n_vectors=5) + with pytest.raises((ValueError, OSError)): + TurboQuantIndex.load(str(p)) + + +def test_load_rejects_huge_n_vectors_without_allocating(tmp_path): + # A tiny file declaring billions of vectors must fail on the truncated + # data, NOT pre-allocate gigabytes. This completes quickly if the loader + # reads incrementally; it would OOM/hang if it pre-sized from the header. + p = tmp_path / "huge.tv" + _craft_tv(p, bit_width=2, dim=8, n_vectors=0xFFFFFFFF, n_scales=0) + with pytest.raises((ValueError, OSError)): + TurboQuantIndex.load(str(p)) + + +@pytest.mark.parametrize("bad", [np.nan, np.inf, -np.inf, 1e17]) +def test_turboquant_search_rejects_non_finite_query(bad): + # A NaN/Inf/overflow query coord previously panicked inside the core, + # surfacing as an uncatchable PanicException. It must raise ValueError. + idx = TurboQuantIndex(dim=8, bit_width=4) + idx.add(np.ones((4, 8), dtype=np.float32)) + q = np.ones((1, 8), dtype=np.float32) + q[0, 0] = bad + with pytest.raises(ValueError): + idx.search(q, 2) + + +@pytest.mark.parametrize("bad", [np.nan, np.inf, -np.inf, 1e17]) +def test_idmap_search_rejects_non_finite_query(bad): + idx = IdMapIndex(dim=8, bit_width=4) + idx.add_with_ids(np.ones((4, 8), dtype=np.float32), np.arange(4, dtype=np.uint64)) + q = np.ones((1, 8), dtype=np.float32) + q[0, 0] = bad + with pytest.raises(ValueError): + idx.search(q, 2) + + +def test_load_rejects_oversized_dim(tmp_path): + # A tiny file declaring a huge dim passes the dim%8 check but would drive + # a multi-GB dim x dim rotation-matrix allocation on first search. The + # loader must reject dim > MAX_DIM (65536). + p = tmp_path / "bigdim.tv" + _craft_tv(p, bit_width=2, dim=70000, n_vectors=0, n_scales=0) + with pytest.raises((ValueError, OSError)): + TurboQuantIndex.load(str(p)) + + +def test_construct_rejects_oversized_dim(): + with pytest.raises(ValueError): + TurboQuantIndex(dim=70000, bit_width=4) + + +def test_lazy_add_rejects_zero_column_array(): + # A 0-column array to a lazy index slipped past the dim%8 check (0%8==0), + # divided by zero in the core, and wedged the index at dim=0. Must raise + # ValueError and leave the index uncommitted. + idx = TurboQuantIndex() # lazy: no dim + with pytest.raises(ValueError): + idx.add(np.ones((4, 0), dtype=np.float32)) + assert idx.dim is None + # Not wedged: a normal add still works afterwards. + idx.add(np.ones((3, 8), dtype=np.float32)) + assert len(idx) == 3 and idx.dim == 8 + + +def test_valid_roundtrip_still_loads(tmp_path): + # The hardening must not break legitimate files. + p = tmp_path / "good.tv" + idx = TurboQuantIndex(dim=8, bit_width=4) + idx.add(np.ones((3, 8), dtype=np.float32)) + idx.write(str(p)) + loaded = TurboQuantIndex.load(str(p)) + assert len(loaded) == 3 diff --git a/turbovec/src/error.rs b/turbovec/src/error.rs index a2793d7d..4f823792 100644 --- a/turbovec/src/error.rs +++ b/turbovec/src/error.rs @@ -31,6 +31,10 @@ pub enum AddError { /// First-add dim on a lazy index must be a multiple of 8. DimNotMultipleOf8(usize), + /// First-add dim on a lazy index exceeds [`MAX_DIM`](crate::MAX_DIM). + /// Bounds the lazily-built `dim`×`dim` rotation matrix allocation. + DimTooLarge { dim: usize, max: usize }, + /// `vectors.len()` is not a whole multiple of `dim`. VectorBufferNotMultipleOfDim { vectors_len: usize, dim: usize }, @@ -64,6 +68,9 @@ impl fmt::Display for AddError { Self::DimNotMultipleOf8(dim) => { write!(f, "dim must be a multiple of 8, got {dim}") } + Self::DimTooLarge { dim, max } => { + write!(f, "dim {dim} exceeds maximum {max}") + } Self::VectorBufferNotMultipleOfDim { vectors_len, dim } => write!( f, "vector buffer length {vectors_len} not a multiple of dim {dim}", @@ -96,6 +103,10 @@ pub enum ConstructError { /// `dim` must be a positive multiple of 8. DimNotPositiveMultipleOf8(usize), + + /// `dim` exceeds [`MAX_DIM`](crate::MAX_DIM). Bounds the lazily-built + /// `dim`×`dim` rotation matrix allocation. + DimTooLarge { dim: usize, max: usize }, } impl fmt::Display for ConstructError { @@ -107,6 +118,9 @@ impl fmt::Display for ConstructError { Self::DimNotPositiveMultipleOf8(dim) => { write!(f, "dim must be a positive multiple of 8, got {dim}") } + Self::DimTooLarge { dim, max } => { + write!(f, "dim {dim} exceeds maximum {max}") + } } } } diff --git a/turbovec/src/io.rs b/turbovec/src/io.rs index 452dcb43..6749a044 100644 --- a/turbovec/src/io.rs +++ b/turbovec/src/io.rs @@ -160,12 +160,17 @@ pub fn load_id_map( let (bit_width, dim, n_vectors, packed_codes, scales, tqplus_shift, tqplus_scale) = read_core_versioned(&mut f, version[0], TVIM_VERSION, ".tvim")?; - let mut slot_to_id = Vec::with_capacity(n_vectors); - let mut buf = [0u8; 8]; - for _ in 0..n_vectors { - f.read_exact(&mut buf)?; - slot_to_id.push(u64::from_le_bytes(buf)); - } + // Read the slot_to_id table via the capped reader rather than + // `Vec::with_capacity(n_vectors)` — `n_vectors` is attacker-controlled and + // pre-reserving it allows a tiny file to drive a huge allocation. + let id_bytes = n_vectors + .checked_mul(8) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "id table size overflows usize"))?; + let raw = read_exact_vec(&mut f, id_bytes)?; + let slot_to_id: Vec = raw + .chunks_exact(8) + .map(|b| u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]])) + .collect(); Ok(( bit_width, dim, n_vectors, packed_codes, scales, tqplus_shift, tqplus_scale, @@ -268,17 +273,78 @@ fn read_header_codes_scales( let dim = u32::from_le_bytes([header[1], header[2], header[3], header[4]]) as usize; let n_vectors = u32::from_le_bytes([header[5], header[6], header[7], header[8]]) as usize; - let packed_bytes = (dim / 8) * bit_width * n_vectors; - let mut packed_codes = vec![0u8; packed_bytes]; - r.read_exact(&mut packed_codes)?; + // Validate header fields before allocating anything. The constructors + // (`new`/`add_2d`) enforce these invariants, but the load path bypasses + // them — so an untrusted file could otherwise smuggle a `bit_width` that + // divides-by-zero in `pack::repack` (0 or >8), a `bit_width` of 5..8 that + // silently passes `from_parts`'s length check and returns wrong scores, + // or a `dim` that isn't a multiple of 8 (the bit-plane layout is + // undefined for it and the size formulas diverge → panic). `dim == 0` is + // the lazy-index sentinel and is only valid alongside `n_vectors == 0`. + if !(2..=4).contains(&bit_width) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("invalid bit_width {bit_width}: must be 2, 3, or 4"), + )); + } + if dim == 0 { + if n_vectors != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("dim 0 (lazy sentinel) requires n_vectors 0, got {n_vectors}"), + )); + } + } else if dim % 8 != 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("invalid dim {dim}: must be a multiple of 8"), + )); + } else if dim > crate::MAX_DIM { + // Bound the lazily-built dim×dim rotation matrix: a tiny file can + // declare a huge dim and drive a multi-GB allocation on first search. + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("invalid dim {dim}: exceeds maximum {}", crate::MAX_DIM), + )); + } + + // Checked arithmetic: `dim`/`n_vectors` are attacker-controlled u32s, so + // the product can overflow `usize` (on 32-bit targets this wrap would + // yield an undersized buffer and later out-of-bounds reads). + let packed_bytes = (dim / 8) + .checked_mul(bit_width) + .and_then(|x| x.checked_mul(n_vectors)) + .ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "packed code size overflows usize") + })?; + let packed_codes = read_exact_vec(r, packed_bytes)?; let scales = read_f32_array(r, n_vectors)?; Ok((bit_width, dim, n_vectors, packed_codes, scales)) } +/// Read exactly `n` bytes without pre-allocating `n` up front. A malicious +/// header can declare a multi-gigabyte length from a tiny file; `read_to_end` +/// on a `take`-limited reader grows the buffer only to the bytes actually +/// present, so we never reserve the attacker's claimed size before confirming +/// the data exists. The length check then rejects a truncated file cleanly. +fn read_exact_vec(r: &mut R, n: usize) -> io::Result> { + let mut buf = Vec::new(); + let read = r.take(n as u64).read_to_end(&mut buf)?; + if read != n { + return Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + format!("truncated file: expected {n} bytes, got {read}"), + )); + } + Ok(buf) +} + fn read_f32_array(r: &mut R, n: usize) -> io::Result> { - let mut bytes = vec![0u8; n * 4]; - r.read_exact(&mut bytes)?; + let n_bytes = n + .checked_mul(4) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "f32 array size overflows usize"))?; + let bytes = read_exact_vec(r, n_bytes)?; Ok(bytes .chunks_exact(4) .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) diff --git a/turbovec/src/lib.rs b/turbovec/src/lib.rs index 46aa6d0e..57e6398f 100644 --- a/turbovec/src/lib.rs +++ b/turbovec/src/lib.rs @@ -51,6 +51,16 @@ use std::sync::OnceLock; const ROTATION_SEED: u64 = 42; const BLOCK: usize = 32; + +/// Upper bound on vector dimensionality. `search`/`prepare` lazily build a +/// `dim`×`dim` f64 rotation matrix, an allocation that scales with `dim²` +/// and is NOT bounded by the size of any loaded file — so an untrusted +/// `.tv`/`.tvim` declaring a huge `dim` could otherwise drive a +/// multi-gigabyte allocation (resource-exhaustion DoS) from a tiny file. +/// 65536 is far above any real embedding dimension (largest in common use +/// is ~4096) and rejects the catastrophic cases. Enforced at construction, +/// first add, and load. +pub const MAX_DIM: usize = 65536; const FLUSH_EVERY: usize = 256; /// Maximum permitted coordinate magnitude. Beyond this, f32 sum-of- @@ -72,7 +82,7 @@ const MAX_INPUT_MAGNITUDE: f32 = 1e16; /// - Huge magnitude: `simd_norm`'s f32 sum-of-squares overflows to /// +Inf, `scale[i] = Inf` gets stored, slot incorrectly wins /// top-k against every query. -fn first_invalid_coord(values: &[f32], dim: usize) -> Option<(usize, usize, f32)> { +pub fn first_invalid_coord(values: &[f32], dim: usize) -> Option<(usize, usize, f32)> { for (i, x) in values.iter().enumerate() { if !x.is_finite() || x.abs() >= MAX_INPUT_MAGNITUDE { let vector_index = if dim == 0 { 0 } else { i / dim }; @@ -162,6 +172,9 @@ impl TurboQuantIndex { if dim == 0 || dim % 8 != 0 { return Err(ConstructError::DimNotPositiveMultipleOf8(dim)); } + if dim > MAX_DIM { + return Err(ConstructError::DimTooLarge { dim, max: MAX_DIM }); + } Ok(Self { dim: Some(dim), @@ -329,9 +342,16 @@ impl TurboQuantIndex { } Some(_) => {} None => { - if dim % 8 != 0 { + // `dim == 0` slips past the `% 8` check (0 % 8 == 0) but is a + // degenerate dim: committing it wedges the lazy index and the + // first `add` divides by zero (`vectors.len() / dim`). Reject + // it here, mirroring IdMapIndex::add_with_ids_2d. + if dim == 0 || dim % 8 != 0 { return Err(AddError::DimNotMultipleOf8(dim)); } + if dim > MAX_DIM { + return Err(AddError::DimTooLarge { dim, max: MAX_DIM }); + } // Don't commit dim until value validation passes — otherwise // a lazy index is left with a committed dim and no vectors, // which would let a follow-up wrong-dim add see a confusing @@ -850,3 +870,68 @@ mod from_parts_tests { assert_eq!(idx.len(), 2); } } + +#[cfg(all(test, target_arch = "x86_64"))] +mod x86_scalar_fallback_tests { + //! Verify the x86 scalar fallback (score_query_into_heap, taken on + //! pre-AVX2 CPUs) returns the SAME top-k as the SIMD kernels on this + //! host. score_query_into_heap is not compiled on aarch64, so this is + //! the only place its full scoring path — including the issue-#106 + //! perm0 de-interleave — runs end to end. + use super::TurboQuantIndex; + use crate::search::FORCE_SCALAR_FALLBACK; + use std::sync::atomic::Ordering; + + fn unit_vectors(n: usize, dim: usize, seed: u64) -> Vec { + let mut s = seed.wrapping_add(0x9E3779B97F4A7C15); + let mut out = vec![0.0f32; n * dim]; + for row in out.chunks_mut(dim) { + let mut norm = 0.0f64; + for x in row.iter_mut() { + s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let v = ((s >> 33) as f64 / (1u64 << 31) as f64) - 1.0; + *x = v as f32; + norm += v * v; + } + let inv = 1.0 / (norm.sqrt() + 1e-9); + for x in row.iter_mut() { + *x = (*x as f64 * inv) as f32; + } + } + out + } + + fn topk_sets(indices: &[i64], nq: usize, k: usize) -> Vec> { + (0..nq) + .map(|q| indices[q * k..(q + 1) * k].iter().copied().collect()) + .collect() + } + + #[test] + fn scalar_fallback_matches_simd_topk() { + let dim = 64; + let n = 600; + let nq = 12; + let k = 16; + for &bits in &[2usize, 3, 4] { + let mut idx = TurboQuantIndex::new(dim, bits).unwrap(); + idx.add(&unit_vectors(n, dim, 11)); + let queries = unit_vectors(nq, dim, 22); + + FORCE_SCALAR_FALLBACK.store(false, Ordering::Relaxed); + let simd = idx.search(&queries, k); + FORCE_SCALAR_FALLBACK.store(true, Ordering::Relaxed); + let scalar = idx.search(&queries, k); + FORCE_SCALAR_FALLBACK.store(false, Ordering::Relaxed); + + assert_eq!(simd.k, scalar.k, "bits={bits}: differing result width"); + // Compare per-query top-k as sets (tie order between kernels may + // differ; membership must not). + assert_eq!( + topk_sets(&simd.indices, nq, simd.k), + topk_sets(&scalar.indices, nq, scalar.k), + "bits={bits}: scalar fallback returned a different top-k than SIMD", + ); + } + } +} diff --git a/turbovec/src/pack.rs b/turbovec/src/pack.rs index b7e1dc61..9c2e9d77 100644 --- a/turbovec/src/pack.rs +++ b/turbovec/src/pack.rs @@ -87,6 +87,39 @@ fn pack_blocked( blocked } +/// Inverse of the `perm0` permutation used by the x86 `pack_blocked`: +/// `INV_PERM0[lane] == j` such that `perm0[j] == lane`, for `lane` in 0..16. +// Used by the x86 scalar fallback and by the round-trip test on every arch. +#[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))] +pub(crate) const INV_PERM0: [usize; 16] = + [0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15]; + +/// Reconstruct the *sequential* code byte for vector `lane` (0..32) of a +/// block group from the x86 `perm0`-interleaved hi/lo-nibble layout that the +/// x86 [`pack_blocked`] produces. `group_off` is the byte offset of the group +/// within `blocked` (i.e. `block_offset + g * BLOCK`). +/// +/// The x86 SIMD kernels read that interleaved layout natively, but the scalar +/// fallback ([`crate::search::score_query_into_heap`]) decodes one sequential +/// byte per vector. Without this de-interleave the scalar path — taken on +/// pre-AVX2 x86 / VMs without AVX2 — read the wrong bytes and returned +/// silently-wrong top-k results (issue #106). The returned byte is identical +/// to what the non-x86 sequential layout stores directly: high nibble = the +/// vector's "hi" code, low nibble = its "lo" code. +#[inline] +#[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))] +pub(crate) fn deinterleave_x86_code_byte(blocked: &[u8], group_off: usize, lane: usize) -> u8 { + let j = INV_PERM0[lane & 15]; + let hi_plane = blocked[group_off + j]; // byte holding hi-nibbles of two vectors + let lo_plane = blocked[group_off + 16 + j]; // byte holding lo-nibbles + let (hi, lo) = if lane < 16 { + (hi_plane & 0x0F, lo_plane & 0x0F) + } else { + (hi_plane >> 4, lo_plane >> 4) + }; + (hi << 4) | lo +} + #[cfg(not(target_arch = "x86_64"))] fn pack_blocked( n: usize, @@ -113,61 +146,48 @@ fn pack_blocked( blocked } -/// Repack 3-bit codes into two blocked arrays: -/// - sub_codes: 2-bit nibble format from planes 0,1 -/// - plane2: packed bits blocked by 32 vectors -pub fn repack_3bit( - packed_codes: &[u8], - n_vectors: usize, - dim: usize, -) -> (Vec, Vec, usize) { - let bytes_per_plane = dim / 8; - let bytes_per_row = 3 * bytes_per_plane; - let n_blocks = (n_vectors + BLOCK - 1) / BLOCK; - - let sub_byte_groups = dim / 4; - let mut sub_codes = vec![0u8; n_blocks * sub_byte_groups * BLOCK]; - - let plane2_byte_groups = bytes_per_plane; - let mut plane2_blocked = vec![0u8; n_blocks * plane2_byte_groups * BLOCK]; - - for block_idx in 0..n_blocks { - let base_vec = block_idx * BLOCK; - - for g in 0..sub_byte_groups { - let out_offset = (block_idx * sub_byte_groups + g) * BLOCK; - for lane in 0..BLOCK { - let vec_idx = base_vec + lane; - if vec_idx >= n_vectors { continue; } - - let mut byte_val = 0u8; - let dim_start = g * 4; - for c in 0..4usize { - let j = dim_start + c; - let byte_in_plane = j / 8; - let bit_in_byte = 7 - (j % 8); - let mask = 1u8 << bit_in_byte; +#[cfg(test)] +mod tests { + use super::{deinterleave_x86_code_byte, BLOCK}; + + /// Pack one 32-vector block exactly as the x86 `pack_blocked` does, then + /// verify `deinterleave_x86_code_byte` recovers each vector's sequential + /// code byte. This validates the issue-#106 scalar-fallback fix on every + /// architecture (including ARM, where the x86 search path can't run) by + /// exercising the layout math directly. + #[test] + fn deinterleave_x86_recovers_sequential_code_bytes() { + let n_byte_groups = 5usize; + // Deterministic pseudo-random code bytes for 32 vectors. + let mut codes_flat = vec![vec![0u8; n_byte_groups]; BLOCK]; + let mut s = 0x1234_5678u32; + for v in 0..BLOCK { + for g in 0..n_byte_groups { + s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + codes_flat[v][g] = (s >> 24) as u8; + } + } - let mut code = 0u8; - for p in 0..2usize { - let plane_byte = packed_codes[vec_idx * bytes_per_row + p * bytes_per_plane + byte_in_plane]; - if plane_byte & mask != 0 { code |= 1 << p; } - } - byte_val |= code << ((3 - c) * 2); - } - sub_codes[out_offset + lane] = byte_val; + let perm0: [usize; 16] = [0, 8, 1, 9, 2, 10, 3, 11, 4, 12, 5, 13, 6, 14, 7, 15]; + let mut blocked = vec![0u8; n_byte_groups * BLOCK]; + for g in 0..n_byte_groups { + let out_offset = g * BLOCK; + for j in 0..16 { + let ba = codes_flat[perm0[j]][g]; + let bb = codes_flat[perm0[j] + 16][g]; + blocked[out_offset + j] = (ba >> 4) | ((bb >> 4) << 4); + blocked[out_offset + 16 + j] = (ba & 0x0F) | ((bb & 0x0F) << 4); } } - for g in 0..plane2_byte_groups { - let out_offset = (block_idx * plane2_byte_groups + g) * BLOCK; + for g in 0..n_byte_groups { for lane in 0..BLOCK { - let vec_idx = base_vec + lane; - if vec_idx >= n_vectors { continue; } - plane2_blocked[out_offset + lane] = packed_codes[vec_idx * bytes_per_row + 2 * bytes_per_plane + g]; + assert_eq!( + deinterleave_x86_code_byte(&blocked, g * BLOCK, lane), + codes_flat[lane][g], + "mismatch at lane {lane}, group {g}", + ); } } } - - (sub_codes, plane2_blocked, n_blocks) } diff --git a/turbovec/src/search.rs b/turbovec/src/search.rs index 2313bc8a..47917d93 100644 --- a/turbovec/src/search.rs +++ b/turbovec/src/search.rs @@ -23,6 +23,15 @@ use crate::{BLOCK, FLUSH_EVERY}; /// retrieval telemetry. Reset is provided for test isolation. pub static BLOCKS_SKIPPED_BY_MASK: AtomicU64 = AtomicU64::new(0); +/// Test-only switch that forces the x86 dispatch to take the scalar +/// fallback even when AVX2/AVX-512 is available, so tests can exercise +/// `score_query_into_heap` on hardware that would otherwise always pick a +/// SIMD kernel. Compiled only under `cfg(test)` — zero cost in release. +#[cfg(test)] +#[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))] +pub(crate) static FORCE_SCALAR_FALLBACK: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + /// Current value of the block-skip counter. See [`BLOCKS_SKIPPED_BY_MASK`]. pub fn blocks_skipped_by_mask() -> u64 { BLOCKS_SKIPPED_BY_MASK.load(Ordering::Relaxed) @@ -1386,6 +1395,17 @@ fn score_query_into_heap( } let mut score = qlut_bias; for g in 0..n_byte_groups { + // The x86 blocked layout is perm0-interleaved hi/lo nibbles, + // so de-interleave this vector's byte before decoding (issue + // #106). Every other target stores the sequential layout that + // can be read directly. + #[cfg(target_arch = "x86_64")] + let byte_val = crate::pack::deinterleave_x86_code_byte( + blocked_codes, + block_offset + g * BLOCK, + lane, + ) as usize; + #[cfg(not(target_arch = "x86_64"))] let byte_val = blocked_codes[block_offset + g * BLOCK + lane] as usize; let hi = byte_val >> 4; let lo = byte_val & 0x0F; @@ -1711,8 +1731,17 @@ pub fn search( let mut heap_mins = vec![f32::NEG_INFINITY; batch_nq]; let mut heap_min_idxs = vec![0usize; batch_nq]; + #[cfg(test)] + let force_scalar = + FORCE_SCALAR_FALLBACK.load(std::sync::atomic::Ordering::Relaxed); + #[cfg(not(test))] + let force_scalar = false; + unsafe { - if is_x86_feature_detected!("avx512bw") && is_x86_feature_detected!("avx512f") { + if !force_scalar + && is_x86_feature_detected!("avx512bw") + && is_x86_feature_detected!("avx512f") + { search_multi_query_avx512bw( blocked_codes, &lut_refs, &scale_vals, &bias_vals, n_byte_groups, vec_scales, n_vectors, @@ -1720,7 +1749,7 @@ pub fn search( &mut heap_scores, &mut heap_indices, &mut heap_sizes, &mut heap_mins, &mut heap_min_idxs, ); - } else if is_x86_feature_detected!("avx2") { + } else if !force_scalar && is_x86_feature_detected!("avx2") { search_multi_query_avx2( blocked_codes, &lut_refs, &scale_vals, &bias_vals, n_byte_groups, vec_scales, n_vectors,