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,