wehub/ryancodrai--turbovec
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
CI / Rust (macos-14) (push) Waiting to run
Details
CI / Rust (windows-latest) (push) Waiting to run
Details
CI / Python (macos-14) (push) Waiting to run
Details
CI / Python (windows-latest) (push) Waiting to run
Details
CI / Rust (ubuntu-latest) (push) Failing after 1s
Details
CI / Python (ubuntu-latest) (push) Failing after 0s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:23:17 +08:00
commit 0d29a02500
129 changed files with 24744 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.cargo/config.toml
+20
View File
@@ -0,0 +1,20 @@
# Target the x86-64-v3 microarchitecture level on all x86_64 builds.
#
# x86-64-v3 corresponds to the generic Haswell/Excavator feature set
# (AVX, AVX2, BMI1, BMI2, FMA, F16C, LZCNT, MOVBE, OSXSAVE) which has
# been ubiquitous since 2013 and is the minimum floor for running
# turbovec's AVX2 SIMD kernel anyway. Setting it here lets LLVM emit
# AVX2/FMA instructions in the surrounding non-intrinsic Rust code
# (rotation loops, LUT build, heap update) which measurably improves
# the non-hot-kernel paths.
#
# We do NOT set x86-64-v4 (AVX-512) here because the AVX-512 kernel is
# already runtime-dispatched via `is_x86_feature_detected!`, and
# raising the baseline to v4 would make the whole crate SIGILL on
# CPUs without AVX-512 even though the AVX2 fallback kernel would
# otherwise handle them fine.
#
# ARM / aarch64 builds are unaffected — they use their own default
# codegen target.
[target.'cfg(target_arch = "x86_64")']
rustflags = ["-C", "target-cpu=x86-64-v3"]
.claude/settings.json
+5
View File
@@ -0,0 +1,5 @@
{
"worktree": {
"bgIsolation": "none"
}
}
.github/CODEOWNERS
+4
View File
@@ -0,0 +1,4 @@
# All changes require review from @RyanCodrai before they can land on main.
# Combined with branch protection on main (require review from Code Owners),
# this means only @RyanCodrai can approve a merge to main.
* @RyanCodrai
.github/PULL_REQUEST_TEMPLATE.md
+30
View File
@@ -0,0 +1,30 @@
<!--
Thanks for the contribution!
The expected workflow is: open an issue describing the change and your
proposed approach, get a 👍 or design discussion, then open this PR
referencing the issue. See CONTRIBUTING.md for the narrow exceptions
(typos, one-line obvious bug fixes, docs-only changes).
-->
## Related issue
Closes #<!-- issue number -->
## Summary
<!-- Bullets describing what changed. Short is good. -->
## Motivation
<!-- Why this change? If the issue already covers this, a one-line pointer is fine. -->
## Test plan
<!--
What you ran to verify. Check the ones you actually verified; add others as relevant.
For recall or speed changes, include before/after numbers.
-->
- [ ] `cargo test -p turbovec --release` passes
- [ ] `pytest turbovec-python/tests/` passes
.github/workflows/ci.yml
+86
View File
@@ -0,0 +1,86 @@
name: CI
on:
pull_request:
push:
branches: [main]
# Cancel any in-progress run on the same ref when a new commit lands —
# saves time when iterating on a PR via force-push.
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
rust:
name: Rust (${{ matrix.os }})
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, macos-14, windows-latest]
steps:
- uses: actions/checkout@v4
- name: Install OpenBLAS
if: runner.os == 'Linux'
run: sudo apt-get update && sudo apt-get install -y libopenblas-dev pkg-config
- name: turbovec test suite
run: cargo test -p turbovec --release
# Standalone cargo project that path-deps turbovec — exercises the same
# link path a downstream `cargo add turbovec` user hits. Catches BLAS
# link-propagation regressions of the class fixed in #69.
- name: Downstream consumer smoke test
run: cargo run --release --manifest-path examples/downstream-smoke/Cargo.toml
python:
name: Python (${{ matrix.os }})
runs-on: ${{ matrix.os }}
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, macos-14, windows-latest]
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.11"
- name: Install OpenBLAS
if: runner.os == 'Linux'
run: sudo apt-get update && sudo apt-get install -y libopenblas-dev pkg-config
- name: Install maturin and pytest
shell: bash
run: |
python -m pip install --upgrade pip
python -m pip install maturin pytest
- name: Build turbovec wheel
shell: bash
working-directory: turbovec-python
run: maturin build --release --out dist
# Install the freshly-built wheel + the four integration extras so the
# haystack / langchain / llama-index / agno test files run instead of
# being skipped by their `pytest.importorskip(...)` guards.
- name: Install wheel + integration extras
shell: bash
working-directory: turbovec-python
run: |
python -m pip install dist/*.whl
python -m pip install \
"langchain-core>=0.3" \
"llama-index-core>=0.11" \
"haystack-ai>=2.0" \
"agno>=2.0"
- name: Run pytest
shell: bash
run: pytest turbovec-python/tests/ -v
.github/workflows/release-crates.yml
+32
View File
@@ -0,0 +1,32 @@
name: Release (crates.io)
on:
push:
tags:
- "v*"
workflow_dispatch:
permissions:
contents: read
jobs:
crates-io:
name: Publish to crates.io
runs-on: ubuntu-latest
if: startsWith(github.ref, 'refs/tags/v')
environment:
name: crates-io
url: https://crates.io/crates/turbovec
permissions:
id-token: write
steps:
- uses: actions/checkout@v4
- name: Install openblas for cargo publish verification
run: sudo apt-get update && sudo apt-get install -y libopenblas-dev pkg-config
- name: Authenticate with crates.io
id: auth
uses: rust-lang/crates-io-auth-action@v1
- name: Publish turbovec
env:
CARGO_REGISTRY_TOKEN: ${{ steps.auth.outputs.token }}
run: cargo publish -p turbovec
.github/workflows/release-pypi.yml
+147
View File
@@ -0,0 +1,147 @@
name: Release (PyPI)
on:
push:
tags:
- "py-v*"
workflow_dispatch:
permissions:
contents: read
jobs:
linux:
name: Linux ${{ matrix.target }}
# Build each target on a runner of its own architecture, so the build
# is native and the in-container `apt-get install libopenblas-dev`
# picks up the right openblas variant without cross-compile gymnastics.
runs-on: ${{ matrix.target == 'aarch64' && 'ubuntu-24.04-arm' || 'ubuntu-latest' }}
strategy:
fail-fast: false
matrix:
target: [x86_64, aarch64]
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.11"
- name: Build wheels
uses: PyO3/maturin-action@v1
with:
target: ${{ matrix.target }}
manylinux: manylinux_2_28
args: --release --out dist
working-directory: turbovec-python
before-script-linux: |
if command -v dnf >/dev/null 2>&1; then
dnf install -y openblas-devel openssl-devel pkgconfig
elif command -v yum >/dev/null 2>&1; then
yum install -y openblas-devel openssl-devel pkgconfig
elif command -v apt-get >/dev/null 2>&1; then
apt-get update && apt-get install -y libopenblas-dev libssl-dev pkg-config
else
echo "No supported package manager found" && exit 1
fi
# Install the freshly-built wheel and run the core test suite on
# the runner (outside the build container). This catches linkage
# regressions that produce structurally-correct .so files which
# fail at Python import time — exactly the class of bug that was
# silently shipping in Linux wheels until this PR.
- name: Install wheel and smoke test
shell: bash
working-directory: turbovec-python
run: |
python -m pip install --upgrade pip
python -m pip install pytest
python -m pip install dist/*.whl
python -m pytest tests/test_index.py tests/test_id_map.py tests/test_filtering.py -v
- name: Upload wheels
uses: actions/upload-artifact@v4
with:
name: wheels-linux-${{ matrix.target }}
path: turbovec-python/dist
macos:
name: macOS aarch64
runs-on: macos-14
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.11"
- name: Build wheels
uses: PyO3/maturin-action@v1
with:
target: aarch64
args: --release --out dist
working-directory: turbovec-python
- name: Upload wheels
uses: actions/upload-artifact@v4
with:
name: wheels-macos-aarch64
path: turbovec-python/dist
windows:
name: Windows x64
runs-on: windows-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.11"
- name: Build wheels
uses: PyO3/maturin-action@v1
with:
target: x64
args: --release --out dist
working-directory: turbovec-python
- name: Install wheel and run tests
shell: bash
working-directory: turbovec-python
run: |
python -m pip install --upgrade pip
python -m pip install pytest
python -m pip install dist/*.whl
python -m pytest tests/test_index.py tests/test_id_map.py tests/test_filtering.py -v
- name: Upload wheels
uses: actions/upload-artifact@v4
with:
name: wheels-windows-x64
path: turbovec-python/dist
sdist:
name: sdist
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Build sdist
uses: PyO3/maturin-action@v1
with:
command: sdist
args: --out dist
working-directory: turbovec-python
- name: Upload sdist
uses: actions/upload-artifact@v4
with:
name: sdist
path: turbovec-python/dist
release:
name: Publish to PyPI
runs-on: ubuntu-latest
needs: [linux, macos, windows, sdist]
if: startsWith(github.ref, 'refs/tags/py-v')
environment:
name: pypi
url: https://pypi.org/project/turbovec
permissions:
id-token: write
steps:
- uses: actions/download-artifact@v4
with:
path: dist
merge-multiple: true
- name: Publish to PyPI
uses: pypa/gh-action-pypi-publish@release/v1
with:
packages-dir: dist
.gitignore
+13
View File
@@ -0,0 +1,13 @@
__pycache__/
.cache/
.venv/
*.pyc
*.bak
*.so
*.dSYM/
*.tq
*.npy
codebooks.npz
target/
evolve/
evolve_rust/
CHANGELOG.md
+873
View File
@@ -0,0 +1,873 @@
# Changelog
All notable changes to turbovec are recorded here. The format is based on
[Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project follows
[Semantic Versioning](https://semver.org/spec/v2.0.0.html).
The Rust crate (`turbovec` on crates.io) and the Python distribution
(`turbovec` on PyPI) version independently. Each release section below
is split by surface — a single feature can affect both, and its bullet
appears under each surface it touches.
## [Unreleased]
## turbovec 0.8.0 (Python package) + turbovec 0.9.0 (Rust crate) — 2026-06-10
Security-audit release. Two adversarial audit passes over the core crate,
the Python binding, and the framework integrations, hardening the
untrusted-file load path and the Python API surface and fixing several
data-integrity bugs in the integration wrappers. Resolves #104, #105, and
#106. No on-disk format change (still `.tv` / `.tvim` v3).
A few fixes change observable behavior — see **Changed** under each surface.
They turn previously-undefined or silently-wrong situations into clean,
typed errors, so the bump is minor rather than patch.
### turbovec — Rust crate (current: 0.8.1 → next: 0.9.0)
#### Fixed
- **Untrusted index files are validated before allocation on load.** A
crafted `.tv` / `.tvim` could previously trigger an integer-overflow in
the packed-size computation, drive a multi-gigabyte allocation from a
tiny file, divide-by-zero in the repack step, or load a structurally
invalid index that returned silently-wrong scores (a `bit_width` of 5–8
passed the old length check). The loader now range-checks `bit_width` and
`dim`, computes every size with checked arithmetic, and reads each payload
through a length-capped reader. (#105)
- **x86 scalar fallback returned wrong results.** On pre-AVX2 x86 (or VMs
without AVX2), `score_query_into_heap` read the perm0-interleaved code
layout as if sequential, producing an incorrect top-k. It now
de-interleaves correctly; verified end-to-end against the SIMD kernels on
AVX-512 hardware. (#106)
#### Changed
- **`AddError` and `ConstructError` are now `#[non_exhaustive]`.** Downstream
`match` on these enums must carry a wildcard arm; in exchange, future error
variants are no longer breaking changes. (The new `DimTooLarge` variant is
why this release is the moment to make the switch.)
- **`dim` is capped at `MAX_DIM` (65536)** at construction, first add, and
load. `search` lazily builds a `dim`×`dim` rotation matrix whose size is
not bounded by any file, so an unbounded `dim` was a load-time
resource-exhaustion vector. Larger dims now return a typed error.
- **A zero-`dim` lazy add is rejected** with `AddError` instead of panicking
with a divide-by-zero and wedging the index.
#### Removed
- Dead, untested `pack::repack_3bit` (no callers; 3-bit goes through
`repack`).
#### Other
- The crate now fails to compile on non-64-bit targets (a `compile_error!`
gated on `target_pointer_width`). The size/offset arithmetic in
`encode`/`pack`/`search` assumes 64-bit `usize`; refusing to build on
32-bit/wasm avoids shipping a silently-unsafe (potential out-of-bounds)
build there.
### turbovec — Python package (current: 0.7.1 → next: 0.8.0)
#### Fixed
- **`search()` no longer panics on NaN / Inf / oversized query
coordinates.** These previously raised an uncatchable `PanicException`
(a `BaseException`); they now raise `ValueError`, matching `add`. (#105)
- **Loading a malformed `.tv` / `.tvim` raises a clean error** instead of
panicking or driving a huge allocation (the Rust load-path hardening
above, surfaced through the binding).
- **agno: duplicate derived `doc_id` no longer orphans vectors.** Two
documents that derive the same id (a repeated `doc.id`, or identical
content) are both kept and both deletable, matching agno's reference
store (LanceDb appends). Previously the earlier vector was counted and
searchable but unreachable by id, and leaked on every upsert. (#104)
- **agno: `delete_by_name` / `delete_by_content_id` / `delete_by_metadata`
no longer over-delete.** When distinct documents collided on a
content-derived `doc_id`, deleting by one attribute also deleted the
id-twin; deletion now targets only the documents matching the predicate.
- **LangChain / Haystack / LlamaIndex: a persisted JSON side-car that is out
of sync with its `.tvim` index now raises a `ValueError` at load** instead
of an opaque `KeyError` deep inside a later query.
- **Internal binding result-shape errors map to `RuntimeError`** rather than
an uncatchable panic.
#### Changed
- `search()` and the index constructors now raise `ValueError` for
non-finite query values and for `dim` outside the supported range, where
some of these inputs previously panicked or were silently accepted.
### Docs
- Corrected stale benchmark figures in the README (recall deltas, ARM/x86
speed ranges) to match the current `benchmarks/results/`; several had
drifted from before the TQ+ calibration step landed.
## turbovec 0.7.1 (Python package) + turbovec 0.8.1 (Rust crate) — 2026-06-09
Bug-fix release. Two data-safety fixes in the Python integration wrappers'
add/upsert paths, plus a source-build fix for the Python extension on macOS.
The Rust crate is functionally unchanged — only non-behavioral cleanups —
but is re-released to keep crates.io in sync with the source tree. No
on-disk format change (still `.tv` / `.tvim` v3).
### turbovec — Rust crate (current: 0.8.0 → next: 0.8.1)
#### Changed
- Internal cleanup only, **no behavior change** — the published crate
behaves identically to 0.8.0. Cleared three build warnings (two unused
bindings in the NEON scoring kernels; the scalar `score_query_into_heap`
fallback is now `cfg`-gated out of `aarch64` builds, where the NEON
kernel is always used and it was dead code) and corrected stale SIMD
module/kernel doc comments.
### turbovec — Python package (current: 0.7.0 → next: 0.7.1)
#### Fixed
- **Intra-batch duplicate ids no longer orphan vectors** in the LangChain
and Haystack integrations. A repeated id within a single `add_texts` /
`add_documents` / `write_documents` call previously added one vector per
row while the id→handle map kept only the last, leaving the earlier
vectors live in search but mapped to the wrong document and unreachable
for delete. Both now resolve duplicates the way their reference stores
do — LangChain (`InMemoryVectorStore`) keeps the last occurrence;
Haystack (`InMemoryDocumentStore`) applies the `DuplicatePolicy` against
the batch-so-far. Fixes #90.
- **Upsert no longer destroys existing data when the new batch fails
validation**, across all four integrations (LangChain, LlamaIndex,
Haystack, Agno). The old vectors for matching ids were deleted *before*
the incoming batch was validated/encoded, so a dimension change or a
non-finite embedding left the store with the originals already gone. The
delete is now deferred until after the add succeeds (Agno captures the
previous generation's handles and removes them after `insert`). Fixes
#89.
- **Plain `cargo build` of the extension now links on macOS.** Building
`turbovec-python` from source failed with "symbol(s) not found for
architecture arm64" because nothing emitted the Python extension-module
linker args (maturin injects them; a bare `cargo build` did not). Added a
`build.rs` calling `pyo3_build_config::add_extension_module_link_args()`.
Prebuilt wheels were unaffected. Fixes #92.
## turbovec 0.7.0 (Python package) + turbovec 0.8.0 (Rust crate) — 2026-05-30
Audit-driven correctness pass on every layer (Rust core, Python bindings,
four integration wrappers). Headline: 14 active bugs found and fixed,
hundreds of regression tests added, doc drift cleaned up across the
public API. No on-disk format change (still `.tv` / `.tvim` v3).
### turbovec — Rust crate (current: 0.7.0 → next: 0.8.0)
#### Added
- **`AddError::InvalidInputValue { vector_index, coord_index, value }`** —
new error variant returned by `TurboQuantIndex::add_2d` and
`IdMapIndex::add_with_ids_2d` when an input coordinate is non-finite
(NaN, +Inf, -Inf) or has magnitude `>= 1e16`. Without this validation
the encode pipeline silently corrupted the index: NaN/Inf propagated
through `simd_norm` and poisoned `vec_scales[slot] = NaN`, making the
slot exist in `len()` but unreachable through `search`; huge magnitudes
overflowed the f32 norm to `+Inf`, making the slot win every query.
- **Scalar fallback in the x86_64 search dispatch.** Previously, `search`
on an x86_64 CPU without AVX-512 BW or AVX2 silently returned empty
top-k results for every query (the SIMD `unsafe { if/else if }` block
had no `else`). Rare in practice on modern hardware but the failure
mode was the worst kind.
#### Changed
- **Breaking**: `AddError` no longer derives `Eq` (the new
`InvalidInputValue` variant carries an `f32`, which isn't `Eq` because
`NaN != NaN`). `PartialEq` is still derived. Downstream code that
pattern-matches `AddError` exhaustively will need to add the new
variant.
- `TurboQuantIndex::add` / `add_2d` / `search` / `search_with_mask` now
reject non-finite / huge-magnitude inputs at entry. `add` and `search`
panic with a clear message (matching their existing precondition-
panic style); `add_2d` and `add_with_ids_2d` return
`Err(InvalidInputValue)` for callers handling untrusted input.
- `TurboQuantIndex::from_parts` asserts structural invariants
(packed_codes / scales / TQ+ length relationships) at entry, catching
any future caller that bypasses the read-layer validation.
- Rustdoc on `add`, `add_2d`, `search`, `search_with_mask`, and
`IdMapIndex::add_with_ids` now documents every panic condition
introduced by the input validation.
#### Fixed
- **`IdMapIndex::add_with_ids_2d` partial-mutation on inner failure.**
ID tables (`id_to_slot` / `slot_to_id`) were mutated before delegating
to the inner `add_2d`. If the inner call returned `Err` (e.g.
`DimMismatch` on a committed-dim index), the ID tables retained `n`
ghost entries pointing at slots that didn't exist in the inner index —
corrupting later `search_with_allowlist` / `remove`. Fixed by capturing
`base_slot` before, running inner add first, mutating ID tables only
on success.
- **v2-loaded index + `add` silently mis-encoded new vectors.** Loading a
pre-TQ+ (v2) file left `tqplus_shift` empty; the next `add` saw
`existing = None`, fit fresh calibration, encoded the new batch with
that calibration — but then silently dropped the fitted shift/scale
because the `n_vectors != 0` else branch only extended `packed_codes`
/ `scales`. The new vectors then got searched against identity
calibration, producing silently wrong scores. Fixed by populating
explicit identity TQ+ in `from_parts` when loading a v2-shaped state.
- **Empty first add froze identity calibration forever.** `add(&[])`
hit the `n < TQPLUS_MIN_SAMPLES` branch in `encode`, returned
identity, and the `n_vectors == 0` branch wrote it to
`self.tqplus_shift`. Every subsequent add — even a million-vector
batch with rich distribution — then saw `existing = Some(identity)`
and silently skipped fresh fitting. Fixed by short-circuiting `add`
to a true no-op when `n == 0`.
### turbovec — Python package (current: 0.6.0 → next: 0.7.0)
#### Changed
- **Breaking** (typed-exception hygiene): `TurboQuantIndex.add` /
`search` and `IdMapIndex.add_with_ids` / `search` now raise
`ValueError` for non-finite or huge-magnitude coordinates, non-
contiguous numpy arrays, and wrong-dim queries. Previously these
surfaced as Rust panics → `PanicException` in Python.
- **Breaking**: `TurboQuantIndex.swap_remove` now raises `IndexError`
for out-of-range indices (was a Rust panic).
- `IdMapIndex.search` and `TurboQuantIndex.search` now return consistent
shapes for empty queries — `(0, min(k, n_vectors, n_allowed))` on
both. Previously `IdMapIndex` returned `(0, k)` (raw `k`), diverging
from `TurboQuantIndex`'s `(0, min(k, n))`. For `IdMapIndex`, the
`effective_k` computation also now dedups the allowlist for the
`nq == 0` path, matching the kernel's mask-based dedup for `nq > 0`.
#### Fixed
- **`turbovec.langchain.TurboQuantVectorStore`**: `similarity_search`,
`similarity_search_with_score`, and `similarity_search_by_vector` now
populate `Document.id` on returned hits (was silently `None`). The
`Document` passed to user-supplied filter callables also carries
`.id` so predicates can filter on it. Fixes #81.
- **`turbovec.haystack.TurboQuantDocumentStore`**: `Document.blob` and
`Document.sparse_embedding` now survive write → retrieval round-trip
(were silently dropped). Docstore schema bumps `v1 → v2` with
backward-compat load. Filter shape validation tightened to match
`InMemoryDocumentStore` (bare `{"field": "x"}` shapes are rejected).
Docstring scoped back from "matches the public surface of
`InMemoryDocumentStore`" since `bm25_retrieval` is not implemented.
- **`turbovec.llama_index.TurboQuantVectorStore`**: full `BaseNode`
fidelity through `query` / `get_nodes` / persist+load. PREVIOUS /
NEXT / PARENT / CHILD relationships, `excluded_embed_metadata_keys` /
`excluded_llm_metadata_keys`, template fields (`text_template`,
`metadata_template`, `metadata_separator`), `start_char_idx` /
`end_char_idx`, and `mimetype` were silently dropped — now preserved
via `node_to_metadata_dict` / `metadata_dict_to_node`. Nodes schema
bumps `v1 → v2` with backward-compat load. Plus:
- `FilterCondition.NOT` now supported (was `NotImplementedError`).
- `FilterOperator.TEXT_MATCH` is now case-sensitive (matches the
reference; previously silently lowercased both sides).
- `FilterOperator.TEXT_MATCH_INSENSITIVE`, `ALL`, `ANY` added.
- `query.mode != VectorStoreQueryMode.DEFAULT` raises
`NotImplementedError` instead of silently behaving as DEFAULT.
- `add()` rejects intra-batch duplicate `node_id`s with a clear
`ValueError` (previously, the index ended up with both vectors but
only the last node_id mapped back to one, orphaning the first handle
and silently returning the second node's payload attached to the
first node's vector).
- **`turbovec.agno.TurboQuantVectorDb`**: `embedder` is now threaded
through returned `Document` objects so `doc.embed()` / `doc.async_embed()`
work on retrieved hits (previously raised "No embedder provided").
Empty query strings short-circuit to `[]` (matching LanceDb).
## turbovec 0.6.0 (Python package) + turbovec 0.7.0 (Rust crate) — 2026-05-27
### turbovec — Rust crate (current: 0.6.0 → next: 0.7.0)
#### Added
- **TQ+ per-coordinate calibration.** Before the data-oblivious rotation,
every coordinate is shifted by its empirical 5th percentile and scaled
so that the 5–95% range maps to `[0, 1]`. The shift/scale pair is
fit incrementally from the cold-path `add` data, so the index stays
online — no separate train pass, no rebuilds as the corpus grows.
At search time, the same affine is applied to the query before the
rotation. Recall@1 lifts across published cells:
- GloVe-200 4-bit: 0.8440 → 0.8498 (+0.6pp)
- OpenAI-1536 2-bit: 0.876 → 0.891 (+1.5pp)
- OpenAI-1536 4-bit: 0.966 → 0.974 (+0.8pp)
- OpenAI-3072 2-bit: 0.911 → 0.929 (+1.8pp)
- OpenAI-3072 4-bit: 0.971 → 0.974 (+0.3pp)
No public API change — TQ+ is always-on. The cost is one extra pass
per `add` batch to update the running quantile estimates, paid once
on the cold path; search latency is essentially unchanged.
- **Cross-arch top-K parity.** The AVX2 and AVX-512 BW kernels now
produce byte-identical top-K result sets to the NEON kernel for any
deterministic input. Per-vector f32 scores still differ by ~1e-5
relative across arches (different SIMD reduction orders), but those
rank swaps are confined to within-tie vectors and never change set
membership. Verified via the new `examples/kernel_xtest.rs` smoke
test (sha256 of sorted-per-query top-K indices matches across all
three SIMD paths).
#### Changed
- **On-disk format version bumped to 3** for both `.tv` and `.tvim`.
v3 appends a TQ+ trailer (per-coord shift + scale arrays) after the
existing scales section. The v3 reader is **backward-compatible**:
v2 files load with empty TQ+ vectors (identity calibration). Files
written by 0.7.x cannot be loaded by 0.6.x or older; there's no
forward-compat shim. Reindexing from source vectors picks up the
TQ+ recall lift; loading an old v2 file gives you the pre-TQ+
numbers.
- **x86 LUT-build is no longer data-dependent.** The AVX2 and AVX-512
BW kernels previously capped `max_lut` at `min(127, 65535 / n_byte_groups)`
to keep their no-flush u8→i16 accumulators in range — which at
d=1536/4-bit clamped to 42, and at d=3072/4-bit to 21, opening a
visible recall gap vs ARM (−1.6pp and −5.5pp respectively). Both
kernels now batch the inner loop by `FLUSH_EVERY=256` byte-groups
and run a mini-epilogue (SUB-trick + i16→f32 + fmadd into per-query
f32 accumulators) at the end of each batch — the same structure
NEON has used since 0.5.x. `max_lut` is now unconditionally 127 on
every arch. x86 speed is essentially flat vs the previous release
(the per-batch flush eliminates the same work from the single final
epilogue).
### turbovec — Python package (current: 0.5.3 → next: 0.6.0)
#### Added
- **TQ+ per-coordinate calibration.** Same kernel-level change as the
Rust crate; Python users see no API change. `TurboQuantIndex.add()`
carries a small extra pass per batch to update the running quantile
estimates (one-shot cold-path cost; search latency unchanged), and
`.search()` returns higher recall on the cells listed above. The
README's "How it works" section documents the calibration step.
#### Changed
- **On-disk format version bumped to 3** for both `.tv` and `.tvim`.
Same forward-compat policy as the Rust crate: old v2 files load
fine into 0.6.0+ (with identity calibration), but indexes written
by 0.6.0+ cannot be loaded by ≤ 0.5.3. Reindex from source vectors
to pick up the recall lift.
#### Fixed
- **x86/ARM recall parity at d=1536 and d=3072, 4-bit.** Previous
releases silently produced lower recall on x86 than ARM at high
dim — most visibly at d=3072/4-bit where x86 measured 0.919 @1 vs
ARM's 0.974 (−5.5pp). Same fix as the Rust crate (porting the
ARM-style periodic accumulator flush to AVX2 and AVX-512 BW). x86
search latency is essentially unchanged.
## turbovec 0.5.3 (Python package) + turbovec 0.6.0 (Rust crate) — 2026-05-25
### turbovec — Rust crate (current: 0.5.0 → next: 0.6.0)
#### Changed
- **BREAKING:** `TurboQuantIndex::new`, `TurboQuantIndex::new_lazy`,
`IdMapIndex::new`, and `IdMapIndex::new_lazy` now return
`Result<Self, ConstructError>` instead of panicking on invalid
input. The new `turbovec::ConstructError` enum covers `bit_width`
out of `{2, 3, 4}` and `dim` not a positive multiple of 8 (which
also closes a latent hole where `dim = 0` was silently accepted —
the previous `dim % 8 == 0` assertion vacuously passed for zero,
then divided-by-zero on the first `add`).
Migration: append `?` (or `.unwrap()` in tests/binaries) to
existing constructor calls. Mirrors the [`AddError`](src/error.rs)
pattern from the previous release.
- **Encode is 2–3× faster on aarch64.** SIMD-ifies the quantize +
scale + bit-pack inner loop via NEON (compare against boundaries
in 8 lanes at a time, weighted horizontal-add for the bit-pack)
and fuses the three passes so there's no intermediate
`codes: Vec<u8>` allocation. Rayon parallelises across rows on
both aarch64 and x86_64; x86_64 keeps the existing scalar inner
loop. Recall is bit-identical to the previous release at every
published cell (verified against `benchmarks/suite/recall_*.py`
on M3 Max). Measured throughput on M2 Pro, single-threaded:
- d=768, 4-bit: 22.5K → 66.3K vec/sec (2.9×)
- d=1536, 4-bit: 9.5K → 21.9K vec/sec (2.3×)
- d=1536, 2-bit: 16.6K → 25.7K vec/sec (1.5×)
- **Codebook is now cached across incremental `add` calls.** The
Lloyd-Max boundaries and centroids are a deterministic function
of `(bit_width, dim)`, so recomputing them on every `add` was
wasted work. They're now stored in `OnceLock` cells (the same
pattern already used for the rotation matrix) and reused across
calls. No behaviour change; faster incremental indexing.
### turbovec — Python package (current: 0.5.2 → next: 0.5.3)
#### Fixed
- **Linux wheels now actually import.** Every Linux wheel since
Linux build support was added had a missing `DT_NEEDED` entry for
`libopenblas`, so `import turbovec` failed at the dynamic linker
step with `undefined symbol: cblas_sgemm` — even on systems that
had OpenBLAS installed. The wheel now declares the dependency
explicitly, and `auditwheel` bundles a self-contained copy of
`libopenblas` (plus its `libgfortran` / `libquadmath` runtime
deps) into `turbovec.libs/`. Linux wheel size grows from ~1.8 MB
to ~11 MB (aarch64) / ~42 MB (x86_64) as a consequence — the
bundled OpenBLAS contains kernel variants for many micro-archs
and dispatches at runtime. The Linux release CI now also runs
`pytest` against the freshly-built wheel on native runners so
this class of bug can't ship silently again.
#### Changed
- **`TurboQuantIndex` and `IdMapIndex` constructors raise
`ValueError` on bad input** (`bit_width` outside `{2, 3, 4}`,
`dim` not a positive multiple of 8, including the previously
silently-accepted `dim = 0` case). Previously these surfaced as
`pyo3_runtime.PanicException`, which subclasses `BaseException`
and so wasn't caught by `except Exception:` — user code can now
recover from a configuration error as a normal usage error.
- **Encode (build-time, not query-time) is faster on aarch64.**
Same kernel-level change as the Rust crate; Python users see no
API change and bit-identical recall at every published cell.
Building an index with `TurboQuantIndex.add()` is ~2–3× faster on
M-series macOS and Linux aarch64. x86_64 sees the Rayon
parallelism but not the SIMD kernel.
## turbovec 0.5.2 (Python package) + turbovec 0.5.0 (Rust crate) — 2026-05-21
### turbovec — Rust crate (current: 0.4.1 → next: 0.5.0)
#### Changed
- **BREAKING:** `TurboQuantIndex::add_2d`, `IdMapIndex::add_with_ids_2d`,
and `IdMapIndex::add_with_ids` now return `Result<(), AddError>`
instead of panicking on invalid input. The new `turbovec::AddError`
enum covers dim mismatch, `dim % 8 != 0` on lazy-commit, vector
buffer length not a multiple of `dim`, ids/vectors count mismatch,
and duplicate ids. The low-level `TurboQuantIndex::add(&[f32])` and
constructor asserts are unchanged — they still panic, since those
signal contract violations rather than user-input errors.
Migration: append `?` (or `.unwrap()` in tests/binaries) to existing
calls. Match on `AddError` if you need to recover from specific
failure modes.
### turbovec — Python package (current: 0.5.1 → next: 0.5.2)
#### Changed
- **Dim mismatch on `add` / `add_with_ids` now raises `ValueError`**
instead of surfacing a `pyo3_runtime.PanicException` with a Rust
backtrace. The previous `PanicException` subclassed `BaseException`
and so was not caught by `except Exception:` — user code can now
recover from a wrong-shape batch as a normal usage error. The same
applies to duplicate ids and length mismatches on
`IdMapIndex.add_with_ids`.
## turbovec 0.5.1 (Python package) + turbovec 0.4.1 (Rust crate) — 2026-05-18
### turbovec — Rust crate (current: 0.4.0 → next: 0.4.1)
#### Added
- **Block-level early exit for selective mask searches** (closes
[#30](https://github.com/RyanCodrai/turbovec/issues/30)). When a
search is issued with `Some(mask)` the SIMD kernels now check
whether each 32-vector block contains any allowed slots before
doing the LUT lookup + popcount + score-decode work for that
block. If not, the entire block is short-circuited at one
integer-load + branch per block. The AVX-512BW path additionally
short-circuits 64-vector pairs at once where possible.
Measured speedup at 1% selectivity, 100K vectors, d=1536 (mask
allowing the last 1K slots): **6.4× on ARM (M3 Max), 12.7× on x86
(Sapphire Rapids c3-standard-8)**. Unmasked search latency is
unchanged (the guard only fires when a mask is passed).
Public API: no change to existing surfaces.
- **`turbovec::search::BLOCKS_SKIPPED_BY_MASK`** — atomic counter
incremented each time a block is short-circuited. Accessors
`blocks_skipped_by_mask()` and `reset_blocks_skipped_by_mask()`
are exposed for hybrid-retrieval telemetry. AVX-512BW pair-level
skips count as 2.
### turbovec — Python package (current: 0.5.0 → next: 0.5.1)
#### Added
- **Block-level early exit for selective `search_with_mask` calls.**
Same kernel-level change as the Rust crate; Python users see
identical API and unchanged unmasked latency. Selective masks now
run substantially faster (≈6–13× at 1% selectivity, scaling with
index size — larger indices amortize fixed per-query cost more
and see larger speedups). Closes
[#30](https://github.com/RyanCodrai/turbovec/issues/30).
## turbovec 0.5.0 (Python package) + turbovec 0.4.0 (Rust crate) — 2026-05-18
> **BREAKING** — on-disk file format version bumped from 1 to 2.
> Existing `.tv` and `.tvim` files written by turbovec ≤ 0.4.3 cannot
> be loaded by 0.5.0+. **Reindex from source vectors to migrate;**
> no in-place migration is provided.
### Migration
If you have indexes built with 0.4.3 or earlier, re-encode them:
```python
import numpy as np
from turbovec import TurboQuantIndex
# Source vectors (the f32 inputs your old index was built from).
vectors = np.load("my_vectors.npy") # shape (n, dim)
# Build a fresh 0.5.0 index. Same API, same recall guarantees, but with
# the new length-renormalization correction applied.
index = TurboQuantIndex(dim=vectors.shape[1], bit_width=4)
index.add(vectors)
index.write("my_index_v2.tv")
```
If you load an old file under 0.5.0+, you will see:
```
this .tv file was written by turbovec ≤ 0.4.3 (format version 1).
It is incompatible with turbovec 0.4.4+ because the per-vector scalar's
meaning changed. Rebuild this index from the source vectors using
turbovec 0.4.4 or later.
```
### turbovec — Rust crate (current: 0.3.0 → next: 0.4.0)
#### Added
- **Length-renormalized scoring.** The per-vector scalar stored in
`TurboQuantIndex` is now `||v|| / <u_rot, x̂>` instead of `||v||`,
giving an unbiased estimator of the inner product. The SIMD kernel
multiplies by this value at the same site it previously used the
norm — no change to kernel speed, storage layout, or public API.
#### Changed
- **On-disk format version bumped to 2** for both `.tv` and `.tvim`.
`.tv` now starts with a 4-byte magic `"TVPI"` + 1-byte version
prefix; `.tvim` keeps its existing magic with version bumped from 1
to 2. Loading a v1 file returns `io::Error` of kind `InvalidData`
with an upgrade-hint message; no in-place migration is provided.
- **`TurboQuantIndex::norms` field renamed to `scales`.** Internal
rename to match the value's new meaning. The SIMD kernel parameter
is `vec_scales` (to disambiguate from the per-query LUT calibration
`scales` parameter inside the same functions).
### turbovec — Python package (current: 0.4.3 → next: 0.5.0)
#### Added
- **Length-renormalized scoring.** Replaces the per-vector `||v||`
scalar with a RaBitQ-style correction `||v|| / <u_rot, x̂>` that
removes the systematic bias of the inner-product estimator. The
SIMD kernel is byte-for-byte unchanged — it multiplies by the new
scalar at the same site it previously used the norm. Recall@1
gains across published benchmarks:
- GloVe-200 2-bit: 0.5053 → 0.5524 (+4.7pp)
- GloVe-200 4-bit: 0.8115 → 0.8440 (+3.3pp)
- OpenAI-1536 2-bit: 0.8700 → 0.9060 (+3.6pp)
- OpenAI-1536 4-bit: 0.9550 → 0.9700 (+1.5pp)
- OpenAI-3072 2-bit: 0.9120 → 0.9240 (+1.2pp)
- OpenAI-3072 4-bit: 0.9670 → 0.9800 (+1.3pp)
Same-session ARM and x86 speed benchmarks confirm no measurable
search-latency change (deltas within FAISS noise floor on every
cell). The correction adds one extra dot product per vector at
encode time — a one-shot cost on the cold path, not visible to
search.
#### Changed
- **On-disk format version bumped to 2** for both `.tv` and `.tvim`.
`.tv` files now start with a 4-byte magic `"TVPI"` + 1-byte
version. `.tvim` files use the existing magic with version byte
bumped from 1 to 2.
- **Loading a turbovec ≤ 0.4.3 index raises with a clear error.**
The per-vector scalar's meaning changed (`||v||` → `||v|| / <u_rot, x̂>`),
so silently re-interpreting v1 files would produce wrong scores.
The new loader detects v1 files by their format signature and
raises `OSError` pointing the caller at rebuilding from source
vectors.
#### Fixed
- **`turbovec.haystack.TurboQuantDocumentStore` clamps cosine scores
to `[-1, 1]` before `scale_score` rescaling.** Cauchy–Schwarz
bounds the true cosine in that range, but the LUT scoring kernel's
float-precision noise can produce values slightly outside it —
most visibly on a self-query, which is algebraically 1.0 but the
kernel produces ~1.00016 after its per-sub-table calibration.
Without the clamp, downstream consumers of `scale_score=True` saw
scores `> 1.0` and the `[0, 1]` contract was violated. Dot-product
path uses a sigmoid that is already bounded; no clamp needed there.
## turbovec 0.4.3 (Python package) — 2026-05-18
### turbovec — Python package (current: 0.4.2 → next: 0.4.3)
#### Added
- **Windows x64 wheel** (closes [#31](https://github.com/RyanCodrai/turbovec/issues/31)).
Prior releases shipped only Linux x86_64/aarch64, macOS aarch64, and an
sdist — Windows users running `pip install turbovec` fell through to
the sdist and hit a `link.exe` build failure unless they had Rust + MSVC
installed locally. The release workflow now also builds a
`cp39-abi3-win_amd64` wheel and validates it by installing and running
the core pytest suite (`test_index.py`, `test_id_map.py`,
`test_filtering.py`) on the build runner before upload. Implementation
in [#33](https://github.com/RyanCodrai/turbovec/pull/33).
Intel Mac (macOS x86_64) was considered alongside Windows but blocked
by GitHub's December 2025 deprecation of free-tier `macos-13` runners;
tracked separately in [#34](https://github.com/RyanCodrai/turbovec/issues/34).
No library changes in this release — same Python API, same on-disk
format, same recall and throughput as 0.4.2. Pure platform-coverage
patch.
## turbovec 0.4.2 (Python package) — 2026-05-17
### turbovec — Python package (current: 0.4.1 → next: 0.4.2)
#### Fixed
- **`numpy` is now a declared runtime dependency.** The Python package
and every integration module imports `numpy` unconditionally, and the
Rust extension's Python surface expects NumPy arrays as input. Prior
releases relied on `numpy` being pulled in transitively via the
framework extras (`langchain-core`, `llama-index-core`, `haystack-ai`).
This broke `pip install turbovec[agno]` in clean environments because
`agno` doesn't depend on `numpy`. `numpy>=1.20` is now declared in
`[project].dependencies`, so it's installed regardless of which extra
(or none) is selected.
## turbovec 0.4.1 (Python package) — 2026-05-17
### turbovec — Python package (current: 0.4.0 → next: 0.4.1)
#### Added
- **Agno integration** (`turbovec.agno`). New `TurboQuantVectorDb` class
implementing Agno's `VectorDb` interface, structurally aligned with
`agno.vectordb.lancedb.LanceDb` (the closest in-tree single-machine
backend). Drop-in for callers that use `LanceDb` as their Agno
knowledge backend.
- Dim is sourced from `embedder.dimensions` (matches `LanceDb`); no
baked-in default.
- Filtered search uses the kernel-level `allowlist=` path: filters
resolve to a handle allowlist before scoring, so selective filters
return up to `limit` results from the filtered set instead of
fewer-than-`limit` from a post-filter.
- JSON side-car persistence (no pickle, no
`allow_dangerous_deserialization` flag).
- Constructor restricts `search_type=vector` and `distance=cosine`
— turbovec doesn't ship a BM25/lexical index and stores
unit-normalized vectors only. Non-vector / non-cosine constructions
raise `ValueError` rather than silently misbehaving.
- Honours `similarity_threshold` (cosine → relevance clamped to
`[0, 1]` via `(s + 1) / 2`), `reranker` (optional rerank pass after
vector retrieval), `content_id` / `content_hash` payload fields.
- Full async surface: `async_*` variants for create/insert/upsert/
search/drop/exists/name_exists, using the embedder's async batch
paths when available.
- Install: `pip install turbovec[agno]`.
## turbovec 0.3.0 (Rust crate) — 2026-05-17
## turbovec 0.3.0 (Rust crate) — 2026-05-17
### turbovec — Rust crate (current: 0.2.0 → next: 0.3.0)
#### Added
- **Search-time filtering.** New methods restrict the returned top-k to
a caller-supplied subset of vectors. The kernel applies the filter at
the heap-update site rather than via post-filtering, so selective
filters return up to `k` results from the allowed set instead of
fewer-than-`k` from an over-fetch pass. Output shape shrinks to
`min(k, n_allowed)` — consistent with the existing `k > len(idx)`
contract; no sentinel padding.
([#21](https://github.com/RyanCodrai/turbovec/issues/21))
- `TurboQuantIndex::search_with_mask(queries, k, mask: Option<&[bool]>)`
— slot bitmask, length equal to `len(idx)`.
- `IdMapIndex::search_with_allowlist(queries, k, allowlist: Option<&[u64]>)`
— external-id allowlist; translated to a slot bitmask internally
via the existing `id_to_slot` map. Panics on empty allowlist or
unknown ids.
- Threaded through every scoring path: NEON (aarch64), AVX2
(x86_64), AVX-512BW (x86_64), and the scalar fallback.
- **Lazy index construction.** The dim can now be deferred and inferred
from the first batch of vectors, rather than committed at construction
time. This is the same ergonomic improvement integration users were
already getting through the framework wrappers, pulled down into the
core so direct Rust users and any future integration get it for free.
- `TurboQuantIndex::new_lazy(bit_width)` and
`IdMapIndex::new_lazy(bit_width)` — construct an empty index with
no committed dim.
- `TurboQuantIndex::add_2d(vectors, dim)` and
`IdMapIndex::add_with_ids_2d(vectors, dim, ids)` — add a flat
vector batch with an explicit dim; locks the index dim on the
first call, validates on subsequent ones. Existing `add(&[f32])` /
`add_with_ids(&[f32], &[u64])` still work on a dim-known index and
panic with a clear message on a lazy uncommitted one.
- `TurboQuantIndex::dim_opt()` / `IdMapIndex::dim_opt()` return
`Option<usize>` — `None` for the lazy uncommitted state. The
existing `dim() -> usize` getters keep returning `usize`, with `0`
as a non-breaking sentinel for the lazy state (the eager
constructor asserts `dim >= 8`, so `0` doesn't collide).
- File format: `.tv` and `.tvim` headers encode the lazy state via
a `dim = 0` sentinel. Files written before this change always have
`dim >= 8` and load cleanly into the eager state.
#### Changed
- `search`, `search_with_mask`, and `prepare` on `TurboQuantIndex`
return empty results / are no-ops when called on a lazy
uncommitted index, rather than panicking.
## turbovec 0.4.0 (Python package) — 2026-05-17
### turbovec — Python package (current: 0.3.0 → next: 0.4.0)
#### Added
- **Search-time filtering.** Same feature surfaced as keyword-only
arguments on `search`:
- `TurboQuantIndex.search(queries, k, *, mask=None)` — `mask` is a
NumPy `bool` array of shape `(len(idx),)`.
- `IdMapIndex.search(queries, k, *, allowlist=None)` — `allowlist`
is a NumPy `uint64` array of external ids.
- Pre-validates shape, dtype, emptiness and unknown ids and raises
`ValueError` / `KeyError` rather than letting the Rust panic
surface as `pyo3.PanicException`.
([#21](https://github.com/RyanCodrai/turbovec/issues/21))
- **Lazy construction.** `TurboQuantIndex(dim=None, bit_width=4)` and
`IdMapIndex(dim=None, bit_width=4)` now accept an optional `dim`.
When omitted, the dim is inferred from the first `.add(...)` /
`.add_with_ids(...)` call using the input array's shape. The
framework integrations all rely on this internally now.
- `.dim` property on both index types now returns `int | None` (was
`int`); `None` means the index hasn't seen its first add yet.
#### Changed
- **Haystack integration** (`turbovec.haystack`):
`TurboQuantDocumentStore` is now a structural drop-in for
`haystack.document_stores.in_memory.InMemoryDocumentStore`. Audited
against `haystack-ai 2.28.0` and brought up to parity. In addition
to the earlier filter-resolution fix:
- `dim` is now optional in the constructor; the index is built
lazily on the first `write_documents`.
- Constructor accepts `embedding_similarity_function`
(`"cosine"` default, since turbovec stores unit-normalized
vectors), `async_executor`, and `return_embedding` for parity
with the reference. `scale_score=True` now uses the right
per-similarity-function formula (`(s + 1) / 2` for cosine,
`expit(s / 100)` for dot product), fixing a pre-existing bug.
- 12 `*_async` variants added (`count_documents_async`,
`filter_documents_async`, `write_documents_async`,
`delete_documents_async`, `delete_all_documents_async`,
`update_by_filter_async`, `count_documents_by_filter_async`,
`count_unique_metadata_by_filter_async`,
`get_metadata_fields_info_async`, `get_metadata_field_min_max_async`,
`get_metadata_field_unique_values_async`, `embedding_retrieval_async`).
- 8 utility methods added (`delete_all_documents`,
`delete_by_filter`, `update_by_filter`, `count_documents_by_filter`,
`count_unique_metadata_by_filter`, `get_metadata_fields_info`,
`get_metadata_field_min_max`, `get_metadata_field_unique_values`),
plus a `storage` property and `shutdown()`.
- `write_documents` now validates its input and raises
`ValueError("Please provide a list of Documents.")` on bad input
instead of an opaque `AttributeError`.
- Persistence methods renamed to match the reference:
`save → save_to_disk`, `load → load_from_disk`. (No deprecation
shims — pre-this-change persisted stores load fine, but the method
names change.)
- **LangChain integration** (`turbovec.langchain`):
`TurboQuantVectorStore` is now a structural drop-in for
`langchain_core.vectorstores.in_memory.InMemoryVectorStore`. Audited
against `langchain_core 0.3.63`. In addition to the earlier filter
fixes:
- `__init__` no longer requires a pre-built `IdMapIndex`. Lazy
construction lets `TurboQuantVectorStore(embedding)` work
directly — same no-arg ergonomics as the reference.
- `_select_relevance_score_fn` override added — maps the raw cosine
similarity into `[0, 1]` so `similarity_search_with_relevance_scores`
and `as_retriever(search_type="similarity_score_threshold")` work.
Result is clamped to `[0, 1]` to absorb the small overshoot caused
by quantization noise.
- `get_by_ids` / `aget_by_ids` implemented from the side-car
docstore.
- `add_documents` overrides the base-class default so partial
`Document.id` is honoured per-document (some ids explicit, others
UUID-generated) instead of being dropped wholesale.
- True async overrides: `aadd_documents`, `aadd_texts` and
`asimilarity_search_with_score` use `aembed_documents` /
`aembed_query` for genuine async embedding generation;
`asimilarity_search`, `asimilarity_search_by_vector`,
`amax_marginal_relevance_search`, `afrom_texts`, `adelete` are
explicit overrides too.
- `delete` now returns `None` (was `bool`) and is a no-op when
called with `ids=None` — matches the reference's contract.
- `max_marginal_relevance_search` / `_by_vector` /
`amax_marginal_relevance_search` raise `NotImplementedError` with
a clear message rather than the base class's bare
`NotImplementedError`. MMR isn't faithfully implementable on a
quantized index because the algorithm requires full-precision
candidate vectors that turbovec discards after encoding.
- Persistence methods renamed: `save_local → dump`, `load_local →
load`, matching the reference.
- **LlamaIndex integration** (`turbovec.llama_index`):
`TurboQuantVectorStore` is now a structural drop-in for
`llama_index.core.vector_stores.simple.SimpleVectorStore`. Audited
against `llama_index.core 0.12.39`. In addition to the earlier
filter fixes:
- `__init__` no longer requires a pre-built `IdMapIndex`;
`TurboQuantVectorStore()` works directly. `from_params(dim=None,
bit_width=4)` is also lazy.
- `get_nodes(node_ids, filters)` implemented (the reference raises
NotImplementedError because it doesn't store nodes; we do).
`clear()` resets state while preserving `bit_width`.
- `to_dict` / `from_dict` for config round-trip.
- `get(text_id)` raises `NotImplementedError` with an explanation —
we can't return the original embedding (quantized away).
- `delete_nodes(node_ids, filters)` now honours `filters` (previously
raised). Both constraints intersect when supplied.
- Async overrides for `async_add`, `adelete`, `adelete_nodes`,
`aclear`, `aquery`, `aget_nodes`.
- **StorageContext compatibility**: new
`from_persist_dir(persist_dir, namespace, fs)` matching the
reference's namespaced-filename convention, so
`StorageContext.from_defaults(persist_dir=...)` works. The
`persist` / `from_persist_path` on-disk layout is now stem-based:
`persist_path` is a path *stem* and we write `{stem}.tvim` +
`{stem}.nodes.json` next to each other. This fits StorageContext's
file-shaped paths and lets multiple namespaced stores share a
directory.
- **JSON side-cars across all three integrations.** Haystack, LangChain
and LlamaIndex persistence now writes a plain-JSON side-car next to
the binary `IdMapIndex` payload instead of a pickle. The
`allow_dangerous_deserialization` flag is gone everywhere — loading
is safe regardless of file provenance. Document / node metadata must
be JSON-serializable, which matches the constraint the reference
in-tree stores already impose. The side-car carries a
`schema_version` field; loaders reject unknown versions instead of
silently misinterpreting bytes.
[Unreleased]: https://github.com/RyanCodrai/turbovec/compare/v0.8.1...HEAD
[py-v0.4.2]: https://github.com/RyanCodrai/turbovec/compare/py-v0.4.1...py-v0.4.2
[py-v0.4.1]: https://github.com/RyanCodrai/turbovec/compare/py-v0.4.0...py-v0.4.1
CONTRIBUTING.md
+40
View File
@@ -0,0 +1,40 @@
# Contributing
Thanks for your interest in turbovec.
## Workflow
1. **Open an issue** describing what you've spotted — a bug, a missing feature, a documentation gap, a performance question. Include enough context that the conversation can start without back-and-forth on what you mean.
2. **Discuss.** If you want to suggest an implementation approach, do that in the issue. This is where the design conversation lives.
3. **Request contributor access** if you want to land code yourself. In your issue or in a follow-up, say so explicitly — e.g. "happy to take this; can I get contributor access?" I'll review the engagement so far and decide. If yes, I'll add you as a collaborator and you can open a PR for the issue.
If you don't request contributor access, that's fine — the issue itself is a valuable contribution, and I (or another contributor) may pick it up.
Only I merge to `main`.
## Why this is the workflow
The contributions that move turbovec forward are **good ideas, clearly articulated** — a sharp framing of a real problem, the right question to ask, an insight about how something should work, a benchmark observation that points at a real gap. That's the work I most need help with, and the work hardest to delegate. A well-written issue is more valuable than a PR.
The reason I've moved to invitation-only PRs is that the cost of reviewing a PR I have to mentally reconstruct from scratch — figure out what it's trying to do, why, whether it's correct, whether it fits the project's direction — is higher than just writing the change myself. This has become particularly acute with AI-assisted PRs that are technically clean but arrive without the design context or reasoning that makes review tractable. When the cognitive load of review exceeds the cognitive load of writing the change, the PR is a net loss for the project.
The "by invitation" gate isn't about credentials — it's about making sure the issue-side work has happened first, so when a PR arrives, review can be about *the code* rather than reconstructing *the why*. Contributors who've done that work via issue engagement are a joy to review. PRs that arrive cold without context aren't.
## For invited contributors: commit and PR conventions
- **One logical change per PR.** Refactors get their own PR, separate from feature work.
- **Commit messages:** short imperative title, body explaining *why* (the *what* is in the diff). Multi-line bodies should preserve formatting — use a HEREDOC if writing from the shell.
- **PRs reference their issue** with `Closes #N` and include a test plan.
- **`Co-Authored-By:` trailers** are fine on commits where Claude or another tool collaborated — leave them in place.
## Integration contributions
If you're adding or modifying an integration (LangChain, LlamaIndex, Haystack, Agno, or a new framework), structurally compare against the canonical in-tree reference store (`InMemoryVectorStore`, `SimpleVectorStore`, `InMemoryDocumentStore`, etc.) for that framework. The wrappers should match the reference's surface and idioms — that's the bar for a drop-in replacement.
## Build, test, bench
See the [Building](README.md#building) and [Running benchmarks](README.md#running-benchmarks) sections of the README. To run the integration test suites (LangChain, LlamaIndex, Haystack, Agno), install the corresponding extras — otherwise they're skipped:
```bash
pip install -e ".[langchain,llama-index,haystack,agno]"
```
Cargo.lock
Generated
+1257
View File
File diff suppressed because it is too large Load Diff
Cargo.toml
+12
View File
@@ -0,0 +1,12 @@
[workspace]
members = ["turbovec", "turbovec-python"]
# Standalone smoke test crate that lives outside the workspace on purpose —
# it must mimic a real downstream `cargo add turbovec` user (no inherited
# workspace settings, no shared lockfile).
exclude = ["examples/downstream-smoke"]
resolver = "2"
[profile.release]
lto = true
codegen-units = 1
opt-level = 3
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 Ryan Codrai
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
README.md
+238
View File
@@ -0,0 +1,238 @@
<p align="center">
<img src="docs/header.png" alt="turbovec — Google's TurboQuant for vector search" width="100%">
</p>
<p align="center">
<a href="https://github.com/RyanCodrai/turbovec/blob/main/LICENSE"><img src="https://img.shields.io/pypi/l/turbovec" alt="License"></a>
<a href="https://pypi.org/project/turbovec/"><img src="https://img.shields.io/pypi/v/turbovec?label=pypi&color=blue" alt="PyPI version"></a>
<a href="https://crates.io/crates/turbovec"><img src="https://img.shields.io/crates/v/turbovec?label=crates.io&color=blue" alt="crates.io version"></a>
<a href="https://arxiv.org/abs/2504.19874"><img src="https://img.shields.io/badge/paper-arXiv-b31b1b.svg" alt="TurboQuant paper"></a>
</p>
---
**A 10 million document corpus takes 31 GB of RAM as float32. turbovec fits it in 4 GB - and searches it faster than FAISS.**
turbovec is a Rust vector index with Python bindings, built on Google Research's [**TurboQuant**](https://arxiv.org/abs/2504.19874) algorithm — a data-oblivious quantizer with near-optimal distortion and no separate training phase.
- **Online ingest.** Add vectors, they're indexed — no train step, no parameter tuning, no rebuilds as the corpus grows.
- **Fast SIMD search.** Hand-written NEON (ARM) and AVX-512BW (x86) kernels beat FAISS IndexPQFastScan by 10–19% on ARM; on x86 they win the 4-bit configs and trail by a few percent on 2-bit.
- **Filter at search time.** Pass an id allowlist (or a slot bitmask) to `search()` and the kernel honours it directly. You always get up to `k` results from the allowed set — no over-fetching, no recall hit on selective filters.
- **Pure local.** No managed service, no data leaving your machine or VPC. Pair with any open-source embedding model for a fully air-gapped RAG stack.
Building RAG where privacy, memory, or latency matters? **You're in the right place.**
## Python
```bash
pip install turbovec
```
```python
from turbovec import TurboQuantIndex
index = TurboQuantIndex(dim=1536, bit_width=4)
index.add(vectors)
index.add(more_vectors)
scores, indices = index.search(query, k=10)
index.write("my_index.tv")
loaded = TurboQuantIndex.load("my_index.tv")
```
Need stable ids that survive deletes? Use `IdMapIndex`:
```python
import numpy as np
from turbovec import IdMapIndex
index = IdMapIndex(dim=1536, bit_width=4)
index.add_with_ids(vectors, np.array([1001, 1002, 1003], dtype=np.uint64))
scores, ids = index.search(query, k=10) # ids are your uint64 external ids
index.remove(1002) # O(1) by id
index.write("my_index.tvim")
loaded = IdMapIndex.load("my_index.tvim")
```
### Hybrid retrieval (filtered search)
Restrict results to a candidate set produced by another system (SQL, BM25, ACL, time window, …):
```python
import numpy as np
from turbovec import IdMapIndex
idx = IdMapIndex(dim=1536, bit_width=4)
idx.add_with_ids(vectors, ids)
# Stage 1: external system narrows to candidate ids.
allowed = np.array(db.execute("SELECT id FROM docs WHERE tenant=?", (t,)).fetchall(),
dtype=np.uint64)
# Stage 2: dense rerank within the candidate set.
scores, ids = idx.search(query, k=10, allowlist=allowed)
```
Filtering happens inside the SIMD kernel at 32-vector block granularity: blocks with no allowed slots are short-circuited before any LUT lookup or scoring work, and individual non-allowed slots inside scored blocks are dropped at heap-insert. Selective allowlists (small fraction of the index allowed) therefore avoid most of the SIMD cost rather than paying it and discarding the result afterwards.
The output length is `min(k, len(allowed))` — when the allowlist is smaller than `k` you get exactly `len(allowed)` results rather than padded fallbacks.
See [`docs/api.md`](docs/api.md) for the full reference.
### Framework integrations
Drop-in replacements for the in-tree reference vector / document stores in each framework. Same public surface, same persistence semantics, same retriever and pipeline wiring — swap the import and keep your pipeline.
- [LangChain](docs/integrations/langchain.md) — `pip install turbovec[langchain]` · replaces `langchain_core.vectorstores.InMemoryVectorStore`
- [LlamaIndex](docs/integrations/llama_index.md) — `pip install turbovec[llama-index]` · replaces `llama_index.core.vector_stores.SimpleVectorStore`
- [Haystack](docs/integrations/haystack.md) — `pip install turbovec[haystack]` · replaces `haystack.document_stores.in_memory.InMemoryDocumentStore`
- [Agno](docs/integrations/agno.md) — `pip install turbovec[agno]` · replaces `agno.vectordb.lancedb.LanceDb`
## Rust
```bash
cargo add turbovec
```
```rust
use turbovec::TurboQuantIndex;
let mut index = TurboQuantIndex::new(1536, 4).unwrap();
index.add(&vectors);
let results = index.search(&queries, 10);
index.write("index.tv").unwrap();
let loaded = TurboQuantIndex::load("index.tv").unwrap();
```
For stable external ids that survive deletes:
```rust
use turbovec::IdMapIndex;
let mut index = IdMapIndex::new(1536, 4).unwrap();
index.add_with_ids(&vectors, &[1001, 1002, 1003]).unwrap();
let (scores, ids) = index.search(&queries, 10);
index.remove(1002);
index.write("index.tvim").unwrap();
let loaded = IdMapIndex::load("index.tvim").unwrap();
```
## Recall
TurboQuant vs FAISS `IndexPQ` (LUT256, nbits=8) — the paper's Section 4.4 baseline. 100K vectors, k=64. FAISS PQ sub-quantizer counts sized to match TurboQuant's bit rate (m=d/4 at 2-bit, m=d/2 at 4-bit).
![Recall GloVe d=200](docs/recall_glove.svg)
![Recall d=1536](docs/recall_d1536.svg)
![Recall d=3072](docs/recall_d3072.svg)
Across OpenAI d=1536 and d=3072, TurboQuant beats FAISS by 0.2–1.9 points at R@1 across 2-bit and 4-bit, and both reach 1.0 by k=8 (≥0.997 already at k=4). GloVe d=200 is the harder regime — at low dim the asymptotic Beta assumption is looser. TurboQuant beats FAISS by 0.9 points at 4-bit and is effectively tied at 2-bit (within 0.1 points) at R@1, both tracking FAISS closely by k≈16.
**A note on baselines.** We compare against FAISS `IndexPQ` (LUT256, nbits=8, float32 LUT) because it's the default production-grade PQ most users would reach for. This is a stronger baseline than the custom u8-LUT PQ in the [TurboQuant paper](https://arxiv.org/abs/2504.19874) — FAISS uses a higher-precision LUT at scoring time and k-means++ for codebook training. We reproduce the paper's TurboQuant numbers on OpenAI d=1536 / d=3072 and hit similar numbers to other community reference implementations on low-dim embeddings (see [`turboquant-py`](https://pypi.org/project/turboquant-py/) at d=384). On GloVe (d=200) — the low-dim regime where the asymptotic Beta assumption is loosest — TurboQuant lands level with FAISS at 2-bit and ahead at 4-bit; TQ+ calibration closes the low-dim gap the base algorithm leaves.
Full results: [d=1536 2-bit](benchmarks/results/recall_d1536_2bit.json), [d=1536 4-bit](benchmarks/results/recall_d1536_4bit.json), [d=3072 2-bit](benchmarks/results/recall_d3072_2bit.json), [d=3072 4-bit](benchmarks/results/recall_d3072_4bit.json), [GloVe 2-bit](benchmarks/results/recall_glove_2bit.json), [GloVe 4-bit](benchmarks/results/recall_glove_4bit.json).
## Compression
![Compression](docs/compression.svg)
## Search Speed
All benchmarks: 100K vectors, 1K queries, k=64, median of 5 runs.
### ARM (Apple M3 Max)
![ARM Speed — Single-threaded](docs/arm_speed_st.svg)
![ARM Speed — Multi-threaded](docs/arm_speed_mt.svg)
On ARM, TurboQuant beats FAISS FastScan by 10–19% across every config.
### x86 (Intel Xeon Platinum 8481C / Sapphire Rapids, 8 vCPUs)
![x86 Speed — Single-threaded](docs/x86_speed_st.svg)
![x86 Speed — Multi-threaded](docs/x86_speed_mt.svg)
On x86, TurboQuant wins the 4-bit configs by up to ~5% (d=3072 multi-threaded ties) and is modestly behind FAISS on 2-bit — most visibly d=1536 single-threaded (~8%), within a few percent on the rest — where FAISS's AVX-512 VBMI path has the edge on the short 2-bit accumulate loop.
## How it works
Each vector is a direction on a high-dimensional hypersphere. TurboQuant compresses these directions using a simple insight: after applying a random rotation, every coordinate follows a known distribution -- regardless of the input data.
**1. Normalize.** Strip the length (norm) from each vector and store it as a single float. Now every vector is a unit direction on the hypersphere.
**2. Random rotation.** Multiply all vectors by the same random orthogonal matrix. After rotation, each coordinate independently follows a Beta distribution that converges to Gaussian N(0, 1/d) in high dimensions. This holds for any input data -- the rotation makes the coordinate distribution predictable.
**3. Per-coordinate calibration (TQ+).** The Beta distribution from step 2 is asymptotic — at finite dimensions, individual coordinates drift from the canonical shape (especially low-bit and word-vector-style embeddings). TQ+ fits two scalars per coordinate — a shift and a scale — during the first add, mapping each coordinate's empirical 5/95% quantiles onto the canonical Beta marginal. The Lloyd-Max codebook then quantizes against the *target* distribution it was designed for. The calibration is frozen after the first add and reused by subsequent adds — no retraining, no rebuilds, no separate train phase. Recall gain: up to +1.4pp at @1 on the cells that drift most (e.g. GloVe at 2-bit).
**4. Lloyd-Max scalar quantization.** Since the distribution is known, we can precompute the optimal way to bucket each coordinate. For 2-bit, that's 4 buckets; for 4-bit, 16 buckets. The [Lloyd-Max algorithm](https://en.wikipedia.org/wiki/Lloyd%27s_algorithm) finds bucket boundaries and centroids that minimize mean squared error. These are computed once from the math, not from the data.
**5. Bit-pack.** Each coordinate is now a small integer (0-3 for 2-bit, 0-15 for 4-bit). Pack these tightly into bytes. A 1536-dim vector goes from 6,144 bytes (FP32) to 384 bytes (2-bit). That's 16x compression.
**6. Length-renormalized scoring.** Scalar quantization systematically underestimates inner products — the reconstructed unit direction is a little shorter than the original. We compute one scalar per vector at encode time — the inner product of the rotated unit vector with its own centroid reconstruction — and store `||v|| / ⟨u, x̂⟩` alongside each compressed vector. The search kernel multiplies the per-candidate score by this scalar before heap insertion, turning the inner-product estimator from downward-biased into unbiased at zero search-time cost and zero extra storage. The recall gain shows up most at low bit widths, where the quantization shrinkage is largest.
Encoding cost: one extra `d`-dimensional dot product per vector to compute `⟨u, x̂⟩`. On 1M vectors at d=1536 this is sub-second of additional encode time — a one-shot price paid at ingest, not at query.
**Search.** Instead of decompressing every database vector, we rotate the query once into the same domain and score directly against the codebook values. The scoring kernel uses SIMD intrinsics (NEON on ARM, AVX-512BW on modern x86 with an AVX2 fallback) with nibble-split lookup tables for maximum throughput.
The Lloyd-Max codebook achieves distortion within a factor of 2.7x of the information-theoretic lower bound (Shannon's distortion-rate limit); the length-renormalization step removes the residual bias the Lloyd-Max codebook introduces on the inner-product estimator itself.
## Building
### Python (via maturin)
```bash
pip install maturin
cd turbovec-python
maturin build --release
pip install target/wheels/*.whl
```
### Rust
```bash
cargo build --release
```
All x86_64 builds target `x86-64-v3` (AVX2 baseline, Haswell 2013+) via `.cargo/config.toml`. Any CPU that can run the AVX2 fallback kernel can run the whole crate — the AVX-512 kernel is gated at runtime via `is_x86_feature_detected!` and only kicks in on hardware that supports it.
## Running benchmarks
Download datasets:
```bash
python3 benchmarks/download_data.py all # all datasets
python3 benchmarks/download_data.py glove # GloVe d=200
python3 benchmarks/download_data.py openai-1536 # OpenAI DBpedia d=1536
python3 benchmarks/download_data.py openai-3072 # OpenAI DBpedia d=3072
```
Each benchmark is a self-contained script in `benchmarks/suite/`. Run any one individually:
```bash
python3 benchmarks/suite/speed_d1536_2bit_arm_mt.py
python3 benchmarks/suite/recall_d1536_2bit.py
python3 benchmarks/suite/compression.py
```
Run all benchmarks for a category:
```bash
for f in benchmarks/suite/speed_*arm*.py; do python3 "$f"; done # all ARM speed
for f in benchmarks/suite/speed_*x86*.py; do python3 "$f"; done # all x86 speed
for f in benchmarks/suite/recall_*.py; do python3 "$f"; done # all recall
python3 benchmarks/suite/compression.py # compression
```
Results are saved as JSON to `benchmarks/results/`. Regenerate charts:
```bash
python3 benchmarks/create_diagrams.py
```
## References
- [TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate](https://arxiv.org/abs/2504.19874) (ICLR 2026) -- the paper this implements
- [RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search](https://arxiv.org/abs/2405.12497) (SIGMOD 2024) -- the source of the per-vector length-renormalization correction adapted in step 5
- [FAISS Fast accumulation of PQ and AQ codes](https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-(FastScan)) -- turbovec's x86 SIMD kernel adapts FastScan's pack layout, nibble-LUT scoring, and u16 accumulator strategy
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`RyanCodrai/turbovec`
- 原始仓库：https://github.com/RyanCodrai/turbovec
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
SECURITY.md
+73
View File
@@ -0,0 +1,73 @@
# Security Policy
Thank you for helping keep turbovec and its users safe. Security reports are
genuinely appreciated.
## Reporting a vulnerability
**Please do not open a public issue for a security vulnerability.** A public
report tips off attackers before a fix is available.
Instead, report it privately through GitHub:
1. Go to the [**Security** tab](https://github.com/RyanCodrai/turbovec/security)
of this repository.
2. Click **"Report a vulnerability"** to open a private advisory visible only
to the maintainers.
This routes the report through GitHub's private vulnerability reporting, where
we can discuss, develop, and review a fix — and, if appropriate, request a CVE
— without disclosing the issue until a patch is ready.
### What to include
A good report makes a fix faster. Where you can, please include:
- the affected component (core Rust crate, Python bindings, or a specific
framework integration) and the version,
- a description of the issue and its impact,
- a minimal reproduction — for a malformed-input bug, the crafted `.tv` /
`.tvim` bytes or the API call sequence that triggers it,
- the platform/architecture if relevant.
## What to expect
- We aim to acknowledge a report within a few days.
- We will confirm the issue, keep you updated on the fix, and coordinate a
disclosure timeline with you.
- With your permission, we will credit you in the published advisory.
Once a fix is released, we publish a GitHub Security Advisory so that the
vulnerability is recorded in the GitHub Advisory Database and downstream users
on crates.io and PyPI are alerted (e.g. via Dependabot) to upgrade.
## Supported versions
turbovec is pre-1.0 and ships frequently. Security fixes target the **latest
released version** of each surface:
- the `turbovec` crate on [crates.io](https://crates.io/crates/turbovec), and
- the `turbovec` distribution on [PyPI](https://pypi.org/project/turbovec/).
Please reproduce on the latest release before reporting where possible.
## Scope
In scope — for example:
- memory-unsafety or out-of-bounds access reachable from the public API,
- a panic, crash, or unbounded allocation triggered by loading an untrusted
`.tv` / `.tvim` index file,
- a panic or silently-wrong result reachable from normal Python/Rust API use
(e.g. malformed query input),
- data-integrity defects in the framework integration wrappers.
Generally out of scope:
- resource exhaustion driven only by a caller's own legitimately-large data on
a supported (64-bit) target,
- behavior on unsupported targets (turbovec requires a 64-bit platform and
refuses to compile elsewhere).
When in doubt, report it — we would rather triage an out-of-scope report than
miss a real one.
benchmarks/create_diagrams.py
+407
View File
@@ -0,0 +1,407 @@
"""Generate benchmark charts as SVG in the turboquant-wasm aesthetic.
Reads JSON files from ./results/ and writes:
../docs/arm_speed_st.svg, ../docs/arm_speed_mt.svg
../docs/x86_speed_st.svg, ../docs/x86_speed_mt.svg
../docs/recall_d1536.svg, ../docs/recall_d3072.svg, ../docs/recall_glove.svg
../docs/compression.svg
"""
import json
import math
import os
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "results")
DOCS_DIR = os.path.join(os.path.dirname(__file__), "..", "docs")
FONT = '-apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif'
C = {
"title": "#0f172a",
"subtitle": "#475569",
"label": "#0f172a",
"secondary": "#475569",
"tick": "#64748b",
"axis": "#334155",
"grid": "#e5e7eb",
"baseline": "#94a3b8",
"tq": "#635bff",
"tq_stroke": "#4338ca",
"tq_text": "#4338ca",
"faiss": "#9aa7b6",
"fp32": "#9aa7b6",
"four_bit": "#1d4ed8",
"two_bit": "#635bff",
"tq_2": "#635bff",
"tq_4": "#0f766e",
"faiss_2": "#9aa7b6",
"faiss_4": "#64748b",
}
def xe(s):
return (
str(s)
.replace("&", "&amp;")
.replace("<", "&lt;")
.replace(">", "&gt;")
.replace('"', "&quot;")
)
def nice_ceil(value):
if value <= 1:
return 1
exponent = math.floor(math.log10(value))
fraction = value / 10 ** exponent
if fraction <= 1:
nf = 1
elif fraction <= 1.5:
nf = 1.5
elif fraction <= 2:
nf = 2
elif fraction <= 5:
nf = 5
else:
nf = 10
return nf * 10 ** exponent
def style_block():
return (
f'<style>\n'
f' .title {{ font: 700 20px {FONT}; fill: {C["title"]}; }}\n'
f' .subtitle {{ font: 400 12px {FONT}; fill: {C["subtitle"]}; }}\n'
f' .panel {{ font: 700 14px {FONT}; fill: {C["title"]}; }}\n'
f' .label {{ font: 600 12px {FONT}; fill: {C["label"]}; }}\n'
f' .secondary {{ font: 400 11px {FONT}; fill: {C["secondary"]}; }}\n'
f' .tick {{ font: 400 11px {FONT}; fill: {C["tick"]}; }}\n'
f' .value {{ font: 700 11px {FONT}; fill: {C["label"]}; }}\n'
f' .value-accent {{ font: 700 11px {FONT}; fill: {C["tq_text"]}; }}\n'
f' .axis {{ font: 600 12px {FONT}; fill: {C["axis"]}; }}\n'
f' .legend {{ font: 600 12px {FONT}; fill: {C["label"]}; }}\n'
f'</style>'
)
def grid_lines(px, py, pw, ph, y_lo, y_hi, fmt, step_count=5):
parts = []
for i in range(step_count + 1):
v = y_lo + (y_hi - y_lo) * i / step_count
y = py + ph - (v - y_lo) / (y_hi - y_lo) * ph
parts.append(
f'<line x1="{px}" y1="{y:.1f}" x2="{px + pw}" y2="{y:.1f}" stroke="{C["grid"]}" stroke-width="1" />'
)
parts.append(
f'<text x="{px - 10}" y="{y + 4:.1f}" text-anchor="end" class="tick">{xe(fmt(v))}</text>'
)
parts.append(
f'<line x1="{px}" y1="{py + ph:.1f}" x2="{px + pw}" y2="{py + ph:.1f}" stroke="{C["baseline"]}" stroke-width="1.5" />'
)
return "\n".join(parts)
def paired_panel(px, py, pw, ph, panel_title, groups, tick_fmt, value_fmt, y_max):
parts = [grid_lines(px, py, pw, ph, 0, y_max, tick_fmt)]
parts.append(f'<text x="{px}" y="{py - 14}" class="panel">{xe(panel_title)}</text>')
n = len(groups)
band = pw / n
bar_w = min(44, band * 0.32)
gap = 6
for i, g in enumerate(groups):
cx = px + band * i + band / 2
tq_x = cx - bar_w - gap / 2
faiss_x = cx + gap / 2
tq_h = (g["tq"] / y_max) * ph
faiss_h = (g["faiss"] / y_max) * ph
tq_y = py + ph - tq_h
faiss_y = py + ph - faiss_h
label_y = py + ph + 22
parts.append(
f'<rect x="{tq_x:.1f}" y="{tq_y:.1f}" width="{bar_w}" height="{tq_h:.1f}" rx="6" '
f'fill="{C["tq"]}" stroke="{C["tq_stroke"]}" stroke-width="1.5" />'
)
parts.append(
f'<rect x="{faiss_x:.1f}" y="{faiss_y:.1f}" width="{bar_w}" height="{faiss_h:.1f}" rx="6" fill="{C["faiss"]}" />'
)
parts.append(
f'<text x="{tq_x + bar_w/2:.1f}" y="{tq_y - 6:.1f}" text-anchor="middle" class="value-accent">{xe(value_fmt(g["tq"]))}</text>'
)
parts.append(
f'<text x="{faiss_x + bar_w/2:.1f}" y="{faiss_y - 6:.1f}" text-anchor="middle" class="value">{xe(value_fmt(g["faiss"]))}</text>'
)
primary, _, secondary = g["label"].partition("|")
parts.append(f'<text x="{cx:.1f}" y="{label_y}" text-anchor="middle" class="label">{xe(primary)}</text>')
if secondary:
parts.append(f'<text x="{cx:.1f}" y="{label_y + 15}" text-anchor="middle" class="secondary">{xe(secondary)}</text>')
return "\n".join(parts)
def legend_tq_faiss(x, y):
parts = [
f'<rect x="{x}" y="{y - 10}" width="14" height="14" rx="3" fill="{C["tq"]}" stroke="{C["tq_stroke"]}" stroke-width="1.5" />',
f'<text x="{x + 22}" y="{y + 1}" class="legend" style="fill: {C["tq_text"]};">TurboQuant</text>',
f'<rect x="{x + 140}" y="{y - 10}" width="14" height="14" rx="3" fill="{C["faiss"]}" />',
f'<text x="{x + 162}" y="{y + 1}" class="legend">FAISS</text>',
]
return "\n".join(parts)
def load_json(name):
with open(os.path.join(RESULTS_DIR, name)) as f:
return json.load(f)
def speed_panels(arch):
panels = {"st": [], "mt": []}
for dim in (1536, 3072):
for bw in (2, 4):
for th in ("st", "mt"):
entry = load_json(f"speed_d{dim}_{bw}bit_{arch}_{th}.json")
panels[th].append(
{
"label": f"d={dim}|{bw}-bit",
"tq": entry["tq_ms_per_query"],
"faiss": entry["faiss_ms_per_query"],
}
)
return panels
def write_speed_panel(arch, hw_label, thread_key, thread_label, tick_fmt, value_fmt, filename):
panels = speed_panels(arch)
width, height = 900, 460
margin = {"top": 82, "right": 32, "bottom": 108, "left": 84}
pw = width - margin["left"] - margin["right"]
ph = height - margin["top"] - margin["bottom"]
px = margin["left"]
py = margin["top"]
y_max = nice_ceil(max(max(g["tq"], g["faiss"]) for g in panels[thread_key]) * 1.22)
parts = [
paired_panel(
px, py, pw, ph, thread_label, panels[thread_key],
tick_fmt=tick_fmt,
value_fmt=value_fmt,
y_max=y_max,
),
f'<text x="26" y="{py + ph/2}" transform="rotate(-90, 26, {py + ph/2})" class="axis">ms / query</text>',
legend_tq_faiss(margin["left"], height - 26),
]
body = "\n".join(parts)
svg = f"""<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}" role="img" aria-label="Search Latency — {xe(hw_label)} — {xe(thread_label)}">
{style_block()}
<rect width="100%" height="100%" fill="#ffffff" />
<text x="{margin["left"]}" y="32" class="title">Search Latency — {xe(hw_label)} — {xe(thread_label)}</text>
<text x="{margin["left"]}" y="52" class="subtitle">100K vectors, 1K queries, k=64, median of 5 runs</text>
{body}
</svg>
"""
out = os.path.join(DOCS_DIR, filename)
with open(out, "w") as f:
f.write(svg)
print(f"wrote {out}")
def line_panel(px, py, pw, ph, panel_title, series, x_values, x_labels, y_lo, y_hi):
parts = [
grid_lines(px, py, pw, ph, y_lo, y_hi, lambda v: f"{v:.2f}"),
f'<text x="{px}" y="{py - 14}" class="panel">{xe(panel_title)}</text>',
]
x_min = math.log2(x_values[0])
x_max = math.log2(x_values[-1])
def xpx(v):
return px + (math.log2(v) - x_min) / (x_max - x_min) * pw
def ypx(v):
return py + ph - (v - y_lo) / (y_hi - y_lo) * ph
for v, lbl in zip(x_values, x_labels):
parts.append(
f'<text x="{xpx(v):.1f}" y="{py + ph + 20}" text-anchor="middle" class="label">{xe(lbl)}</text>'
)
for s in series:
color = s["color"]
dash = ' stroke-dasharray="6 4"' if s.get("dashed") else ""
points = [(xpx(x), ypx(y)) for x, y in zip(x_values, s["values"])]
path = "M " + " L ".join(f"{x:.1f},{y:.1f}" for x, y in points)
parts.append(
f'<path d="{path}" fill="none" stroke="{color}" stroke-width="2.25"{dash} />'
)
for x, y in points:
parts.append(f'<circle cx="{x:.1f}" cy="{y:.1f}" r="3.5" fill="{color}" />')
return "\n".join(parts)
def write_recall_panel(dim_key, dim_label, filename, y_lo=0.85):
width, height = 900, 460
margin = {"top": 82, "right": 32, "bottom": 108, "left": 84}
pw = width - margin["left"] - margin["right"]
ph = height - margin["top"] - margin["bottom"]
px = margin["left"]
py = margin["top"]
x_values = [1, 2, 4, 8, 16, 32, 64]
x_labels = ["1", "2", "4", "8", "16", "32", "64"]
# Draw FAISS lines first (background), then TurboQuant on top — emphasises
# the TQ series when lines overlap or cross at high-K.
faiss_series = []
tq_series = []
for bw_key, bw_label in [("2bit", "2-bit"), ("4bit", "4-bit")]:
data = load_json(f"recall_{dim_key}_{bw_key}.json")
tq_vals = [float(data["tq_recalls"][str(k)]) for k in x_values]
faiss_vals = [float(data["faiss_recalls"][str(k)]) for k in x_values]
tq_color = C["tq_2"] if bw_key == "2bit" else C["tq_4"]
faiss_color = C["faiss_2"] if bw_key == "2bit" else C["faiss_4"]
tq_series.append({"label": f"TQ {bw_label}", "values": tq_vals, "color": tq_color})
faiss_series.append({"label": f"FAISS {bw_label}", "values": faiss_vals, "color": faiss_color, "dashed": True})
series = faiss_series + tq_series
parts = [
line_panel(px, py, pw, ph, dim_label, series, x_values, x_labels, y_lo, 1.005),
f'<text x="{px - 62}" y="{py + ph/2}" transform="rotate(-90, {px - 62}, {py + ph/2})" class="axis">recall@1@k</text>',
f'<text x="{px + pw/2}" y="{py + ph + 48}" text-anchor="middle" class="axis">k</text>',
]
legend_y = height - 26
lx = margin["left"]
items = [
("TQ 2-bit", C["tq_2"], False),
("TQ 4-bit", C["tq_4"], False),
("FAISS 2-bit", C["faiss_2"], True),
("FAISS 4-bit", C["faiss_4"], True),
]
for i, (lbl, col, dash) in enumerate(items):
cx = lx + i * 140
dash_attr = ' stroke-dasharray="6 4"' if dash else ""
parts.append(
f'<line x1="{cx}" y1="{legend_y - 2}" x2="{cx + 24}" y2="{legend_y - 2}" stroke="{col}" stroke-width="2.25"{dash_attr} />'
)
parts.append(f'<circle cx="{cx + 12}" cy="{legend_y - 2}" r="3.5" fill="{col}" />')
parts.append(f'<text x="{cx + 32}" y="{legend_y + 1}" class="legend">{xe(lbl)}</text>')
body = "\n".join(parts)
svg = f"""<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}" role="img" aria-label="Recall — {xe(dim_label)}">
{style_block()}
<rect width="100%" height="100%" fill="#ffffff" />
<text x="{margin["left"]}" y="32" class="title">Recall — {xe(dim_label)}</text>
<text x="{margin["left"]}" y="52" class="subtitle">100K vectors, k=64 search. recall@1@k measures how often the true top-1 result appears in the top-k returned.</text>
{body}
</svg>
"""
out = os.path.join(DOCS_DIR, filename)
with open(out, "w") as f:
f.write(svg)
print(f"wrote {out}")
def write_compression_chart(filename):
datasets = [
("GloVe|d=200", 76.3, 9.9, 5.1),
("OpenAI|d=1536", 585.9, 73.6, 37.0),
("OpenAI|d=3072", 1171.9, 146.9, 73.6),
]
width, height = 900, 460
margin = {"top": 82, "right": 32, "bottom": 108, "left": 84}
pw = width - margin["left"] - margin["right"]
ph = height - margin["top"] - margin["bottom"]
px = margin["left"]
py = margin["top"]
y_max = nice_ceil(max(d[1] for d in datasets) * 1.15)
parts = [grid_lines(px, py, pw, ph, 0, y_max, lambda v: f"{v:.0f}")]
n = len(datasets)
band = pw / n
bar_w = min(56, band * 0.22)
gap = 10
for i, (label, fp32, four, two) in enumerate(datasets):
cx = px + band * i + band / 2
x_fp = cx - 1.5 * bar_w - gap
x_4 = cx - 0.5 * bar_w
x_2 = cx + 0.5 * bar_w + gap
def draw(xbar, val, color, accent=False):
h = (val / y_max) * ph
y = py + ph - h
stroke = (
f' stroke="{C["tq_stroke"]}" stroke-width="1.5"' if accent else ""
)
value_cls = "value-accent" if accent else "value"
return "\n".join(
[
f'<rect x="{xbar:.1f}" y="{y:.1f}" width="{bar_w}" height="{h:.1f}" rx="6" fill="{color}"{stroke} />',
f'<text x="{xbar + bar_w/2:.1f}" y="{y - 6:.1f}" text-anchor="middle" class="{value_cls}">{xe(f"{val:.0f}")}</text>',
]
)
parts.append(draw(x_fp, fp32, C["fp32"]))
parts.append(draw(x_4, four, C["four_bit"]))
parts.append(draw(x_2, two, C["two_bit"], accent=True))
label_y = py + ph + 22
primary, _, secondary = label.partition("|")
parts.append(f'<text x="{cx:.1f}" y="{label_y}" text-anchor="middle" class="label">{xe(primary)}</text>')
if secondary:
parts.append(f'<text x="{cx:.1f}" y="{label_y + 15}" text-anchor="middle" class="secondary">{xe(secondary)}</text>')
parts.append(
f'<text x="26" y="{py + ph/2}" transform="rotate(-90, 26, {py + ph/2})" class="axis">Index size (MB)</text>'
)
legend_y = height - 26
lx = margin["left"]
items = [
("FP32", C["fp32"], False),
("4-bit", C["four_bit"], False),
("2-bit", C["two_bit"], True),
]
for i, (lbl, col, accent) in enumerate(items):
lcx = lx + i * 120
stroke = f' stroke="{C["tq_stroke"]}" stroke-width="1.5"' if accent else ""
parts.append(f'<rect x="{lcx}" y="{legend_y - 10}" width="14" height="14" rx="3" fill="{col}"{stroke} />')
parts.append(f'<text x="{lcx + 22}" y="{legend_y + 1}" class="legend">{xe(lbl)}</text>')
body = "\n".join(parts)
svg = f"""<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}" role="img" aria-label="Index Size — TurboQuant">
{style_block()}
<rect width="100%" height="100%" fill="#ffffff" />
<text x="{margin["left"]}" y="32" class="title">Index Size — 100K vectors</text>
<text x="{margin["left"]}" y="52" class="subtitle">TurboQuant packs vectors ~16× smaller than FP32 at 2-bit with comparable recall</text>
{body}
</svg>
"""
out = os.path.join(DOCS_DIR, filename)
with open(out, "w") as f:
f.write(svg)
print(f"wrote {out}")
if __name__ == "__main__":
os.makedirs(DOCS_DIR, exist_ok=True)
write_speed_panel("arm", "ARM (Apple M3 Max)", "st", "Single-threaded",
tick_fmt=lambda v: f"{v:.1f}", value_fmt=lambda v: f"{v:.2f}",
filename="arm_speed_st.svg")
write_speed_panel("arm", "ARM (Apple M3 Max)", "mt", "Multi-threaded",
tick_fmt=lambda v: f"{v:.2f}", value_fmt=lambda v: f"{v:.3f}",
filename="arm_speed_mt.svg")
write_speed_panel("x86", "x86 (Intel Sapphire Rapids, 8 vCPUs)", "st", "Single-threaded",
tick_fmt=lambda v: f"{v:.1f}", value_fmt=lambda v: f"{v:.2f}",
filename="x86_speed_st.svg")
write_speed_panel("x86", "x86 (Intel Sapphire Rapids, 8 vCPUs)", "mt", "Multi-threaded",
tick_fmt=lambda v: f"{v:.2f}", value_fmt=lambda v: f"{v:.3f}",
filename="x86_speed_mt.svg")
write_recall_panel("d1536", "d=1536", "recall_d1536.svg")
write_recall_panel("d3072", "d=3072", "recall_d3072.svg")
write_recall_panel("glove", "GloVe d=200", "recall_glove.svg", y_lo=0.4)
write_compression_chart("compression.svg")
benchmarks/download_data.py
+64
View File
@@ -0,0 +1,64 @@
"""
Download benchmark datasets.
Usage:
python3 download_data.py glove openai-1536 openai-3072
python3 download_data.py all
"""
import os
import subprocess
import sys
import numpy as np
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
GLOVE_PATH = os.path.join(DATA_DIR, "glove-200-angular.hdf5")
GLOVE_URL = "http://ann-benchmarks.com/glove-200-angular.hdf5"
def download_glove():
if os.path.exists(GLOVE_PATH):
print(f"Already downloaded: {GLOVE_PATH}")
return
os.makedirs(DATA_DIR, exist_ok=True)
print(f"Downloading {GLOVE_URL}...")
subprocess.run(["curl", "-L", "-o", GLOVE_PATH, GLOVE_URL], check=True)
print(f"Saved: {GLOVE_PATH} ({os.path.getsize(GLOVE_PATH) / 1024 / 1024:.0f} MB)")
def download_openai(dim):
from datasets import load_dataset
path = os.path.join(DATA_DIR, f"openai-{dim}.npy")
if os.path.exists(path):
print(f"Already downloaded: {path}")
return
os.makedirs(DATA_DIR, exist_ok=True)
name = f"Qdrant/dbpedia-entities-openai3-text-embedding-3-large-{dim}-1M"
col = f"text-embedding-3-large-{dim}-embedding"
print(f"Downloading {name}...")
ds = load_dataset(name, split="train")
ds.set_format("numpy")
vecs = ds[col].astype(np.float32)
np.save(path, vecs)
print(f"Saved: {path} ({os.path.getsize(path) / 1024 / 1024:.0f} MB)")
TARGETS = {
"glove": download_glove,
"openai-1536": lambda: download_openai(1536),
"openai-3072": lambda: download_openai(3072),
}
if __name__ == "__main__":
args = sys.argv[1:] if len(sys.argv) > 1 else ["all"]
if "all" in args:
args = list(TARGETS.keys())
for name in args:
if name not in TARGETS:
print(f"Unknown dataset: {name}")
print(f"Available: {', '.join(TARGETS.keys())}")
continue
TARGETS[name]()
benchmarks/results/compression.json
+50
View File
@@ -0,0 +1,50 @@
{
"glove_d200_2bit": {
"n": 100000,
"dim": 200,
"bit_width": 2,
"fp32_mb": 76.3,
"index_mb": 5.1,
"ratio": 14.8
},
"glove_d200_4bit": {
"n": 100000,
"dim": 200,
"bit_width": 4,
"fp32_mb": 76.3,
"index_mb": 9.9,
"ratio": 7.7
},
"openai_d1536_2bit": {
"n": 100000,
"dim": 1536,
"bit_width": 2,
"fp32_mb": 585.9,
"index_mb": 37.0,
"ratio": 15.8
},
"openai_d1536_4bit": {
"n": 100000,
"dim": 1536,
"bit_width": 4,
"fp32_mb": 585.9,
"index_mb": 73.6,
"ratio": 8.0
},
"openai_d3072_2bit": {
"n": 100000,
"dim": 3072,
"bit_width": 2,
"fp32_mb": 1171.9,
"index_mb": 73.6,
"ratio": 15.9
},
"openai_d3072_4bit": {
"n": 100000,
"dim": 3072,
"bit_width": 4,
"fp32_mb": 1171.9,
"index_mb": 146.9,
"ratio": 8.0
}
}
benchmarks/results/recall_d1536_2bit.json
+25
View File
@@ -0,0 +1,25 @@
{
"dataset": "openai-1536",
"dim": 1536,
"bit_width": 2,
"faiss_variant": "IndexPQ(m=384, nbits=8)",
"seed": 42,
"tq_recalls": {
"1": 0.891,
"2": 0.979,
"4": 0.998,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
},
"faiss_recalls": {
"1": 0.872,
"2": 0.977,
"4": 0.997,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
}
}
benchmarks/results/recall_d1536_4bit.json
+25
View File
@@ -0,0 +1,25 @@
{
"dataset": "openai-1536",
"dim": 1536,
"bit_width": 4,
"faiss_variant": "IndexPQ(m=768, nbits=8)",
"seed": 42,
"tq_recalls": {
"1": 0.974,
"2": 0.998,
"4": 1.0,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
},
"faiss_recalls": {
"1": 0.966,
"2": 0.998,
"4": 1.0,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
}
}
benchmarks/results/recall_d3072_2bit.json
+25
View File
@@ -0,0 +1,25 @@
{
"dataset": "openai-3072",
"dim": 3072,
"bit_width": 2,
"faiss_variant": "IndexPQ(m=768, nbits=8)",
"seed": 42,
"tq_recalls": {
"1": 0.929,
"2": 0.988,
"4": 0.999,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
},
"faiss_recalls": {
"1": 0.912,
"2": 0.986,
"4": 1.0,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
}
}
benchmarks/results/recall_d3072_4bit.json
+25
View File
@@ -0,0 +1,25 @@
{
"dataset": "openai-3072",
"dim": 3072,
"bit_width": 4,
"faiss_variant": "IndexPQ(m=1536, nbits=8)",
"seed": 42,
"tq_recalls": {
"1": 0.974,
"2": 0.999,
"4": 1.0,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
},
"faiss_recalls": {
"1": 0.972,
"2": 0.998,
"4": 1.0,
"8": 1.0,
"16": 1.0,
"32": 1.0,
"64": 1.0
}
}
benchmarks/results/recall_glove_2bit.json
+25
View File
@@ -0,0 +1,25 @@
{
"dataset": "glove",
"dim": 200,
"bit_width": 2,
"faiss_variant": "IndexPQ(m=50, nbits=8)",
"seed": 42,
"tq_recalls": {
"1": 0.5637,
"2": 0.7209,
"4": 0.8443,
"8": 0.9236,
"16": 0.9676,
"32": 0.9892,
"64": 0.9966
},
"faiss_recalls": {
"1": 0.5643,
"2": 0.7188,
"4": 0.8446,
"8": 0.9252,
"16": 0.97,
"32": 0.9908,
"64": 0.9981
}
}
benchmarks/results/recall_glove_4bit.json
+25
View File
@@ -0,0 +1,25 @@
{
"dataset": "glove",
"dim": 200,
"bit_width": 4,
"faiss_variant": "IndexPQ(m=100, nbits=8)",
"seed": 42,
"tq_recalls": {
"1": 0.8498,
"2": 0.9584,
"4": 0.9936,
"8": 0.9998,
"16": 1.0,
"32": 1.0,
"64": 1.0
},
"faiss_recalls": {
"1": 0.841,
"2": 0.9515,
"4": 0.9914,
"8": 0.9986,
"16": 1.0,
"32": 1.0,
"64": 1.0
}
}
benchmarks/results/speed_d1536_2bit_arm_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 2,
"arch": "arm",
"threading": "mt",
"tq_ms_per_query": 0.103,
"faiss_ms_per_query": 0.115
}
benchmarks/results/speed_d1536_2bit_arm_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 2,
"arch": "arm",
"threading": "st",
"tq_ms_per_query": 1.083,
"faiss_ms_per_query": 1.235
}
benchmarks/results/speed_d1536_2bit_x86_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 2,
"arch": "x86",
"threading": "mt",
"tq_ms_per_query": 0.304,
"faiss_ms_per_query": 0.295
}
benchmarks/results/speed_d1536_2bit_x86_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 2,
"arch": "x86",
"threading": "st",
"tq_ms_per_query": 1.271,
"faiss_ms_per_query": 1.172
}
benchmarks/results/speed_d1536_4bit_arm_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 4,
"arch": "arm",
"threading": "mt",
"tq_ms_per_query": 0.185,
"faiss_ms_per_query": 0.22
}
benchmarks/results/speed_d1536_4bit_arm_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 4,
"arch": "arm",
"threading": "st",
"tq_ms_per_query": 1.992,
"faiss_ms_per_query": 2.45
}
benchmarks/results/speed_d1536_4bit_x86_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 4,
"arch": "x86",
"threading": "mt",
"tq_ms_per_query": 0.576,
"faiss_ms_per_query": 0.59
}
benchmarks/results/speed_d1536_4bit_x86_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 1536,
"bit_width": 4,
"arch": "x86",
"threading": "st",
"tq_ms_per_query": 2.439,
"faiss_ms_per_query": 2.56
}
benchmarks/results/speed_d3072_2bit_arm_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 2,
"arch": "arm",
"threading": "mt",
"tq_ms_per_query": 0.201,
"faiss_ms_per_query": 0.224
}
benchmarks/results/speed_d3072_2bit_arm_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 2,
"arch": "arm",
"threading": "st",
"tq_ms_per_query": 2.124,
"faiss_ms_per_query": 2.439
}
benchmarks/results/speed_d3072_2bit_x86_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 2,
"arch": "x86",
"threading": "mt",
"tq_ms_per_query": 0.626,
"faiss_ms_per_query": 0.59
}
benchmarks/results/speed_d3072_2bit_x86_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 2,
"arch": "x86",
"threading": "st",
"tq_ms_per_query": 2.657,
"faiss_ms_per_query": 2.582
}
benchmarks/results/speed_d3072_4bit_arm_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 4,
"arch": "arm",
"threading": "mt",
"tq_ms_per_query": 0.375,
"faiss_ms_per_query": 0.448
}
benchmarks/results/speed_d3072_4bit_arm_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 4,
"arch": "arm",
"threading": "st",
"tq_ms_per_query": 3.968,
"faiss_ms_per_query": 4.925
}
benchmarks/results/speed_d3072_4bit_x86_mt.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 4,
"arch": "x86",
"threading": "mt",
"tq_ms_per_query": 1.177,
"faiss_ms_per_query": 1.177
}
benchmarks/results/speed_d3072_4bit_x86_st.json
+8
View File
@@ -0,0 +1,8 @@
{
"dim": 3072,
"bit_width": 4,
"arch": "x86",
"threading": "st",
"tq_ms_per_query": 5.342,
"faiss_ms_per_query": 5.474
}
benchmarks/suite/compression.py
+90
View File
@@ -0,0 +1,90 @@
#!/usr/bin/env python3
"""Compression benchmark: TQ index file size vs FP32."""
import os, json, tempfile
import numpy as np
import h5py
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
GLOVE_PATH = os.path.join(DATA_DIR, "glove-200-angular.hdf5")
RESULTS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "results")
RESULT_FILE = os.path.join(RESULTS_DIR, "compression.json")
def load_openai(dim, seed=42):
path = os.path.join(DATA_DIR, f"openai-{dim}.npy")
all_vecs = np.load(path)
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def load_glove(seed=42):
f = h5py.File(GLOVE_PATH, "r")
all_train = f["train"][:].astype(np.float32)
queries = f["test"][:].astype(np.float32)
rng = np.random.RandomState(seed)
idx = rng.choice(len(all_train), 100_000, replace=False)
database = all_train[idx]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def measure_index_size(database, dim, bit_width):
from turbovec import TurboQuantIndex
index = TurboQuantIndex(dim, bit_width)
index.add(database)
with tempfile.NamedTemporaryFile(suffix=".tv", delete=False) as tmp:
path = tmp.name
try:
index.write(path)
return os.path.getsize(path)
finally:
os.remove(path)
def main():
datasets = [
("glove_d200", 200, load_glove),
("openai_d1536", 1536, lambda: load_openai(1536)),
("openai_d3072", 3072, lambda: load_openai(3072)),
]
results = {}
for name, dim, loader in datasets:
print(f"\nLoading {name}...")
database, _ = loader()
n = len(database)
fp32_bytes = n * dim * 4
for bit_width in [2, 4]:
key = f"{name}_{bit_width}bit"
print(f" {key}...", end=" ", flush=True)
index_bytes = measure_index_size(database, dim, bit_width)
fp32_mb = fp32_bytes / (1024 * 1024)
index_mb = index_bytes / (1024 * 1024)
results[key] = {
"n": n,
"dim": dim,
"bit_width": bit_width,
"fp32_mb": round(fp32_mb, 1),
"index_mb": round(index_mb, 1),
"ratio": round(fp32_mb / index_mb, 1),
}
print(f"{fp32_mb:.1f} MB -> {index_mb:.1f} MB ({fp32_mb / index_mb:.1f}x)")
os.makedirs(RESULTS_DIR, exist_ok=True)
with open(RESULT_FILE, "w") as f:
json.dump(results, f, indent=2)
print(f"\nResults saved to {RESULT_FILE}")
print(json.dumps(results, indent=2))
if __name__ == "__main__":
main()
benchmarks/suite/recall_d1536_2bit.py
+81
View File
@@ -0,0 +1,81 @@
#!/usr/bin/env python3
"""Recall benchmark: OpenAI d=1536, 2-bit (TQ vs FAISS PQ with LUT256).
Matches the paper's Section 4.4 setup: FAISS `IndexPQ` with 256 codewords
per sub-quantizer (nbits=8), grouping 4 coordinates per sub at 2-bit
(m = dim / 4). Not FastScan (LUT16) which the paper rejects.
"""
import os, json, time, numpy as np, faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "..", "results")
DIM = 1536
BIT_WIDTH = 2
K = 64
K_VALUES = [1, 2, 4, 8, 16, 32, 64]
SEED = 42
def load_openai(dim):
path = os.path.join(DATA_DIR, f"openai-{dim}.npy")
all_vecs = np.load(path)
rng = np.random.RandomState(SEED)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]].astype(np.float32)
queries = all_vecs[idx[100_000:101_000]].astype(np.float32)
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def recall_at_1_at_k(true_top1, predicted_indices, k):
return float(np.mean([true_top1[i] in predicted_indices[i, :k] for i in range(len(true_top1))]))
def main():
print(f"=== OpenAI d={DIM} {BIT_WIDTH}-bit (seed={SEED}) ===")
m = DIM // 4
nbits = 8
database, queries = load_openai(DIM)
true_top1 = np.argmax(queries @ database.T, axis=1)
t0 = time.time()
index_tq = TurboQuantIndex(DIM, bit_width=BIT_WIDTH)
index_tq.add(database)
_, tq_indices = index_tq.search(queries, k=K)
tq_indices = np.array(tq_indices)
tq_recalls = {str(k): round(recall_at_1_at_k(true_top1, tq_indices, k), 4) for k in K_VALUES}
print(f" TQ ({time.time() - t0:.1f}s) recall@1 = {tq_recalls['1']:.4f}")
t0 = time.time()
index_faiss = faiss.IndexPQ(DIM, m, nbits, faiss.METRIC_INNER_PRODUCT)
index_faiss.train(database)
index_faiss.add(database)
_, faiss_ids = index_faiss.search(queries, K)
faiss_recalls = {str(k): round(recall_at_1_at_k(true_top1, faiss_ids, k), 4) for k in K_VALUES}
print(f" FAISS ({time.time() - t0:.1f}s) recall@1 = {faiss_recalls['1']:.4f}")
results = {
"dataset": f"openai-{DIM}",
"dim": DIM,
"bit_width": BIT_WIDTH,
"faiss_variant": f"IndexPQ(m={m}, nbits={nbits})",
"seed": SEED,
"tq_recalls": tq_recalls,
"faiss_recalls": faiss_recalls,
}
print("\nTQ: ", tq_recalls)
print("FAISS:", faiss_recalls)
os.makedirs(RESULTS_DIR, exist_ok=True)
out_path = os.path.join(RESULTS_DIR, "recall_d1536_2bit.json")
with open(out_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nSaved to {out_path}")
if __name__ == "__main__":
main()
benchmarks/suite/recall_d1536_4bit.py
+81
View File
@@ -0,0 +1,81 @@
#!/usr/bin/env python3
"""Recall benchmark: OpenAI d=1536, 4-bit (TQ vs FAISS PQ with LUT256).
Matches the paper's Section 4.4 setup: FAISS `IndexPQ` with 256 codewords
per sub-quantizer (nbits=8), grouping 2 coordinates per sub at 4-bit
(m = dim / 2).
"""
import os, json, time, numpy as np, faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "..", "results")
DIM = 1536
BIT_WIDTH = 4
K = 64
K_VALUES = [1, 2, 4, 8, 16, 32, 64]
SEED = 42
def load_openai(dim):
path = os.path.join(DATA_DIR, f"openai-{dim}.npy")
all_vecs = np.load(path)
rng = np.random.RandomState(SEED)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]].astype(np.float32)
queries = all_vecs[idx[100_000:101_000]].astype(np.float32)
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def recall_at_1_at_k(true_top1, predicted_indices, k):
return float(np.mean([true_top1[i] in predicted_indices[i, :k] for i in range(len(true_top1))]))
def main():
print(f"=== OpenAI d={DIM} {BIT_WIDTH}-bit (seed={SEED}) ===")
m = DIM // 2
nbits = 8
database, queries = load_openai(DIM)
true_top1 = np.argmax(queries @ database.T, axis=1)
t0 = time.time()
index_tq = TurboQuantIndex(DIM, bit_width=BIT_WIDTH)
index_tq.add(database)
_, tq_indices = index_tq.search(queries, k=K)
tq_indices = np.array(tq_indices)
tq_recalls = {str(k): round(recall_at_1_at_k(true_top1, tq_indices, k), 4) for k in K_VALUES}
print(f" TQ ({time.time() - t0:.1f}s) recall@1 = {tq_recalls['1']:.4f}")
t0 = time.time()
index_faiss = faiss.IndexPQ(DIM, m, nbits, faiss.METRIC_INNER_PRODUCT)
index_faiss.train(database)
index_faiss.add(database)
_, faiss_ids = index_faiss.search(queries, K)
faiss_recalls = {str(k): round(recall_at_1_at_k(true_top1, faiss_ids, k), 4) for k in K_VALUES}
print(f" FAISS ({time.time() - t0:.1f}s) recall@1 = {faiss_recalls['1']:.4f}")
results = {
"dataset": f"openai-{DIM}",
"dim": DIM,
"bit_width": BIT_WIDTH,
"faiss_variant": f"IndexPQ(m={m}, nbits={nbits})",
"seed": SEED,
"tq_recalls": tq_recalls,
"faiss_recalls": faiss_recalls,
}
print("\nTQ: ", tq_recalls)
print("FAISS:", faiss_recalls)
os.makedirs(RESULTS_DIR, exist_ok=True)
out_path = os.path.join(RESULTS_DIR, "recall_d1536_4bit.json")
with open(out_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nSaved to {out_path}")
if __name__ == "__main__":
main()
benchmarks/suite/recall_d3072_2bit.py
+76
View File
@@ -0,0 +1,76 @@
#!/usr/bin/env python3
"""Recall benchmark: OpenAI d=3072, 2-bit (TQ vs FAISS PQ with LUT256)."""
import os, json, time, numpy as np, faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "..", "results")
DIM = 3072
BIT_WIDTH = 2
K = 64
K_VALUES = [1, 2, 4, 8, 16, 32, 64]
SEED = 42
def load_openai(dim):
path = os.path.join(DATA_DIR, f"openai-{dim}.npy")
all_vecs = np.load(path)
rng = np.random.RandomState(SEED)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]].astype(np.float32)
queries = all_vecs[idx[100_000:101_000]].astype(np.float32)
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def recall_at_1_at_k(true_top1, predicted_indices, k):
return float(np.mean([true_top1[i] in predicted_indices[i, :k] for i in range(len(true_top1))]))
def main():
print(f"=== OpenAI d={DIM} {BIT_WIDTH}-bit (seed={SEED}) ===")
m = DIM // 4
nbits = 8
database, queries = load_openai(DIM)
true_top1 = np.argmax(queries @ database.T, axis=1)
t0 = time.time()
index_tq = TurboQuantIndex(DIM, bit_width=BIT_WIDTH)
index_tq.add(database)
_, tq_indices = index_tq.search(queries, k=K)
tq_indices = np.array(tq_indices)
tq_recalls = {str(k): round(recall_at_1_at_k(true_top1, tq_indices, k), 4) for k in K_VALUES}
print(f" TQ ({time.time() - t0:.1f}s) recall@1 = {tq_recalls['1']:.4f}")
t0 = time.time()
index_faiss = faiss.IndexPQ(DIM, m, nbits, faiss.METRIC_INNER_PRODUCT)
index_faiss.train(database)
index_faiss.add(database)
_, faiss_ids = index_faiss.search(queries, K)
faiss_recalls = {str(k): round(recall_at_1_at_k(true_top1, faiss_ids, k), 4) for k in K_VALUES}
print(f" FAISS ({time.time() - t0:.1f}s) recall@1 = {faiss_recalls['1']:.4f}")
results = {
"dataset": f"openai-{DIM}",
"dim": DIM,
"bit_width": BIT_WIDTH,
"faiss_variant": f"IndexPQ(m={m}, nbits={nbits})",
"seed": SEED,
"tq_recalls": tq_recalls,
"faiss_recalls": faiss_recalls,
}
print("\nTQ: ", tq_recalls)
print("FAISS:", faiss_recalls)
os.makedirs(RESULTS_DIR, exist_ok=True)
out_path = os.path.join(RESULTS_DIR, "recall_d3072_2bit.json")
with open(out_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nSaved to {out_path}")
if __name__ == "__main__":
main()
benchmarks/suite/recall_d3072_4bit.py
+76
View File
@@ -0,0 +1,76 @@
#!/usr/bin/env python3
"""Recall benchmark: OpenAI d=3072, 4-bit (TQ vs FAISS PQ with LUT256)."""
import os, json, time, numpy as np, faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "..", "results")
DIM = 3072
BIT_WIDTH = 4
K = 64
K_VALUES = [1, 2, 4, 8, 16, 32, 64]
SEED = 42
def load_openai(dim):
path = os.path.join(DATA_DIR, f"openai-{dim}.npy")
all_vecs = np.load(path)
rng = np.random.RandomState(SEED)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]].astype(np.float32)
queries = all_vecs[idx[100_000:101_000]].astype(np.float32)
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def recall_at_1_at_k(true_top1, predicted_indices, k):
return float(np.mean([true_top1[i] in predicted_indices[i, :k] for i in range(len(true_top1))]))
def main():
print(f"=== OpenAI d={DIM} {BIT_WIDTH}-bit (seed={SEED}) ===")
m = DIM // 2
nbits = 8
database, queries = load_openai(DIM)
true_top1 = np.argmax(queries @ database.T, axis=1)
t0 = time.time()
index_tq = TurboQuantIndex(DIM, bit_width=BIT_WIDTH)
index_tq.add(database)
_, tq_indices = index_tq.search(queries, k=K)
tq_indices = np.array(tq_indices)
tq_recalls = {str(k): round(recall_at_1_at_k(true_top1, tq_indices, k), 4) for k in K_VALUES}
print(f" TQ ({time.time() - t0:.1f}s) recall@1 = {tq_recalls['1']:.4f}")
t0 = time.time()
index_faiss = faiss.IndexPQ(DIM, m, nbits, faiss.METRIC_INNER_PRODUCT)
index_faiss.train(database)
index_faiss.add(database)
_, faiss_ids = index_faiss.search(queries, K)
faiss_recalls = {str(k): round(recall_at_1_at_k(true_top1, faiss_ids, k), 4) for k in K_VALUES}
print(f" FAISS ({time.time() - t0:.1f}s) recall@1 = {faiss_recalls['1']:.4f}")
results = {
"dataset": f"openai-{DIM}",
"dim": DIM,
"bit_width": BIT_WIDTH,
"faiss_variant": f"IndexPQ(m={m}, nbits={nbits})",
"seed": SEED,
"tq_recalls": tq_recalls,
"faiss_recalls": faiss_recalls,
}
print("\nTQ: ", tq_recalls)
print("FAISS:", faiss_recalls)
os.makedirs(RESULTS_DIR, exist_ok=True)
out_path = os.path.join(RESULTS_DIR, "recall_d3072_4bit.json")
with open(out_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nSaved to {out_path}")
if __name__ == "__main__":
main()
benchmarks/suite/recall_glove_2bit.py
+82
View File
@@ -0,0 +1,82 @@
#!/usr/bin/env python3
"""Recall benchmark: GloVe d=200, 2-bit (TQ vs FAISS PQ with LUT256).
Uses FAISS `IndexPQ` (not FastScan) to stay compatible with GloVe's d=200,
which isn't m%32-aligned. Matches the paper's Section 4.4 configuration:
4 coordinates per sub-quantizer at 2-bit (m = d / 4 = 50), 256 codewords.
"""
import os, json, time, numpy as np, h5py, faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
GLOVE_PATH = os.path.join(DATA_DIR, "glove-200-angular.hdf5")
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "..", "results")
DIM = 200
BIT_WIDTH = 2
K = 64
K_VALUES = [1, 2, 4, 8, 16, 32, 64]
SEED = 42
def load_glove():
f = h5py.File(GLOVE_PATH, "r")
all_train = f["train"][:].astype(np.float32)
queries = f["test"][:].astype(np.float32)
rng = np.random.RandomState(SEED)
idx = rng.choice(len(all_train), 100_000, replace=False)
database = all_train[idx]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def recall_at_1_at_k(true_top1, predicted_indices, k):
return float(np.mean([true_top1[i] in predicted_indices[i, :k] for i in range(len(true_top1))]))
def main():
print(f"=== GloVe d={DIM} {BIT_WIDTH}-bit (seed={SEED}) ===")
m = DIM // 4
nbits = 8
database, queries = load_glove()
true_top1 = np.argmax(queries @ database.T, axis=1)
t0 = time.time()
index_tq = TurboQuantIndex(DIM, bit_width=BIT_WIDTH)
index_tq.add(database)
_, tq_indices = index_tq.search(queries, k=K)
tq_indices = np.array(tq_indices)
tq_recalls = {str(k): round(recall_at_1_at_k(true_top1, tq_indices, k), 4) for k in K_VALUES}
print(f" TQ ({time.time() - t0:.1f}s) recall@1 = {tq_recalls['1']:.4f}")
t0 = time.time()
index_faiss = faiss.IndexPQ(DIM, m, nbits, faiss.METRIC_INNER_PRODUCT)
index_faiss.train(database)
index_faiss.add(database)
_, faiss_ids = index_faiss.search(queries, K)
faiss_recalls = {str(k): round(recall_at_1_at_k(true_top1, faiss_ids, k), 4) for k in K_VALUES}
print(f" FAISS ({time.time() - t0:.1f}s) recall@1 = {faiss_recalls['1']:.4f}")
results = {
"dataset": "glove",
"dim": DIM,
"bit_width": BIT_WIDTH,
"faiss_variant": f"IndexPQ(m={m}, nbits={nbits})",
"seed": SEED,
"tq_recalls": tq_recalls,
"faiss_recalls": faiss_recalls,
}
print("\nTQ: ", tq_recalls)
print("FAISS:", faiss_recalls)
os.makedirs(RESULTS_DIR, exist_ok=True)
out_path = os.path.join(RESULTS_DIR, "recall_glove_2bit.json")
with open(out_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nSaved to {out_path}")
if __name__ == "__main__":
main()
benchmarks/suite/recall_glove_4bit.py
+82
View File
@@ -0,0 +1,82 @@
#!/usr/bin/env python3
"""Recall benchmark: GloVe d=200, 4-bit (TQ vs FAISS PQ with LUT256).
Uses FAISS `IndexPQ` (not FastScan) to stay compatible with GloVe's d=200,
which isn't m%32-aligned. Matches the paper's Section 4.4 configuration:
2 coordinates per sub-quantizer at 4-bit (m = d / 2 = 100), 256 codewords.
"""
import os, json, time, numpy as np, h5py, faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
GLOVE_PATH = os.path.join(DATA_DIR, "glove-200-angular.hdf5")
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "..", "results")
DIM = 200
BIT_WIDTH = 4
K = 64
K_VALUES = [1, 2, 4, 8, 16, 32, 64]
SEED = 42
def load_glove():
f = h5py.File(GLOVE_PATH, "r")
all_train = f["train"][:].astype(np.float32)
queries = f["test"][:].astype(np.float32)
rng = np.random.RandomState(SEED)
idx = rng.choice(len(all_train), 100_000, replace=False)
database = all_train[idx]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
def recall_at_1_at_k(true_top1, predicted_indices, k):
return float(np.mean([true_top1[i] in predicted_indices[i, :k] for i in range(len(true_top1))]))
def main():
print(f"=== GloVe d={DIM} {BIT_WIDTH}-bit (seed={SEED}) ===")
m = DIM // 2
nbits = 8
database, queries = load_glove()
true_top1 = np.argmax(queries @ database.T, axis=1)
t0 = time.time()
index_tq = TurboQuantIndex(DIM, bit_width=BIT_WIDTH)
index_tq.add(database)
_, tq_indices = index_tq.search(queries, k=K)
tq_indices = np.array(tq_indices)
tq_recalls = {str(k): round(recall_at_1_at_k(true_top1, tq_indices, k), 4) for k in K_VALUES}
print(f" TQ ({time.time() - t0:.1f}s) recall@1 = {tq_recalls['1']:.4f}")
t0 = time.time()
index_faiss = faiss.IndexPQ(DIM, m, nbits, faiss.METRIC_INNER_PRODUCT)
index_faiss.train(database)
index_faiss.add(database)
_, faiss_ids = index_faiss.search(queries, K)
faiss_recalls = {str(k): round(recall_at_1_at_k(true_top1, faiss_ids, k), 4) for k in K_VALUES}
print(f" FAISS ({time.time() - t0:.1f}s) recall@1 = {faiss_recalls['1']:.4f}")
results = {
"dataset": "glove",
"dim": DIM,
"bit_width": BIT_WIDTH,
"faiss_variant": f"IndexPQ(m={m}, nbits={nbits})",
"seed": SEED,
"tq_recalls": tq_recalls,
"faiss_recalls": faiss_recalls,
}
print("\nTQ: ", tq_recalls)
print("FAISS:", faiss_recalls)
os.makedirs(RESULTS_DIR, exist_ok=True)
out_path = os.path.join(RESULTS_DIR, "recall_glove_4bit.json")
with open(out_path, "w") as f:
json.dump(results, f, indent=2)
print(f"\nSaved to {out_path}")
if __name__ == "__main__":
main()
benchmarks/suite/speed_d1536_2bit_arm_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_2bit_arm_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d1536_2bit_arm_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_2bit_arm_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d1536_2bit_x86_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_2bit_x86_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d1536_2bit_x86_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_2bit_x86_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d1536_4bit_arm_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_4bit_arm_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d1536_4bit_arm_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_4bit_arm_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d1536_4bit_x86_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_4bit_x86_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d1536_4bit_x86_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 1536, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d1536_4bit_x86_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_2bit_arm_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_2bit_arm_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_2bit_arm_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_2bit_arm_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_2bit_x86_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_2bit_x86_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_2bit_x86_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 2
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM // 2
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_2bit_x86_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_4bit_arm_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_4bit_arm_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_4bit_arm_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "arm", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_4bit_arm_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_4bit_x86_mt.py
+49
View File
@@ -0,0 +1,49 @@
import os, time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "mt",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_4bit_x86_mt.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
benchmarks/suite/speed_d3072_4bit_x86_st.py
+52
View File
@@ -0,0 +1,52 @@
import os
os.environ["RAYON_NUM_THREADS"] = "1"
import time, json, numpy as np
import faiss
from turbovec import TurboQuantIndex
DATA_DIR = os.path.expanduser("~/data/py-turboquant")
DIM, BIT_WIDTH = 3072, 4
def load_openai(dim, seed=42):
all_vecs = np.load(os.path.join(DATA_DIR, f"openai-{dim}.npy"))
rng = np.random.RandomState(seed)
idx = rng.permutation(len(all_vecs))
database = all_vecs[idx[:100_000]]
queries = all_vecs[idx[100_000:101_000]]
database /= np.linalg.norm(database, axis=-1, keepdims=True)
queries /= np.linalg.norm(queries, axis=-1, keepdims=True)
return database, queries
database, queries = load_openai(DIM)
faiss.omp_set_num_threads(1)
# TurboQuant
tq = TurboQuantIndex(dim=DIM, bit_width=BIT_WIDTH)
tq.add(database)
tq.search(queries[:1], k=64) # warmup
tq_times = []
for _ in range(5):
t0 = time.perf_counter()
tq.search(queries, k=64)
tq_times.append((time.perf_counter() - t0) / len(queries) * 1000)
tq_ms = sorted(tq_times)[2]
# FAISS PQ
m_pq = DIM
pq = faiss.IndexPQFastScan(DIM, m_pq, 4)
pq.train(database)
pq.add(database)
pq.search(queries[:1], 64) # warmup
faiss_times = []
for _ in range(5):
t0 = time.perf_counter()
pq.search(queries, 64)
faiss_times.append((time.perf_counter() - t0) / len(queries) * 1000)
faiss_ms = sorted(faiss_times)[2]
result = {"dim": DIM, "bit_width": BIT_WIDTH, "arch": "x86", "threading": "st",
"tq_ms_per_query": round(tq_ms, 3), "faiss_ms_per_query": round(faiss_ms, 3)}
out = os.path.join(os.path.dirname(__file__), "..", "results", "speed_d3072_4bit_x86_st.json")
os.makedirs(os.path.dirname(out), exist_ok=True)
json.dump(result, open(out, "w"), indent=2)
print(json.dumps(result, indent=2))
docs/api.md
+182
View File
@@ -0,0 +1,182 @@
# API Reference
turbovec exposes two index types and one serialization format per type.
- [`TurboQuantIndex`](#turboquantindex) — positional index, O(1) `swap_remove` delete.
- [`IdMapIndex`](#idmapindex) — stable external `u64` ids on top of `TurboQuantIndex`.
- [File formats](#file-formats) — `.tv` and `.tvim`.
All examples below are Python. The Rust API mirrors it — see each type's rustdoc for the exact signatures.
---
## `TurboQuantIndex`
Positional index. Each vector is identified by its insertion slot (`0..n`). Fast and small, but external references to slots are invalidated by `swap_remove`. If you need stable ids, use [`IdMapIndex`](#idmapindex).
```python
from turbovec import TurboQuantIndex
idx = TurboQuantIndex(dim=1536, bit_width=4)
idx.add(vectors) # np.ndarray of shape (n, dim), float32
scores, indices = idx.search(queries, k=10)
idx.swap_remove(5) # O(1); the previously-last vector moves into slot 5
idx.write("index.tv") # .tv format
loaded = TurboQuantIndex.load("index.tv")
```
`dim` is optional. Omit it to let the index pick up the dimensionality from the first batch of vectors:
```python
idx = TurboQuantIndex(bit_width=4) # dim inferred on first add
idx.add(vectors) # locks dim to vectors.shape[1]
```
Before the first add, `idx.dim` is `None`, `len(idx)` is `0`, and `search()` returns empty results.
### Methods
| Method | Notes |
|---|---|
| `TurboQuantIndex(dim=None, bit_width=4)` | `bit_width ∈ {2, 3, 4}`. `dim` must be a positive multiple of 8 and `≤ 65536` (`MAX_DIM`). `dim` is optional; when omitted it is inferred from the first `add` call. |
| `add(vectors)` | `vectors` is a contiguous float32 array of shape `(n, dim)`. On a lazy index the first call locks `dim`; subsequent calls must match. Raises `ValueError` on dim mismatch, a zero-width (0-column) batch, or any coordinate that is non-finite (NaN/Inf) or `\|value\| ≥ 1e16`. |
| `search(queries, k, *, mask=None)` | Returns `(scores, indices)`, both shape `(nq, effective_k)`. Indices are `int64` slot positions. `mask` is an optional `bool` array of length `len(idx)`; when given, only slots with `mask[i] == True` contribute. `effective_k = min(k, mask.sum())`. Raises `ValueError` on a non-finite or `\|value\| ≥ 1e16` query coordinate. |
| `swap_remove(idx)` | O(1). Moves the last vector into `idx`; returns the previous position of that moved vector (so external refs can be updated if needed). |
| `prepare()` | Optional. Eagerly builds the rotation matrix, Lloyd-Max centroids and SIMD-blocked layout so the first `search` call doesn't pay the one-time cost. No-op on a lazy index that hasn't seen its first add. |
| `write(path)` / `load(path)` | `.tv` format. |
| `len(idx)` / `idx.dim` / `idx.bit_width` | Introspection. `idx.dim` returns `int` once committed, or `None` on a lazy index that hasn't seen its first add. |
### `swap_remove` semantics
`swap_remove(i)` is named to match Rust's [`Vec::swap_remove`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.swap_remove): the last element moves into slot `i`, and the vector is truncated by one. It is **not** a shift (FAISS's `IndexPQ::remove_ids` behaviour). Order is not preserved; slot indices of vectors you didn't delete may now point at different vectors than before.
Use [`IdMapIndex`](#idmapindex) if external references have to stay stable across deletes.
---
## `IdMapIndex`
Stable-id wrapper around `TurboQuantIndex`. Roughly equivalent to FAISS's `IndexIDMap2` — hash-table backed, O(1) `remove(id)`.
```python
import numpy as np
from turbovec import IdMapIndex
idx = IdMapIndex(dim=1536, bit_width=4)
idx.add_with_ids(vectors, np.array([1001, 1002, 1003], dtype=np.uint64))
scores, ids = idx.search(queries, k=10) # ids are uint64 external ids
idx.remove(1002) # O(1) by id
assert 1003 in idx # __contains__ sugar
idx.write("index.tvim") # .tvim format
loaded = IdMapIndex.load("index.tvim")
```
As with [`TurboQuantIndex`](#turboquantindex), `dim` is optional and gets inferred from the first `add_with_ids` call:
```python
idx = IdMapIndex(bit_width=4) # dim inferred on first add
idx.add_with_ids(vectors, ids) # locks dim to vectors.shape[1]
```
### Methods
| Method | Notes |
|---|---|
| `IdMapIndex(dim=None, bit_width=4)` | `bit_width ∈ {2, 3, 4}`; `dim` must be a positive multiple of 8 and `≤ 65536`. `dim` is optional; when omitted it is inferred from the first `add_with_ids` call. |
| `add_with_ids(vectors, ids)` | `ids` is a `uint64` array with length `vectors.shape[0]`. On a lazy index the first call locks `dim`. Raises `ValueError` on dim mismatch, duplicate ids, `len(ids) != vectors.shape[0]`, a zero-width batch, or a non-finite / `\|value\| ≥ 1e16` coordinate. |
| `remove(id) -> bool` | `True` if the id was present and removed, `False` otherwise. O(1). |
| `search(queries, k, *, allowlist=None)` | Returns `(scores, ids)` — `ids` are `uint64` external ids. `allowlist` is an optional `uint64` array of ids; when given, results are restricted to those ids and `effective_k = min(k, len(allowlist))`. Raises `ValueError` on an empty allowlist or a non-finite / `\|value\| ≥ 1e16` query coordinate, and `KeyError` on unknown ids. |
| `contains(id)` / `id in idx` | Membership. |
| `write(path)` / `load(path)` | `.tvim` format. |
| `len(idx)` / `idx.dim` / `idx.bit_width` / `prepare()` | Same as `TurboQuantIndex`. |
### When to use which
- `TurboQuantIndex` — you never delete, or you're fine with positional ids.
- `IdMapIndex` — you need stable external ids (e.g. string-id → vector mapping maintained by the caller).
All the framework integrations (LangChain, LlamaIndex, Haystack) use `IdMapIndex` internally for exactly this reason.
---
## Filtering
Both index types support restricting the returned top-`k` to a caller-supplied subset of vectors. Unlike post-filtering (search then drop), the kernel never inserts disallowed vectors into the per-query heap, so you always get up to `k` results from the allowed set rather than fewer.
```python
# IdMapIndex — allowlist of external ids (typical use)
allowed = np.array([1003, 1010, 1042], dtype=np.uint64)
scores, ids = idx.search(queries, k=10, allowlist=allowed)
# scores.shape == (nq, min(k, len(allowed))) == (nq, 3)
# TurboQuantIndex — bool mask over slots
mask = np.ones(len(idx), dtype=bool)
mask[disabled_slots] = False
scores, slots = idx.search(queries, k=10, mask=mask)
```
The output shape is `(nq, min(k, n_allowed))` — same shrinking behaviour you already see when `k > len(idx)`. No `-1` / `NaN` padding; pad on the caller side if you need a fixed-width batch.
Common use cases:
- Hybrid retrieval where a SQL/BM25 stage produces a candidate id set.
- Access control or multi-tenant queries (only return ids the caller can see).
- Time-windowed search (e.g. only documents from the last 7 days).
---
## File formats
### `.tv` — `TurboQuantIndex`
```
┌──────────────────────────────────────┐
│ magic "TVPI" (4 bytes) │
│ version u8 = 3 │
├──────────────────────────────────────┤
│ core header │
│ bit_width (u8) │
│ dim (u32 LE) │
│ n_vectors (u32 LE) │
├──────────────────────────────────────┤
│ packed codes │
│ (dim / 8) * bit_width * n_vectors │
├──────────────────────────────────────┤
│ scales (n_vectors × f32 LE) │
│ per-vector length-renormalization │
├──────────────────────────────────────┤
│ TQ+ trailer │
│ n_calib (u32 LE) — 0 or dim │
│ shift (n_calib × f32 LE) │
│ scale (n_calib × f32 LE) │
└──────────────────────────────────────┘
```
### `.tvim` — `IdMapIndex`
```
┌──────────────────────────────────────┐
│ magic "TVIM" (4 bytes) │
│ version u8 = 3 │
├──────────────────────────────────────┤
│ core payload (same as .tv: │
│ header + codes + scales + TQ+) │
├──────────────────────────────────────┤
│ slot_to_id (n_vectors × u64 LE) │
└──────────────────────────────────────┘
```
On load, the reverse `id → slot` map is rebuilt in memory. Duplicate ids in the `slot_to_id` table are rejected as corrupt.
Both `.tv` and `.tvim` loads validate the header **before allocating**: `bit_width` must be 2/3/4, `dim` a positive multiple of 8 and `≤ 65536`, and every payload size is computed with checked arithmetic and read through a length-capped reader. A malformed or untrusted file therefore raises a clean error rather than panicking, dividing by zero, or driving an oversized allocation.
`n_calib = 0` in the TQ+ trailer means identity calibration (a lazy index with no `add` yet, or a pre-TQ+ index that was re-saved); otherwise it equals `dim`. Loading a version-2 file (no TQ+ trailer) is still supported and is read as identity calibration; version 1 (headerless, no magic) is rejected.
`dim = 0` in the core header signals a lazy uncommitted index. It is only valid alongside `n_vectors = 0`; on load it produces an index whose `dim` is `None` until the first `add` / `add_with_ids` call.
Both formats carry a magic + version byte and are stable across minor versions. Breaking changes bump the version byte.
docs/arm_speed_mt.svg
+61
View File
@@ -0,0 +1,61 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Search Latency — ARM (Apple M3 Max) — Multi-threaded">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Search Latency — ARM (Apple M3 Max) — Multi-threaded</text>
<text x="84" y="52" class="subtitle">100K vectors, 1K queries, k=64, median of 5 runs</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0.00</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">0.20</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">0.40</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">0.60</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">0.80</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">1.00</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<text x="84" y="68" class="panel">Multi-threaded</text>
<rect x="135.0" y="324.2" width="44" height="27.8" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="185.0" y="320.9" width="44" height="31.1" rx="6" fill="#9aa7b6" />
<text x="157.0" y="318.2" text-anchor="middle" class="value-accent">0.103</text>
<text x="207.0" y="314.9" text-anchor="middle" class="value">0.115</text>
<text x="182.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="182.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="331.0" y="302.1" width="44" height="50.0" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="381.0" y="292.6" width="44" height="59.4" rx="6" fill="#9aa7b6" />
<text x="353.0" y="296.1" text-anchor="middle" class="value-accent">0.185</text>
<text x="403.0" y="286.6" text-anchor="middle" class="value">0.220</text>
<text x="378.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="378.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<rect x="527.0" y="297.7" width="44" height="54.3" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="577.0" y="291.5" width="44" height="60.5" rx="6" fill="#9aa7b6" />
<text x="549.0" y="291.7" text-anchor="middle" class="value-accent">0.201</text>
<text x="599.0" y="285.5" text-anchor="middle" class="value">0.224</text>
<text x="574.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="574.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="723.0" y="250.8" width="44" height="101.2" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="773.0" y="231.0" width="44" height="121.0" rx="6" fill="#9aa7b6" />
<text x="745.0" y="244.8" text-anchor="middle" class="value-accent">0.375</text>
<text x="795.0" y="225.0" text-anchor="middle" class="value">0.448</text>
<text x="770.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="770.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<text x="26" y="217.0" transform="rotate(-90, 26, 217.0)" class="axis">ms / query</text>
<rect x="84" y="424" width="14" height="14" rx="3" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="106" y="435" class="legend" style="fill: #4338ca;">TurboQuant</text>
<rect x="224" y="424" width="14" height="14" rx="3" fill="#9aa7b6" />
<text x="246" y="435" class="legend">FAISS</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 4.8 KiB

docs/arm_speed_st.svg
+61
View File
@@ -0,0 +1,61 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Search Latency — ARM (Apple M3 Max) — Single-threaded">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Search Latency — ARM (Apple M3 Max) — Single-threaded</text>
<text x="84" y="52" class="subtitle">100K vectors, 1K queries, k=64, median of 5 runs</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0.0</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">2.0</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">4.0</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">6.0</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">8.0</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">10.0</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<text x="84" y="68" class="panel">Single-threaded</text>
<rect x="135.0" y="322.8" width="44" height="29.2" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="185.0" y="318.7" width="44" height="33.3" rx="6" fill="#9aa7b6" />
<text x="157.0" y="316.8" text-anchor="middle" class="value-accent">1.08</text>
<text x="207.0" y="312.7" text-anchor="middle" class="value">1.24</text>
<text x="182.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="182.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="331.0" y="298.2" width="44" height="53.8" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="381.0" y="285.9" width="44" height="66.2" rx="6" fill="#9aa7b6" />
<text x="353.0" y="292.2" text-anchor="middle" class="value-accent">1.99</text>
<text x="403.0" y="279.9" text-anchor="middle" class="value">2.45</text>
<text x="378.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="378.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<rect x="527.0" y="294.7" width="44" height="57.3" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="577.0" y="286.1" width="44" height="65.9" rx="6" fill="#9aa7b6" />
<text x="549.0" y="288.7" text-anchor="middle" class="value-accent">2.12</text>
<text x="599.0" y="280.1" text-anchor="middle" class="value">2.44</text>
<text x="574.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="574.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="723.0" y="244.9" width="44" height="107.1" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="773.0" y="219.0" width="44" height="133.0" rx="6" fill="#9aa7b6" />
<text x="745.0" y="238.9" text-anchor="middle" class="value-accent">3.97</text>
<text x="795.0" y="213.0" text-anchor="middle" class="value">4.92</text>
<text x="770.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="770.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<text x="26" y="217.0" transform="rotate(-90, 26, 217.0)" class="axis">ms / query</text>
<rect x="84" y="424" width="14" height="14" rx="3" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="106" y="435" class="legend" style="fill: #4338ca;">TurboQuant</text>
<rect x="224" y="424" width="14" height="14" rx="3" fill="#9aa7b6" />
<text x="246" y="435" class="legend">FAISS</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 4.8 KiB

docs/compression.svg
+62
View File
@@ -0,0 +1,62 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Index Size — TurboQuant">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Index Size — 100K vectors</text>
<text x="84" y="52" class="subtitle">TurboQuant packs vectors ~16× smaller than FP32 at 2-bit with comparable recall</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">300</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">600</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">900</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">1200</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">1500</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<rect x="120.7" y="338.3" width="56" height="13.7" rx="6" fill="#9aa7b6" />
<text x="148.7" y="332.3" text-anchor="middle" class="value">76</text>
<rect x="186.7" y="350.2" width="56" height="1.8" rx="6" fill="#1d4ed8" />
<text x="214.7" y="344.2" text-anchor="middle" class="value">10</text>
<rect x="252.7" y="351.1" width="56" height="0.9" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="280.7" y="345.1" text-anchor="middle" class="value-accent">5</text>
<text x="214.7" y="374" text-anchor="middle" class="label">GloVe</text>
<text x="214.7" y="389" text-anchor="middle" class="secondary">d=200</text>
<rect x="382.0" y="246.5" width="56" height="105.5" rx="6" fill="#9aa7b6" />
<text x="410.0" y="240.5" text-anchor="middle" class="value">586</text>
<rect x="448.0" y="338.8" width="56" height="13.2" rx="6" fill="#1d4ed8" />
<text x="476.0" y="332.8" text-anchor="middle" class="value">74</text>
<rect x="514.0" y="345.3" width="56" height="6.7" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="542.0" y="339.3" text-anchor="middle" class="value-accent">37</text>
<text x="476.0" y="374" text-anchor="middle" class="label">OpenAI</text>
<text x="476.0" y="389" text-anchor="middle" class="secondary">d=1536</text>
<rect x="643.3" y="141.1" width="56" height="210.9" rx="6" fill="#9aa7b6" />
<text x="671.3" y="135.1" text-anchor="middle" class="value">1172</text>
<rect x="709.3" y="325.6" width="56" height="26.4" rx="6" fill="#1d4ed8" />
<text x="737.3" y="319.6" text-anchor="middle" class="value">147</text>
<rect x="775.3" y="338.8" width="56" height="13.2" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="803.3" y="332.8" text-anchor="middle" class="value-accent">74</text>
<text x="737.3" y="374" text-anchor="middle" class="label">OpenAI</text>
<text x="737.3" y="389" text-anchor="middle" class="secondary">d=3072</text>
<text x="26" y="217.0" transform="rotate(-90, 26, 217.0)" class="axis">Index size (MB)</text>
<rect x="84" y="424" width="14" height="14" rx="3" fill="#9aa7b6" />
<text x="106" y="435" class="legend">FP32</text>
<rect x="204" y="424" width="14" height="14" rx="3" fill="#1d4ed8" />
<text x="226" y="435" class="legend">4-bit</text>
<rect x="324" y="424" width="14" height="14" rx="3" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="346" y="435" class="legend">2-bit</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 4.7 KiB

docs/header.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 313 KiB

docs/integrations/agno.md
+143
View File
@@ -0,0 +1,143 @@
# Agno integration
`turbovec.agno.TurboQuantVectorDb` is an [Agno](https://github.com/agno-agi/agno) `VectorDb` backed by an `IdMapIndex`. It implements the same public surface as `agno.vectordb.lancedb.LanceDb` (the closest in-tree single-machine backend) so this can be swapped in wherever LanceDb is used.
## Install
```bash
pip install turbovec[agno]
```
## Basic usage
```python
from agno.agent import Agent
from agno.knowledge import Knowledge
from agno.knowledge.embedder.openai import OpenAIEmbedder
from turbovec.agno import TurboQuantVectorDb
vector_db = TurboQuantVectorDb(embedder=OpenAIEmbedder())
knowledge = Knowledge(vector_db=vector_db)
knowledge.load_text("Turbovec compresses vectors to 4 bits per dimension.")
agent = Agent(knowledge=knowledge)
agent.print_response("What does turbovec do?")
```
## Constructor
```python
TurboQuantVectorDb(
*,
id: Optional[str] = None,
name: Optional[str] = None,
description: Optional[str] = None,
similarity_threshold: Optional[float] = None,
embedder: Embedder, # required
bit_width: int = 4,
search_type: SearchType = SearchType.vector,
distance: Distance = Distance.cosine,
reranker: Optional[Reranker] = None,
path: Optional[str] = None,
)
```
| Parameter | Notes |
|---|---|
| `embedder` | **Required.** Source of truth for the embedding dimension — `embedder.dimensions` sizes the underlying quantized index. |
| `bit_width` | Quantization width per coordinate; one of `{2, 4}`. |
| `search_type` | Only `SearchType.vector` is supported. Constructing with `keyword` or `hybrid` raises `ValueError` (keyword/hybrid would require an external BM25/lexical index that turbovec doesn't ship). |
| `distance` | Only `Distance.cosine` is supported. turbovec stores unit-normalized vectors and the kernel's raw score is cosine similarity directly. |
| `similarity_threshold` | Optional. Scores are mapped from cosine `[-1, 1]` to relevance `[0, 1]` via `(s + 1) / 2`; results below the threshold are dropped. |
| `reranker` | Optional Agno reranker applied to the result set after vector retrieval. |
| `path` | Optional directory for save/load persistence. When given, `create()` loads existing data from this path if present. |
## Insert / upsert
`insert` and `upsert` follow the same `(content_hash, documents, filters)` signature as `LanceDb`. The internal `doc_id` is derived as `md5(f"{base_id}_{content_hash}")` where `base_id` is `doc.id` (or `md5(content)` when missing). The contract: the same `(base_id, content_hash)` pair always produces the same internal id, and the same `base_id` with a *different* `content_hash` is treated as a new entry — letting you keep content versions side-by-side.
Because `doc_id` is derived from `base_id` + `content_hash` (not from `name`, `content_id`, or metadata), two documents can collide on the same `doc_id` — a repeated explicit `doc.id`, or two documents with identical content and no id. When that happens **both are stored and both remain individually deletable** — keep-all, matching `LanceDb`'s append-only behavior. (This differs from the LangChain store, which keeps the last write per id.)
```python
from agno.knowledge.document import Document
docs = [Document(id="doc-1", name="paper.pdf", content="...", meta_data={"source": "arxiv"})]
vector_db.insert(content_hash="v1", documents=docs)
# Same doc with a new content_hash → new stored entry.
vector_db.insert(content_hash="v2", documents=docs)
```
Documents without embeddings are embedded via `self.embedder` before insertion. If embedding fails (`get_embedding` returns `None`) the call raises `ValueError` rather than silently dropping the document.
## Filtered search
Filters are resolved to an allowlist **before** scoring — the kernel only ever inserts allowed candidates into the per-query heap. You always get up to `limit` results from the filtered set; no over-fetching, no recall hit on selective filters.
```python
results = vector_db.search(
"quantum computing applications",
limit=5,
filters={"source": "arxiv", "year": 2024}, # AND of exact equality
)
```
Dict filters use AND-of-exact-equality on `Document.meta_data`. List-style `FilterExpr` filters (Agno's structured filter type) are silently ignored, matching `LanceDb`'s behaviour.
## Existence checks
```python
vector_db.name_exists("paper.pdf") # bool — by Document.name
vector_db.id_exists("derived-md5-id") # bool — by the internally-derived id
vector_db.content_hash_exists("v1") # O(1) — set lookup, not a scan
```
## Delete
```python
vector_db.delete_by_id(derived_id) # by internal id
vector_db.delete_by_name("paper.pdf") # by Document.name
vector_db.delete_by_metadata({"source": "web"}) # AND-of-equality on meta_data
vector_db.delete_by_content_id("cid-42") # by Document.content_id
vector_db.drop() # clear all
vector_db.delete() # alias for drop(), returns True
```
Each `delete_by_*` returns `True` iff at least one document was removed. `delete_by_name` / `delete_by_content_id` / `delete_by_metadata` remove only the documents matching that exact predicate, even when other stored documents share the same derived `doc_id`. `delete_by_id` removes every document under that internal id.
## update_metadata
```python
vector_db.update_metadata("cid-42", {"reviewed": True})
```
Merges the given metadata into `meta_data` of every document with the matching `content_id`. Overrides the base class's no-op warning.
## Save / load
```python
vector_db = TurboQuantVectorDb(embedder=embedder, path="./my-store")
vector_db.create() # loads from path if existing
# ... insert documents ...
vector_db.save() # persists to path
```
Writes two files under the given folder path:
- `index.tvim` — the `IdMapIndex` payload.
- `docstore.json` — JSON-encoded document text, metadata, and id maps.
Document metadata must be JSON-serializable — same constraint Agno's `LanceDb` imposes on its payload column. The side-car carries a `schema_version` field; loaders refuse to deserialize unknown versions, and validate that the side-car's id maps are consistent with the loaded `index.tvim` (a mismatched or out-of-sync pair raises at load rather than failing later at query time).
## Async
The lifecycle, write, and read methods have async counterparts: `async_create`, `async_drop`, `async_exists`, `async_name_exists`, `async_get_count`, `async_insert`, `async_upsert`, `async_search`. The remaining methods (the `delete_by_*` family, `update_metadata`, `save`, `id_exists`, `content_hash_exists`, `optimize`) are sync-only. When the embedder exposes `async_get_embedding` / `async_get_embeddings_batch_and_usage`, the async paths use it for genuine async embedding generation.
## Known limitations
- **Vector search only.** `search_type=SearchType.keyword` and `SearchType.hybrid` are not supported (would require an external BM25 / lexical index). Constructor raises `ValueError` on those.
- **Cosine distance only.** `Distance.cosine` is the only supported metric. turbovec stores unit-normalized vectors; other distances would require non-trivial scoring changes.
- **Embeddings are not retained after quantization.** Stored vectors are the quantized form; the original full-precision embedding can't be recovered.
- **JSON-serializable metadata only.** Non-JSON-serializable values fail at `save()` time.
docs/integrations/haystack.md
+171
View File
@@ -0,0 +1,171 @@
# Haystack integration
`turbovec.haystack.TurboQuantDocumentStore` is a Haystack 2.x [`DocumentStore`](https://docs.haystack.deepset.ai/docs/document-store) backed by an `IdMapIndex`. It implements the same public surface as `haystack.document_stores.in_memory.InMemoryDocumentStore` and can be used as a drop-in replacement wherever the in-memory store is used.
## Install
```bash
pip install turbovec[haystack]
```
## Basic usage
```python
from haystack import Document
from turbovec.haystack import TurboQuantDocumentStore
store = TurboQuantDocumentStore()
store.write_documents([
Document(content="...", embedding=[...], meta={"source": "a"}),
Document(content="...", embedding=[...], meta={"source": "b"}),
])
results = store.embedding_retrieval(query_embedding=[...], top_k=5)
```
Documents must have pre-computed embeddings — `TurboQuantDocumentStore` doesn't invoke an embedder. Pipe a Haystack embedder component upstream if your documents arrive without embeddings.
## Constructor
```python
TurboQuantDocumentStore(
dim: Optional[int] = None,
bit_width: int = 4,
*,
embedding_similarity_function: Literal["dot_product", "cosine"] = "cosine",
async_executor: Optional[ThreadPoolExecutor] = None,
return_embedding: bool = False,
)
```
| Parameter | Notes |
|---|---|
| `dim` | Optional. When omitted the vector dimensionality is inferred from the first `write_documents` call. |
| `bit_width` | Quantization width per coordinate; one of `{2, 4}`. |
| `embedding_similarity_function` | Drives the `scale_score=True` formula on retrieval. Defaults to `"cosine"` (right for unit-normalized embeddings); `"dot_product"` uses Haystack's `expit(s / 100)` formula. |
| `async_executor` | Optional `ThreadPoolExecutor` for the `*_async` methods. If omitted, a single-threaded executor is created and cleaned up with the store. |
| `return_embedding` | Accepted for API parity with `InMemoryDocumentStore`. The full-precision embedding is never available (quantized away), so `Document.embedding` on retrieved docs is always `None` regardless of the flag. |
## `DuplicatePolicy`
`write_documents` takes a `policy` argument controlling how id collisions are handled:
```python
from haystack.document_stores.types import DuplicatePolicy
store.write_documents(docs, policy=DuplicatePolicy.FAIL) # raise if any id collides
store.write_documents(docs, policy=DuplicatePolicy.SKIP) # silently skip colliding ids
store.write_documents(docs, policy=DuplicatePolicy.OVERWRITE) # remove-then-re-add colliding ids
# DuplicatePolicy.NONE is treated as FAIL.
```
Returns the number of documents actually written (so `SKIP` may return less than `len(docs)`).
## Delete
```python
store.delete_documents(["id-1", "id-2"]) # by id; missing ids are silently ignored
store.delete_by_filter(filters) # by filter; returns count
store.delete_all_documents() # clear everything
```
`delete_documents` and `delete_by_filter` are O(1) per matching document via the inner `IdMapIndex`.
## Filters
`filter_documents(filters)`, `embedding_retrieval(..., filters=...)`, and the other filter-aware helpers accept the full [Haystack filter DSL](https://docs.haystack.deepset.ai/docs/metadata-filtering):
```python
filters = {
"operator": "AND",
"conditions": [
{"field": "meta.source", "operator": "==", "value": "manual"},
{"field": "meta.version", "operator": ">=", "value": 2},
],
}
# All docs matching the filter (no vector search):
docs = store.filter_documents(filters=filters)
# Top-k nearest to a query, filtered:
results = store.embedding_retrieval(
query_embedding=[...],
top_k=5,
filters=filters,
)
```
Filter evaluation is delegated to `haystack.utils.filters.document_matches_filter` — anything Haystack's own stores support, we support.
For `embedding_retrieval`, filters are resolved to an allowlist **before** scoring rather than via post-filtering. Selective filters return up to `top_k` matches from the filtered set; you never get fewer than `top_k` results just because the filter happened to exclude the top-scoring candidates.
## Metadata helpers
```python
store.count_documents_by_filter(filters) # int
store.count_unique_metadata_by_filter(filters, ["source", "tag"]) # dict[str, int]
store.update_by_filter(filters, {"reviewed": True}) # bulk metadata update; returns count
store.get_metadata_fields_info()
# {"source": {"type": "keyword"}, "version": {"type": "int"}, ...}
store.get_metadata_field_min_max("version") # {"min": 1, "max": 5}
store.get_metadata_field_unique_values("source")
# (["a", "b", "c"], 3)
```
`update_by_filter` updates metadata only — embeddings are quantized at write time and not re-encoded.
## Async
Every public method has an `*_async` variant:
```python
await store.write_documents_async(docs)
results = await store.embedding_retrieval_async(query_embedding=q, top_k=5)
await store.delete_documents_async(["id-1"])
```
By default they run on a single-threaded executor owned by the store. Pass an `async_executor=` to the constructor to share an executor across stores (or to use more workers).
## Save / load
```python
store.save_to_disk("./my-store")
# ... later ...
store = TurboQuantDocumentStore.load_from_disk("./my-store")
```
Writes two files under the given folder path:
- `index.tvim` — the `IdMapIndex` payload (quantized vectors + id maps).
- `docstore.json` — JSON-encoded document text, metadata, and id maps.
Document metadata must be JSON-serializable — the same constraint `InMemoryDocumentStore.save_to_disk` imposes. If the `docstore.json` side-car is out of sync with its `index.tvim` (a partial copy, a stale backup, tampering), `load_from_disk` raises a `ValueError` immediately rather than failing later with a `KeyError` at query time.
## Using in a Haystack Pipeline
`TurboQuantDocumentStore` implements `to_dict` / `from_dict` so it can be serialized as part of a Haystack `Pipeline`. `to_dict` captures the component *config* (`dim`, `bit_width`, `embedding_similarity_function`, `return_embedding`); persisting the stored documents is the job of `save_to_disk` / `load_from_disk`.
Plug into a standard RAG pipeline the same way you'd use `InMemoryDocumentStore`:
```python
from haystack import Pipeline
from haystack.components.embedders import SentenceTransformersDocumentEmbedder
from haystack.components.writers import DocumentWriter
store = TurboQuantDocumentStore() # dim inferred from first batch
indexing = Pipeline()
indexing.add_component("embedder", SentenceTransformersDocumentEmbedder(
model="sentence-transformers/all-MiniLM-L6-v2",
))
indexing.add_component("writer", DocumentWriter(document_store=store))
indexing.connect("embedder.documents", "writer.documents")
indexing.run({"embedder": {"documents": my_docs}})
```
## Known limitations
- **Embeddings are not retained.** `embedding_retrieval(..., return_embedding=True)` is accepted for signature compatibility but `Document.embedding` is always `None` on retrieved docs — turbovec discards the full-precision vector after quantization.
- **JSON-serializable metadata only.** Document metadata is stored as JSON in the side-car. Non-JSON-serializable values (custom objects, sets, etc.) fail at save time — the same constraint `InMemoryDocumentStore.save_to_disk` imposes.
- **`dim` is locked on the first add.** Subsequent calls with a different shape raise `ValueError`. If you need to change `dim`, construct a fresh store.
docs/integrations/langchain.md
+144
View File
@@ -0,0 +1,144 @@
# LangChain integration
`turbovec.langchain.TurboQuantVectorStore` is a [LangChain `VectorStore`](https://python.langchain.com/docs/integrations/vectorstores/) backed by an `IdMapIndex`. It implements the same public surface as `langchain_core.vectorstores.in_memory.InMemoryVectorStore` and can be used as a drop-in replacement wherever the in-memory store is used.
## Install
```bash
pip install turbovec[langchain]
```
## Basic usage
```python
from langchain_huggingface import HuggingFaceEmbeddings
from turbovec.langchain import TurboQuantVectorStore
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
store = TurboQuantVectorStore.from_texts(
texts=["Document 1...", "Document 2...", "Document 3..."],
embedding=embeddings,
bit_width=4,
)
retriever = store.as_retriever(search_kwargs={"k": 5})
```
The dimensionality of the underlying quantized index is inferred from the embedding model on the first `add_*` call — no need to specify it up front.
## Construction
```python
# No-arg: lazy. dim is inferred from the first add.
store = TurboQuantVectorStore(embeddings)
# from_texts: same lazy behaviour, plus immediate ingest.
store = TurboQuantVectorStore.from_texts(texts, embeddings, bit_width=4)
# Pre-built index: bring your own IdMapIndex (e.g. one loaded from disk).
from turbovec import IdMapIndex
store = TurboQuantVectorStore(embeddings, index=IdMapIndex(1536, 4))
```
`bit_width` is `2` or `4` and is fixed once the index is created.
## Adding with explicit ids
```python
store.add_texts(
texts=["a", "b", "c"],
ids=["doc-a", "doc-b", "doc-c"],
metadatas=[{"source": "x"}, {"source": "y"}, {"source": "z"}],
)
# add_documents honours per-Document.id, falling back to a UUID per
# document if .id is missing — partial ids are not dropped wholesale.
store.add_documents([
Document(id="explicit", page_content="..."),
Document(page_content="..."), # gets a UUID
])
```
If an id is already present, `add_texts` **upserts** — the existing entry is removed and the new one added with the same id. This matches the typical user expectation that re-indexing a document with the same id should replace it, not duplicate it.
Async equivalents (`aadd_texts`, `aadd_documents`) use the embedding model's `aembed_documents` so they benefit from concurrent embedding generation when the model supports it.
## Search
```python
# By string query (uses the embedding function)
docs = store.similarity_search("what is turbovec?", k=5)
# With scores
docs_and_scores = store.similarity_search_with_score("...", k=5)
# By raw vector
import numpy as np
qvec = np.random.randn(768).astype(np.float32)
qvec /= np.linalg.norm(qvec)
docs = store.similarity_search_by_vector(qvec.tolist(), k=5)
```
Scores are raw inner products. Because vectors are L2-normalized on insert, inner product equals cosine similarity — higher is better, range `[-1, 1]`.
`similarity_search_with_relevance_scores` and `as_retriever(search_type="similarity_score_threshold")` work: the raw cosine is mapped to `[0, 1]` via `(sim + 1) / 2` (clamped to absorb the tiny overshoot caused by quantization noise).
Async equivalents (`asimilarity_search`, `asimilarity_search_with_score`, `asimilarity_search_by_vector`, `aget_by_ids`) are all implemented.
## Filters
`similarity_search`, `similarity_search_with_score`, and `similarity_search_by_vector` all accept a `filter` keyword:
```python
# Dict — AND of exact equality on Document.metadata.
docs = store.similarity_search(
"query", k=5, filter={"source": "manual", "version": 2},
)
# Callable — predicate over the Document.
docs = store.similarity_search(
"query", k=5, filter=lambda doc: doc.metadata.get("score", 0) > 0.8,
)
```
The callable form matches the `Callable[[Document], bool]` convention used by `InMemoryVectorStore`, so predicates ported from there work unchanged.
Filters are resolved to an id allowlist **before** scoring; the kernel only ever inserts allowed documents into the per-query heap. You get up to `k` results from the filtered set, never fewer than `k` because the filter happened to exclude the top-scoring candidates.
## Document retrieval by id
```python
docs = store.get_by_ids(["doc-a", "doc-c"])
# Missing ids are silently skipped.
```
`aget_by_ids` is also available.
## Delete
```python
store.delete(["doc-a", "doc-b"]) # missing ids silently skipped, returns None
```
Delete is O(1) per id. `delete(None)` is a no-op (matches the `InMemoryVectorStore` contract).
## Save / load
```python
store.dump("./my-store")
# ... later ...
store = TurboQuantVectorStore.load("./my-store", embedding=embeddings)
```
Writes two files under the given folder path:
- `index.tvim` — the `IdMapIndex` payload (see [api.md](../api.md#tvim--idmapindex)).
- `docstore.json` — JSON-encoded document text, metadata, and id maps.
Document metadata must be JSON-serializable — the same constraint `InMemoryVectorStore.dump` imposes. If the `docstore.json` side-car is out of sync with its `index.tvim` (a partial copy, a stale backup, tampering), `load` raises a `ValueError` immediately rather than failing later with a `KeyError` at query time.
## Known limitations
- **Max-marginal-relevance search is not supported.** `max_marginal_relevance_search` and its variants raise `NotImplementedError` with an explanation. MMR requires the full-precision embedding of each candidate to compute pairwise diversity; turbovec discards full-precision vectors after quantization. If you need MMR, keep a parallel store with the raw embeddings and run MMR over that.
- **Embeddings are not retained.** `search` returns `Document` objects with `page_content` and `metadata`, but the original embedding is not recoverable.
- **JSON-serializable metadata only.** Non-JSON-serializable values (custom objects, sets, etc.) fail at save time — same constraint as the in-tree reference store.
docs/integrations/llama_index.md
+218
View File
@@ -0,0 +1,218 @@
# LlamaIndex integration
`turbovec.llama_index.TurboQuantVectorStore` is a LlamaIndex [`BasePydanticVectorStore`](https://docs.llamaindex.ai/en/stable/module_guides/storing/vector_stores/) backed by an `IdMapIndex`. It implements the same public surface as `llama_index.core.vector_stores.simple.SimpleVectorStore` and can be used as a drop-in replacement wherever the simple in-memory store is used.
## Install
```bash
pip install turbovec[llama-index]
```
## Basic usage
```python
from llama_index.core import VectorStoreIndex, StorageContext
from turbovec.llama_index import TurboQuantVectorStore
vector_store = TurboQuantVectorStore()
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
retriever = index.as_retriever(similarity_top_k=5)
```
The vector dimensionality is inferred from the embedding model on the first `add()` call.
## Construction
```python
# No-arg: lazy. dim is inferred from the first add.
vector_store = TurboQuantVectorStore()
# from_params: same lazy behaviour, plus an explicit bit_width.
vector_store = TurboQuantVectorStore.from_params(bit_width=4)
# Pre-built index: bring your own IdMapIndex (e.g. one you loaded from disk).
from turbovec import IdMapIndex
vector_store = TurboQuantVectorStore(index=IdMapIndex(1536, 4))
```
`bit_width` is `2` or `4` and is fixed once the index is created.
## The two `delete` signatures
LlamaIndex's vector-store protocol has two distinct delete entry points:
### `delete(ref_doc_id: str)` — remove an entire source document
Removes **every node** whose `ref_doc_id` matches. Use this when you want to delete a whole parent document and its chunks in one call.
```python
vector_store.delete("my-source-document-123")
```
Missing `ref_doc_id`s are silently ignored.
### `delete_nodes(node_ids, filters)` — remove specific chunks
Removes nodes matching either `node_ids`, `filters`, or both (intersected). Missing `node_id`s are silently ignored.
```python
# By node_id
vector_store.delete_nodes(node_ids=["abc-123", "def-456"])
# By metadata filter
from llama_index.core.vector_stores.types import (
MetadataFilter, MetadataFilters, FilterOperator,
)
filters = MetadataFilters(
filters=[MetadataFilter(key="tier", value="archived", operator=FilterOperator.EQ)],
)
vector_store.delete_nodes(filters=filters)
# Both: intersect — delete only nodes in this list that ALSO match the filter
vector_store.delete_nodes(node_ids=["abc-123"], filters=filters)
```
### `clear()` — drop everything
```python
vector_store.clear()
```
Resets the store while preserving the configured `bit_width`. The cleared store is immediately usable for new adds; `dim` is inferred again from the next batch.
## Query
LlamaIndex calls `query(VectorStoreQuery)` internally. If you've gone through `VectorStoreIndex.from_documents(...)`, you won't call this directly — the retriever does. For direct use:
```python
from llama_index.core.vector_stores.types import VectorStoreQuery
result = vector_store.query(VectorStoreQuery(
query_embedding=[...],
similarity_top_k=5,
))
# result.nodes, result.similarities, result.ids
```
`query_embedding` is **required**. turbovec doesn't embed query text itself; the calling component (retriever / query engine) is responsible for that.
### Filtered query
`VectorStoreQuery` accepts `filters`, `node_ids`, and `doc_ids`. All three intersect when more than one is supplied:
```python
from llama_index.core.vector_stores.types import (
MetadataFilter, MetadataFilters, FilterCondition, FilterOperator,
VectorStoreQuery,
)
filters = MetadataFilters(
filters=[
MetadataFilter(key="tier", value="pro", operator=FilterOperator.EQ),
MetadataFilter(key="year", value=2024, operator=FilterOperator.GTE),
],
condition=FilterCondition.AND,
)
result = vector_store.query(VectorStoreQuery(
query_embedding=[...],
similarity_top_k=5,
filters=filters,
node_ids=["chunk-1", "chunk-2", "chunk-3"], # restrict to these chunks
doc_ids=["src-doc-42"], # restrict to chunks of this source doc
))
```
Supported operators on `MetadataFilter`: `EQ`, `NE`, `GT`, `LT`, `GTE`, `LTE`, `IN`, `NIN`, `TEXT_MATCH`, `TEXT_MATCH_INSENSITIVE`, `CONTAINS`, `ANY`, `ALL`, `IS_EMPTY`. Conditions: `AND`, `OR`, `NOT`. Nested `MetadataFilters` work.
Filter semantics match `SimpleVectorStore`'s reference implementation — notably, every operator except `IS_EMPTY` returns `False` when the filter key is missing from the document's metadata, and `TEXT_MATCH` is case-sensitive (use `TEXT_MATCH_INSENSITIVE` for a case-insensitive substring match).
Filters are resolved to a handle allowlist **before** scoring. Selective filters return up to `similarity_top_k` matches from the filtered set; you never get fewer just because the filter happened to exclude the top-scoring candidates.
## Get nodes
```python
nodes = vector_store.get_nodes(node_ids=["chunk-1", "chunk-2"])
nodes = vector_store.get_nodes(filters=filters)
nodes = vector_store.get_nodes(node_ids=["chunk-1", "chunk-2"], filters=filters) # intersect
```
Returns a `List[BaseNode]` reconstructed from the side-car. Missing `node_id`s are silently skipped.
## Upsert semantics
Calling `add()` with a node whose `node_id` already exists **replaces** the existing entry. Matches LlamaIndex user expectation when re-indexing the same chunks.
A `node_id` repeated **within a single `add()` batch** raises `ValueError` — deduplicate before calling. (This differs from the LangChain and Haystack stores, which silently keep the last occurrence; here it's a hard error so an accidental duplicate doesn't quietly drop a node.)
```python
node = TextNode(text="v1", embedding=[...])
vector_store.add([node])
# Same node_id, different text/embedding → replaces.
updated = TextNode(text="v2", id_=node.node_id, embedding=[...])
vector_store.add([updated])
assert len(vector_store._index) == 1
```
## Async
Every public method has an async counterpart, suitable for use in LlamaIndex's async retriever / query-engine paths:
```python
await vector_store.async_add(nodes)
result = await vector_store.aquery(VectorStoreQuery(...))
fetched = await vector_store.aget_nodes(node_ids=[...])
await vector_store.adelete("ref-doc-id")
await vector_store.adelete_nodes(node_ids=[...])
await vector_store.aclear()
```
## Persist / load
### Direct (file-stem) interface
```python
vector_store.persist("./store/vectors.json")
# ... later ...
vector_store = TurboQuantVectorStore.from_persist_path("./store/vectors.json")
```
`persist_path` is treated as a path *stem* — the binary index and JSON side-car are written next to each other as `{stem}.tvim` and `{stem}.nodes.json`. The extension on `persist_path` (e.g. `.json`, as LlamaIndex's StorageContext default uses) is replaced. Node metadata must be JSON-serializable. If the `{stem}.nodes.json` side-car is out of sync with its `{stem}.tvim` index (a partial copy, a stale backup, tampering), `from_persist_path` raises a `ValueError` immediately rather than failing later with a `KeyError` at query time.
### Via `StorageContext`
The store works with `StorageContext.from_defaults(persist_dir=...)` the same way `SimpleVectorStore` does:
```python
# Persist
storage_context.persist(persist_dir="./store")
# Load
vector_store = TurboQuantVectorStore.from_persist_dir(persist_dir="./store")
storage_context = StorageContext.from_defaults(
vector_store=vector_store,
persist_dir="./store",
)
```
`from_persist_dir(persist_dir, namespace="default", fs=None)` constructs the namespaced filename (`{persist_dir}/{namespace}__vector_store.json`) and delegates to `from_persist_path`. Multiple namespaced stores can share a persist directory.
### Config-only round-trip
```python
config = vector_store.to_dict() # {"bit_width": 4, "dim": 1536}
fresh = TurboQuantVectorStore.from_dict(config) # empty store with the same config
```
`to_dict` / `from_dict` serialize only the store's configuration. Node data round-trips through `persist` / `from_persist_path`.
## Known limitations
- **MMR is not supported.** Max-marginal-relevance retrieval requires the full-precision embedding of each candidate to compute pairwise diversity; turbovec discards full-precision vectors after quantization.
- **`get(text_id)` raises** rather than returning a vector — same reason. The full-precision embedding is not recoverable.
- **`fsspec` filesystems are not supported.** `persist`, `from_persist_path`, and `from_persist_dir` accept a local path. Pass `fs=None` (the default).
- **JSON-serializable metadata only.** Node metadata is stored as JSON in the side-car. Non-JSON-serializable values fail at persist time — same constraint as `SimpleVectorStore.persist`.
- **`stores_text = True`.** Unlike `SimpleVectorStore`, we keep node text in the side-car so query results return populated `TextNode`s without depending on a separate docstore. If you're swapping this in for `SimpleVectorStore` and your pipeline expects text to live elsewhere, the difference is harmless — the framework treats `stores_text` as informational.
docs/recall_d1536.svg
+85
View File
@@ -0,0 +1,85 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Recall — d=1536">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Recall — d=1536</text>
<text x="84" y="52" class="subtitle">100K vectors, k=64 search. recall@1@k measures how often the true top-1 result appears in the top-k returned.</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0.85</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">0.88</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">0.91</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">0.94</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">0.97</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">1.00</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<text x="84" y="68" class="panel">d=1536</text>
<text x="84.0" y="372" text-anchor="middle" class="label">1</text>
<text x="214.7" y="372" text-anchor="middle" class="label">2</text>
<text x="345.3" y="372" text-anchor="middle" class="label">4</text>
<text x="476.0" y="372" text-anchor="middle" class="label">8</text>
<text x="606.7" y="372" text-anchor="middle" class="label">16</text>
<text x="737.3" y="372" text-anchor="middle" class="label">32</text>
<text x="868.0" y="372" text-anchor="middle" class="label">64</text>
<path d="M 84.0,313.7 L 214.7,130.8 L 345.3,95.9 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#9aa7b6" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="84.0" cy="313.7" r="3.5" fill="#9aa7b6" />
<circle cx="214.7" cy="130.8" r="3.5" fill="#9aa7b6" />
<circle cx="345.3" cy="95.9" r="3.5" fill="#9aa7b6" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#9aa7b6" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#9aa7b6" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#9aa7b6" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#9aa7b6" />
<path d="M 84.0,149.9 L 214.7,94.2 L 345.3,90.7 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#64748b" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="84.0" cy="149.9" r="3.5" fill="#64748b" />
<circle cx="214.7" cy="94.2" r="3.5" fill="#64748b" />
<circle cx="345.3" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#64748b" />
<path d="M 84.0,280.6 L 214.7,127.3 L 345.3,94.2 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#635bff" stroke-width="2.25" />
<circle cx="84.0" cy="280.6" r="3.5" fill="#635bff" />
<circle cx="214.7" cy="127.3" r="3.5" fill="#635bff" />
<circle cx="345.3" cy="94.2" r="3.5" fill="#635bff" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#635bff" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#635bff" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#635bff" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#635bff" />
<path d="M 84.0,136.0 L 214.7,94.2 L 345.3,90.7 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#0f766e" stroke-width="2.25" />
<circle cx="84.0" cy="136.0" r="3.5" fill="#0f766e" />
<circle cx="214.7" cy="94.2" r="3.5" fill="#0f766e" />
<circle cx="345.3" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#0f766e" />
<text x="22" y="217.0" transform="rotate(-90, 22, 217.0)" class="axis">recall@1@k</text>
<text x="476.0" y="400" text-anchor="middle" class="axis">k</text>
<line x1="84" y1="432" x2="108" y2="432" stroke="#635bff" stroke-width="2.25" />
<circle cx="96" cy="432" r="3.5" fill="#635bff" />
<text x="116" y="435" class="legend">TQ 2-bit</text>
<line x1="224" y1="432" x2="248" y2="432" stroke="#0f766e" stroke-width="2.25" />
<circle cx="236" cy="432" r="3.5" fill="#0f766e" />
<text x="256" y="435" class="legend">TQ 4-bit</text>
<line x1="364" y1="432" x2="388" y2="432" stroke="#9aa7b6" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="376" cy="432" r="3.5" fill="#9aa7b6" />
<text x="396" y="435" class="legend">FAISS 2-bit</text>
<line x1="504" y1="432" x2="528" y2="432" stroke="#64748b" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="516" cy="432" r="3.5" fill="#64748b" />
<text x="536" y="435" class="legend">FAISS 4-bit</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 6.0 KiB

docs/recall_d3072.svg
+85
View File
@@ -0,0 +1,85 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Recall — d=3072">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Recall — d=3072</text>
<text x="84" y="52" class="subtitle">100K vectors, k=64 search. recall@1@k measures how often the true top-1 result appears in the top-k returned.</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0.85</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">0.88</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">0.91</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">0.94</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">0.97</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">1.00</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<text x="84" y="68" class="panel">d=3072</text>
<text x="84.0" y="372" text-anchor="middle" class="label">1</text>
<text x="214.7" y="372" text-anchor="middle" class="label">2</text>
<text x="345.3" y="372" text-anchor="middle" class="label">4</text>
<text x="476.0" y="372" text-anchor="middle" class="label">8</text>
<text x="606.7" y="372" text-anchor="middle" class="label">16</text>
<text x="737.3" y="372" text-anchor="middle" class="label">32</text>
<text x="868.0" y="372" text-anchor="middle" class="label">64</text>
<path d="M 84.0,244.0 L 214.7,115.1 L 345.3,90.7 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#9aa7b6" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="84.0" cy="244.0" r="3.5" fill="#9aa7b6" />
<circle cx="214.7" cy="115.1" r="3.5" fill="#9aa7b6" />
<circle cx="345.3" cy="90.7" r="3.5" fill="#9aa7b6" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#9aa7b6" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#9aa7b6" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#9aa7b6" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#9aa7b6" />
<path d="M 84.0,139.5 L 214.7,94.2 L 345.3,90.7 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#64748b" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="84.0" cy="139.5" r="3.5" fill="#64748b" />
<circle cx="214.7" cy="94.2" r="3.5" fill="#64748b" />
<circle cx="345.3" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#64748b" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#64748b" />
<path d="M 84.0,214.4 L 214.7,111.6 L 345.3,92.5 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#635bff" stroke-width="2.25" />
<circle cx="84.0" cy="214.4" r="3.5" fill="#635bff" />
<circle cx="214.7" cy="111.6" r="3.5" fill="#635bff" />
<circle cx="345.3" cy="92.5" r="3.5" fill="#635bff" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#635bff" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#635bff" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#635bff" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#635bff" />
<path d="M 84.0,136.0 L 214.7,92.5 L 345.3,90.7 L 476.0,90.7 L 606.7,90.7 L 737.3,90.7 L 868.0,90.7" fill="none" stroke="#0f766e" stroke-width="2.25" />
<circle cx="84.0" cy="136.0" r="3.5" fill="#0f766e" />
<circle cx="214.7" cy="92.5" r="3.5" fill="#0f766e" />
<circle cx="345.3" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="476.0" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="606.7" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="737.3" cy="90.7" r="3.5" fill="#0f766e" />
<circle cx="868.0" cy="90.7" r="3.5" fill="#0f766e" />
<text x="22" y="217.0" transform="rotate(-90, 22, 217.0)" class="axis">recall@1@k</text>
<text x="476.0" y="400" text-anchor="middle" class="axis">k</text>
<line x1="84" y1="432" x2="108" y2="432" stroke="#635bff" stroke-width="2.25" />
<circle cx="96" cy="432" r="3.5" fill="#635bff" />
<text x="116" y="435" class="legend">TQ 2-bit</text>
<line x1="224" y1="432" x2="248" y2="432" stroke="#0f766e" stroke-width="2.25" />
<circle cx="236" cy="432" r="3.5" fill="#0f766e" />
<text x="256" y="435" class="legend">TQ 4-bit</text>
<line x1="364" y1="432" x2="388" y2="432" stroke="#9aa7b6" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="376" cy="432" r="3.5" fill="#9aa7b6" />
<text x="396" y="435" class="legend">FAISS 2-bit</text>
<line x1="504" y1="432" x2="528" y2="432" stroke="#64748b" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="516" cy="432" r="3.5" fill="#64748b" />
<text x="536" y="435" class="legend">FAISS 4-bit</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 6.0 KiB

docs/recall_glove.svg
+85
View File
@@ -0,0 +1,85 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Recall — GloVe d=200">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Recall — GloVe d=200</text>
<text x="84" y="52" class="subtitle">100K vectors, k=64 search. recall@1@k measures how often the true top-1 result appears in the top-k returned.</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0.40</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">0.52</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">0.64</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">0.76</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">0.88</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">1.00</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<text x="84" y="68" class="panel">GloVe d=200</text>
<text x="84.0" y="372" text-anchor="middle" class="label">1</text>
<text x="214.7" y="372" text-anchor="middle" class="label">2</text>
<text x="345.3" y="372" text-anchor="middle" class="label">4</text>
<text x="476.0" y="372" text-anchor="middle" class="label">8</text>
<text x="606.7" y="372" text-anchor="middle" class="label">16</text>
<text x="737.3" y="372" text-anchor="middle" class="label">32</text>
<text x="868.0" y="372" text-anchor="middle" class="label">64</text>
<path d="M 84.0,278.7 L 214.7,209.7 L 345.3,153.6 L 476.0,117.6 L 606.7,97.6 L 737.3,88.3 L 868.0,85.1" fill="none" stroke="#9aa7b6" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="84.0" cy="278.7" r="3.5" fill="#9aa7b6" />
<circle cx="214.7" cy="209.7" r="3.5" fill="#9aa7b6" />
<circle cx="345.3" cy="153.6" r="3.5" fill="#9aa7b6" />
<circle cx="476.0" cy="117.6" r="3.5" fill="#9aa7b6" />
<circle cx="606.7" cy="97.6" r="3.5" fill="#9aa7b6" />
<circle cx="737.3" cy="88.3" r="3.5" fill="#9aa7b6" />
<circle cx="868.0" cy="85.1" r="3.5" fill="#9aa7b6" />
<path d="M 84.0,155.2 L 214.7,105.9 L 345.3,88.1 L 476.0,84.9 L 606.7,84.2 L 737.3,84.2 L 868.0,84.2" fill="none" stroke="#64748b" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="84.0" cy="155.2" r="3.5" fill="#64748b" />
<circle cx="214.7" cy="105.9" r="3.5" fill="#64748b" />
<circle cx="345.3" cy="88.1" r="3.5" fill="#64748b" />
<circle cx="476.0" cy="84.9" r="3.5" fill="#64748b" />
<circle cx="606.7" cy="84.2" r="3.5" fill="#64748b" />
<circle cx="737.3" cy="84.2" r="3.5" fill="#64748b" />
<circle cx="868.0" cy="84.2" r="3.5" fill="#64748b" />
<path d="M 84.0,278.9 L 214.7,208.8 L 345.3,153.7 L 476.0,118.3 L 606.7,98.7 L 737.3,89.1 L 868.0,85.7" fill="none" stroke="#635bff" stroke-width="2.25" />
<circle cx="84.0" cy="278.9" r="3.5" fill="#635bff" />
<circle cx="214.7" cy="208.8" r="3.5" fill="#635bff" />
<circle cx="345.3" cy="153.7" r="3.5" fill="#635bff" />
<circle cx="476.0" cy="118.3" r="3.5" fill="#635bff" />
<circle cx="606.7" cy="98.7" r="3.5" fill="#635bff" />
<circle cx="737.3" cy="89.1" r="3.5" fill="#635bff" />
<circle cx="868.0" cy="85.7" r="3.5" fill="#635bff" />
<path d="M 84.0,151.3 L 214.7,102.8 L 345.3,87.1 L 476.0,84.3 L 606.7,84.2 L 737.3,84.2 L 868.0,84.2" fill="none" stroke="#0f766e" stroke-width="2.25" />
<circle cx="84.0" cy="151.3" r="3.5" fill="#0f766e" />
<circle cx="214.7" cy="102.8" r="3.5" fill="#0f766e" />
<circle cx="345.3" cy="87.1" r="3.5" fill="#0f766e" />
<circle cx="476.0" cy="84.3" r="3.5" fill="#0f766e" />
<circle cx="606.7" cy="84.2" r="3.5" fill="#0f766e" />
<circle cx="737.3" cy="84.2" r="3.5" fill="#0f766e" />
<circle cx="868.0" cy="84.2" r="3.5" fill="#0f766e" />
<text x="22" y="217.0" transform="rotate(-90, 22, 217.0)" class="axis">recall@1@k</text>
<text x="476.0" y="400" text-anchor="middle" class="axis">k</text>
<line x1="84" y1="432" x2="108" y2="432" stroke="#635bff" stroke-width="2.25" />
<circle cx="96" cy="432" r="3.5" fill="#635bff" />
<text x="116" y="435" class="legend">TQ 2-bit</text>
<line x1="224" y1="432" x2="248" y2="432" stroke="#0f766e" stroke-width="2.25" />
<circle cx="236" cy="432" r="3.5" fill="#0f766e" />
<text x="256" y="435" class="legend">TQ 4-bit</text>
<line x1="364" y1="432" x2="388" y2="432" stroke="#9aa7b6" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="376" cy="432" r="3.5" fill="#9aa7b6" />
<text x="396" y="435" class="legend">FAISS 2-bit</text>
<line x1="504" y1="432" x2="528" y2="432" stroke="#64748b" stroke-width="2.25" stroke-dasharray="6 4" />
<circle cx="516" cy="432" r="3.5" fill="#64748b" />
<text x="536" y="435" class="legend">FAISS 4-bit</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 6.0 KiB

docs/x86_speed_mt.svg
+61
View File
@@ -0,0 +1,61 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Search Latency — x86 (Intel Sapphire Rapids, 8 vCPUs) — Multi-threaded">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Search Latency — x86 (Intel Sapphire Rapids, 8 vCPUs) — Multi-threaded</text>
<text x="84" y="52" class="subtitle">100K vectors, 1K queries, k=64, median of 5 runs</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0.00</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">0.30</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">0.60</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">0.90</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">1.20</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">1.50</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<text x="84" y="68" class="panel">Multi-threaded</text>
<rect x="135.0" y="297.3" width="44" height="54.7" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="185.0" y="298.9" width="44" height="53.1" rx="6" fill="#9aa7b6" />
<text x="157.0" y="291.3" text-anchor="middle" class="value-accent">0.304</text>
<text x="207.0" y="292.9" text-anchor="middle" class="value">0.295</text>
<text x="182.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="182.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="331.0" y="248.3" width="44" height="103.7" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="381.0" y="245.8" width="44" height="106.2" rx="6" fill="#9aa7b6" />
<text x="353.0" y="242.3" text-anchor="middle" class="value-accent">0.576</text>
<text x="403.0" y="239.8" text-anchor="middle" class="value">0.590</text>
<text x="378.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="378.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<rect x="527.0" y="239.3" width="44" height="112.7" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="577.0" y="245.8" width="44" height="106.2" rx="6" fill="#9aa7b6" />
<text x="549.0" y="233.3" text-anchor="middle" class="value-accent">0.626</text>
<text x="599.0" y="239.8" text-anchor="middle" class="value">0.590</text>
<text x="574.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="574.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="723.0" y="140.1" width="44" height="211.9" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="773.0" y="140.1" width="44" height="211.9" rx="6" fill="#9aa7b6" />
<text x="745.0" y="134.1" text-anchor="middle" class="value-accent">1.177</text>
<text x="795.0" y="134.1" text-anchor="middle" class="value">1.177</text>
<text x="770.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="770.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<text x="26" y="217.0" transform="rotate(-90, 26, 217.0)" class="axis">ms / query</text>
<rect x="84" y="424" width="14" height="14" rx="3" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="106" y="435" class="legend" style="fill: #4338ca;">TurboQuant</text>
<rect x="224" y="424" width="14" height="14" rx="3" fill="#9aa7b6" />
<text x="246" y="435" class="legend">FAISS</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 4.8 KiB

docs/x86_speed_st.svg
+61
View File
@@ -0,0 +1,61 @@
<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" width="900" height="460" viewBox="0 0 900 460" role="img" aria-label="Search Latency — x86 (Intel Sapphire Rapids, 8 vCPUs) — Single-threaded">
<style>
.title { font: 700 20px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.subtitle { font: 400 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.panel { font: 700 14px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.label { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.secondary { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #475569; }
.tick { font: 400 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #64748b; }
.value { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
.value-accent { font: 700 11px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #4338ca; }
.axis { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #334155; }
.legend { font: 600 12px -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; fill: #0f172a; }
</style>
<rect width="100%" height="100%" fill="#ffffff" />
<text x="84" y="32" class="title">Search Latency — x86 (Intel Sapphire Rapids, 8 vCPUs) — Single-threaded</text>
<text x="84" y="52" class="subtitle">100K vectors, 1K queries, k=64, median of 5 runs</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="356.0" text-anchor="end" class="tick">0.0</text>
<line x1="84" y1="298.0" x2="868" y2="298.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="302.0" text-anchor="end" class="tick">2.0</text>
<line x1="84" y1="244.0" x2="868" y2="244.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="248.0" text-anchor="end" class="tick">4.0</text>
<line x1="84" y1="190.0" x2="868" y2="190.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="194.0" text-anchor="end" class="tick">6.0</text>
<line x1="84" y1="136.0" x2="868" y2="136.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="140.0" text-anchor="end" class="tick">8.0</text>
<line x1="84" y1="82.0" x2="868" y2="82.0" stroke="#e5e7eb" stroke-width="1" />
<text x="74" y="86.0" text-anchor="end" class="tick">10.0</text>
<line x1="84" y1="352.0" x2="868" y2="352.0" stroke="#94a3b8" stroke-width="1.5" />
<text x="84" y="68" class="panel">Single-threaded</text>
<rect x="135.0" y="317.7" width="44" height="34.3" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="185.0" y="320.4" width="44" height="31.6" rx="6" fill="#9aa7b6" />
<text x="157.0" y="311.7" text-anchor="middle" class="value-accent">1.27</text>
<text x="207.0" y="314.4" text-anchor="middle" class="value">1.17</text>
<text x="182.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="182.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="331.0" y="286.1" width="44" height="65.9" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="381.0" y="282.9" width="44" height="69.1" rx="6" fill="#9aa7b6" />
<text x="353.0" y="280.1" text-anchor="middle" class="value-accent">2.44</text>
<text x="403.0" y="276.9" text-anchor="middle" class="value">2.56</text>
<text x="378.0" y="374" text-anchor="middle" class="label">d=1536</text>
<text x="378.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<rect x="527.0" y="280.3" width="44" height="71.7" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="577.0" y="282.3" width="44" height="69.7" rx="6" fill="#9aa7b6" />
<text x="549.0" y="274.3" text-anchor="middle" class="value-accent">2.66</text>
<text x="599.0" y="276.3" text-anchor="middle" class="value">2.58</text>
<text x="574.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="574.0" y="389" text-anchor="middle" class="secondary">2-bit</text>
<rect x="723.0" y="207.8" width="44" height="144.2" rx="6" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<rect x="773.0" y="204.2" width="44" height="147.8" rx="6" fill="#9aa7b6" />
<text x="745.0" y="201.8" text-anchor="middle" class="value-accent">5.34</text>
<text x="795.0" y="198.2" text-anchor="middle" class="value">5.47</text>
<text x="770.0" y="374" text-anchor="middle" class="label">d=3072</text>
<text x="770.0" y="389" text-anchor="middle" class="secondary">4-bit</text>
<text x="26" y="217.0" transform="rotate(-90, 26, 217.0)" class="axis">ms / query</text>
<rect x="84" y="424" width="14" height="14" rx="3" fill="#635bff" stroke="#4338ca" stroke-width="1.5" />
<text x="106" y="435" class="legend" style="fill: #4338ca;">TurboQuant</text>
<rect x="224" y="424" width="14" height="14" rx="3" fill="#9aa7b6" />
<text x="246" y="435" class="legend">FAISS</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 4.8 KiB

examples/downstream-smoke/Cargo.lock
Generated
+1135
View File
File diff suppressed because it is too large Load Diff
examples/downstream-smoke/Cargo.toml
+15
View File
@@ -0,0 +1,15 @@
[package]
name = "downstream-smoke"
version = "0.0.0"
edition = "2021"
publish = false
# Standalone crate (NOT a workspace member) that depends on turbovec via
# a path dep. This exercises the same code path a downstream cargo user
# hits when they `cargo add turbovec` and write their first program:
# the binary must link `cblas_sgemm` without any `extern crate blas_src;`
# ceremony or extra Cargo.toml entries.
#
# Run from this directory: `cargo run --release`.
[dependencies]
turbovec = { path = "../../turbovec" }
examples/downstream-smoke/src/main.rs
+66
View File
@@ -0,0 +1,66 @@
//! Smoke test for the downstream `cargo add turbovec` experience.
//!
//! Exercises the public API end-to-end (construct, add, prepare, search,
//! write, load) with no BLAS-related setup. If this binary links and
//! runs, the link-directive propagation from turbovec's build.rs is
//! working and downstream users won't hit the `cblas_sgemm` error.
use turbovec::TurboQuantIndex;
const DIM: usize = 64;
const N_DB: usize = 256;
const N_QUERIES: usize = 4;
const K: usize = 5;
fn unit_vectors(n: usize, dim: usize, seed: u64) -> Vec<f32> {
let mut state = seed | 1;
let mut next = || {
state ^= state << 13;
state ^= state >> 7;
state ^= state << 17;
(state as f32 / u64::MAX as f32) * 2.0 - 1.0
};
let mut out = vec![0.0f32; n * dim];
for v in out.chunks_mut(dim) {
for x in v.iter_mut() {
*x = next();
}
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt().max(1e-12);
for x in v.iter_mut() {
*x /= norm;
}
}
out
}
fn main() {
let db = unit_vectors(N_DB, DIM, 1);
let queries = unit_vectors(N_QUERIES, DIM, 2);
let mut index = TurboQuantIndex::new(DIM, 4).expect("construct");
index.add(&db);
index.prepare();
let results = index.search(&queries, K);
assert_eq!(results.nq, N_QUERIES);
assert_eq!(results.k, K);
for q in 0..N_QUERIES {
let idxs = results.indices_for_query(q);
assert_eq!(idxs.len(), K);
for &i in idxs {
assert!((0..N_DB as i64).contains(&i), "out-of-range index {i}");
}
}
let tmp = std::env::temp_dir().join("turbovec_downstream_smoke.tv");
index.write(&tmp).expect("write");
let loaded = TurboQuantIndex::load(&tmp).expect("load");
assert_eq!(loaded.len(), N_DB);
assert_eq!(loaded.dim(), DIM);
let _ = std::fs::remove_file(&tmp);
println!(
"downstream-smoke: OK ({} db vectors, {} queries, top-{} search, write/load round-trip)",
N_DB, N_QUERIES, K
);
}
turbovec-python/Cargo.toml
+17
View File
@@ -0,0 +1,17 @@
[package]
name = "turbovec-python"
version = "0.8.0"
edition = "2021"
rust-version = "1.70"
[lib]
name = "_turbovec"
crate-type = ["cdylib"]
[dependencies]
turbovec-core = { package = "turbovec", path = "../turbovec" }
pyo3 = { version = "0.27.0", features = ["extension-module", "abi3-py39"] }
numpy = "0.27.0"
[build-dependencies]
pyo3-build-config = "0.27.0"
turbovec-python/README.md
Symlink
+1
View File
@@ -0,0 +1 @@
../README.md
turbovec-python/build.rs
+10
View File
@@ -0,0 +1,10 @@
fn main() {
// Emit the platform-correct linker arguments for a Python extension
// module. On macOS this passes `-undefined dynamic_lookup` so symbols
// from the Python interpreter (e.g. `Py_True`) resolve at load time
// instead of failing the link step. Without it, a plain `cargo build`
// on macOS fails with "symbol(s) not found for architecture arm64"
// (issue #92). Building via maturin already injects these args; this
// makes a bare `cargo build` work too.
pyo3_build_config::add_extension_module_link_args();
}
turbovec-python/pyproject.toml
+62
View File
@@ -0,0 +1,62 @@
[build-system]
requires = ["maturin>=1.12,<2.0"]
build-backend = "maturin"
[project]
name = "turbovec"
version = "0.8.0"
description = "Fast vector quantization with 2-4 bit compression and SIMD search"
readme = "README.md"
requires-python = ">=3.9"
license = { text = "MIT" }
authors = [
{ name = "Ryan Codrai" },
]
keywords = [
"vector-search",
"quantization",
"nearest-neighbor",
"ann",
"simd",
"rag",
"embeddings",
]
classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Developers",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: MIT License",
"Operating System :: POSIX :: Linux",
"Operating System :: MacOS",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: 3.14",
"Programming Language :: Rust",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Database",
]
dependencies = [
"numpy>=1.20",
]
[project.urls]
Homepage = "https://github.com/RyanCodrai/turbovec"
Repository = "https://github.com/RyanCodrai/turbovec"
Issues = "https://github.com/RyanCodrai/turbovec/issues"
[project.optional-dependencies]
langchain = ["langchain-core>=0.3"]
llama-index = ["llama-index-core>=0.11"]
haystack = ["haystack-ai>=2.0"]
agno = ["agno>=2.0"]
[tool.maturin]
manifest-path = "Cargo.toml"
features = ["pyo3/extension-module"]
python-source = "python"
module-name = "turbovec._turbovec"
turbovec-python/python/turbovec/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from ._turbovec import IdMapIndex, TurboQuantIndex
__all__ = ["IdMapIndex", "TurboQuantIndex"]
turbovec-python/python/turbovec/_dedup.py
+71
View File
@@ -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"]
turbovec-python/python/turbovec/_persist.py
+55
View File
@@ -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"]
turbovec-python/python/turbovec/agno.py
+788
View File
@@ -0,0 +1,788 @@
"""Agno VectorDb backed by turbovec's quantized index.
Install with: ``pip install turbovec[agno]``.
Implements Agno's ``VectorDb`` interface and matches the public surface
of ``agno.vectordb.lancedb.LanceDb`` (the closest in-tree single-machine
backend) so this can be swapped in wherever ``LanceDb`` is used.
"""
from __future__ import annotations
import json
from hashlib import md5
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Set, Union
import numpy as np
from ._turbovec import IdMapIndex
try:
from agno.knowledge.document import Document
from agno.knowledge.embedder import Embedder
from agno.knowledge.reranker.base import Reranker
from agno.vectordb.base import VectorDb
from agno.vectordb.distance import Distance
from agno.vectordb.search import SearchType
except ImportError as exc:
raise ImportError(
"agno is required to use turbovec.agno. "
"Install with: pip install turbovec[agno]"
) from exc
_INDEX_FILENAME = "index.tvim"
_STORE_FILENAME = "docstore.json"
# Bump when docstore.json shape changes; loader refuses unknown versions.
_DOCSTORE_SCHEMA_VERSION = 1
class TurboQuantVectorDb(VectorDb):
"""Agno VectorDb backed by a :class:`IdMapIndex`.
Vectors are quantized to 2-4 bits per dimension. The public surface
mirrors ``agno.vectordb.lancedb.LanceDb`` so this is a drop-in
replacement wherever a single-machine LanceDb is used. Search-time
filtering is resolved to an allowlist *before* scoring (kernel-level)
rather than via post-filtering, so selective filters return up to
``limit`` results from the filtered set instead of fewer.
Example::
from agno.knowledge.embedder.openai import OpenAIEmbedder
from turbovec.agno import TurboQuantVectorDb
vector_db = TurboQuantVectorDb(embedder=OpenAIEmbedder())
vector_db.create()
# ... use as a normal Agno VectorDb ...
"""
def __init__(
self,
*,
id: Optional[str] = None,
name: Optional[str] = None,
description: Optional[str] = None,
similarity_threshold: Optional[float] = None,
embedder: Optional[Embedder] = None,
bit_width: int = 4,
search_type: SearchType = SearchType.vector,
distance: Distance = Distance.cosine,
reranker: Optional[Reranker] = None,
path: Optional[str] = None,
) -> None:
"""
:param embedder: Required. Agno embedder used to encode documents
and queries. ``embedder.dimensions`` must be set — it's the
sole source of truth for the underlying quantized index's
dimensionality.
:param bit_width: Quantization width (2 or 4).
:param search_type: Only :class:`SearchType.vector` is supported;
other values raise :class:`ValueError`. (Keyword/hybrid search
would require an external BM25/lexical index.)
:param distance: Only :class:`Distance.cosine` is supported.
turbovec stores unit-normalized vectors, so the kernel's raw
score is cosine similarity directly.
:param reranker: Optional Agno reranker applied to the result set
after vector retrieval.
:param path: Optional directory for save/load persistence. When
given to the constructor, :meth:`create` loads existing data
from this path if present.
"""
super().__init__(
id=id,
name=name,
description=description,
similarity_threshold=similarity_threshold,
)
if embedder is None:
raise ValueError(
"`embedder` is required; turbovec needs the embedder's "
"`dimensions` to size the underlying index."
)
if embedder.dimensions is None:
raise ValueError("Embedder.dimensions must be set.")
if bit_width not in (2, 4):
raise ValueError(f"bit_width must be 2 or 4, got {bit_width}")
if search_type != SearchType.vector:
raise ValueError(
f"TurboQuantVectorDb only supports search_type=SearchType.vector; "
f"got {search_type}. Use LanceDb / Chroma / etc. for keyword "
f"or hybrid search."
)
if distance != Distance.cosine:
raise ValueError(
f"TurboQuantVectorDb only supports distance=Distance.cosine; "
f"got {distance}. turbovec stores unit-normalized vectors."
)
self.embedder: Embedder = embedder
self.dimensions: int = embedder.dimensions
self.bit_width = bit_width
self.search_type = search_type
self.distance = distance
self.reranker = reranker
self.path: Optional[str] = path
# Lazy: the underlying IdMapIndex is created by `create()`, not
# in __init__. This matches LanceDb's `exists()` contract: a
# 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 -> 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
self._next_u64: int = 0
# Auxiliary indexes for O(1) protocol queries
self._content_hashes: Set[str] = set()
self._name_to_ids: Dict[str, Set[str]] = {}
# ---- handle allocation ------------------------------------------------
def _issue_handle(self) -> int:
self._next_u64 += 1
return self._next_u64
# ---- VectorDb protocol: lifecycle ------------------------------------
def create(self) -> None:
"""Create the underlying index if it doesn't already exist.
Idempotent — calling on an already-created store is a no-op.
If ``path`` was set on the constructor and a previous save exists
under it, ``create()`` loads that save; otherwise it instantiates
a fresh empty index sized to ``embedder.dimensions``.
"""
if self._index is not None:
return
# Try loading from path first if one was set; fall through to a
# fresh index if the path doesn't contain a previous save.
if self.path is not None and Path(self.path).is_dir():
try:
self._load_from(Path(self.path))
return
except FileNotFoundError:
pass
self._index = IdMapIndex(self.dimensions, self.bit_width)
async def async_create(self) -> None:
self.create()
def drop(self) -> None:
"""Drop the underlying index. After this call ``exists()`` returns
``False`` until ``create()`` is called again — matches LanceDb's
contract where ``drop()`` removes the table entirely."""
self._index = None
self._str_to_u64.clear()
self._u64_to_doc.clear()
self._next_u64 = 0
self._content_hashes.clear()
self._name_to_ids.clear()
async def async_drop(self) -> None:
self.drop()
def exists(self) -> bool:
"""True iff the underlying index has been created via ``create()``
and not subsequently dropped. Matches LanceDb's
"table-exists-in-connection" semantic; *does not* mean
"has any documents" — call ``get_count()`` for that."""
return self._index is not None
async def async_exists(self) -> bool:
return self.exists()
def delete(self) -> bool:
"""Returns ``False``. The Agno protocol declares this abstract
method but LanceDb (the drop-in reference) unconditionally
returns False — actual destruction goes through ``drop()``."""
return False
def optimize(self) -> None:
"""No-op. The underlying quantized index doesn't have a
post-write optimization step. Matches LanceDb's ``optimize()``
which is also a no-op."""
return None
def get_count(self) -> int:
"""Number of documents currently stored."""
if self._index is None:
return 0
return len(self._index)
async def async_get_count(self) -> int:
return self.get_count()
# ---- VectorDb protocol: existence checks ------------------------------
def name_exists(self, name: str) -> bool:
if self._index is None:
return False
return name in self._name_to_ids
async def async_name_exists(self, name: str) -> bool:
# LanceDb raises NotImplementedError here; we have a trivial sync
# backing call, so we return the real answer. Intentional deviation.
return self.name_exists(name)
def id_exists(self, id: str) -> bool:
if self._index is None:
return False
return id in self._str_to_u64
def content_hash_exists(self, content_hash: str) -> bool:
if self._index is None:
return False
return content_hash in self._content_hashes
# ---- VectorDb protocol: insert / upsert -------------------------------
@staticmethod
def _derive_doc_id(doc: Document, content_hash: str, cleaned_content: str) -> str:
"""Match LanceDb's id-derivation contract so the same doc with the
same content_hash produces the same stable doc_id across stores."""
base_id = doc.id or md5(cleaned_content.encode()).hexdigest()
return md5(f"{base_id}_{content_hash}".encode()).hexdigest()
def _embed_missing(self, documents: List[Document]) -> None:
"""Populate embeddings on any documents that don't have one. Uses
the embedder's batch path when available."""
to_embed = [
doc
for doc in documents
if doc.embedding is None
or (isinstance(doc.embedding, list) and len(doc.embedding) == 0)
]
if not to_embed:
return
if (
getattr(self.embedder, "enable_batch", False)
and hasattr(self.embedder, "get_embeddings_batch_and_usage")
):
contents = [doc.content for doc in to_embed]
embeddings, usages = self.embedder.get_embeddings_batch_and_usage(contents)
for j, doc in enumerate(to_embed):
if j < len(embeddings):
doc.embedding = embeddings[j]
doc.usage = usages[j] if j < len(usages) else None
else:
for doc in to_embed:
doc.embed(embedder=self.embedder)
async def _embed_missing_async(self, documents: List[Document]) -> None:
to_embed = [
doc
for doc in documents
if doc.embedding is None
or (isinstance(doc.embedding, list) and len(doc.embedding) == 0)
]
if not to_embed:
return
if (
getattr(self.embedder, "enable_batch", False)
and hasattr(self.embedder, "async_get_embeddings_batch_and_usage")
):
contents = [doc.content for doc in to_embed]
embeddings, usages = await self.embedder.async_get_embeddings_batch_and_usage(
contents
)
for j, doc in enumerate(to_embed):
if j < len(embeddings):
doc.embedding = embeddings[j]
doc.usage = usages[j] if j < len(usages) else None
else:
# Embedder has no async batch path — fall back to sync.
self._embed_missing(to_embed)
def insert(
self,
content_hash: str,
documents: List[Document],
filters: Optional[Dict[str, Any]] = None,
) -> None:
if not documents:
return
if self._index is None:
# Match LanceDb's "table not initialized" handling: do not
# silently auto-create. Callers must invoke create() first.
raise RuntimeError(
"TurboQuantVectorDb not initialized — call create() before insert()."
)
# Merge `filters` into each document's metadata (matches LanceDb).
if filters:
for doc in documents:
meta = dict(doc.meta_data) if doc.meta_data else {}
meta.update(filters)
doc.meta_data = meta
self._embed_missing(documents)
# Raise on any document that still lacks an embedding rather than
# silently dropping — silent drops mask data-pipeline bugs.
missing = [doc for doc in documents if not doc.embedding]
if missing:
ids = [doc.id or "<no id>" for doc in missing]
raise ValueError(
f"failed to embed {len(missing)} document(s): {ids}"
)
# Batch the entire `documents` list into a single add_with_ids call.
# Per-document inserts would invalidate the SIMD-blocked cache
# between every doc.
vectors = np.asarray([doc.embedding for doc in documents], dtype=np.float32)
if vectors.ndim != 2:
raise ValueError(
f"expected 2D embedding batch, got {vectors.ndim}D"
)
if vectors.shape[1] != self.dimensions:
raise ValueError(
f"embedding dim {vectors.shape[1]} does not match "
f"index dim {self.dimensions}"
)
if not vectors.flags["C_CONTIGUOUS"]:
vectors = np.ascontiguousarray(vectors)
handles = np.array(
[self._issue_handle() for _ in documents], dtype=np.uint64
)
self._index.add_with_ids(vectors, handles)
for doc, handle in zip(documents, handles):
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.setdefault(doc_id, set()).add(h)
self._u64_to_doc[h] = {
"id": doc_id,
"name": doc.name,
"content": cleaned,
"meta_data": dict(doc.meta_data) if doc.meta_data else {},
"usage": doc.usage,
"content_id": doc.content_id,
"content_hash": content_hash,
}
self._content_hashes.add(content_hash)
if doc.name:
self._name_to_ids.setdefault(doc.name, set()).add(doc_id)
async def async_insert(
self,
content_hash: str,
documents: List[Document],
filters: Optional[Dict[str, Any]] = None,
) -> None:
if not documents:
return
await self._embed_missing_async(documents)
# Now every doc should have an embedding; insert delegates to sync.
self.insert(content_hash, documents, filters)
def upsert_available(self) -> bool:
return True
def upsert(
self,
content_hash: str,
documents: List[Document],
filters: Optional[Dict[str, Any]] = None,
) -> None:
# Match LanceDb's semantic: replace all documents previously
# stored under this content_hash with the incoming batch. Not
# "replace by derived doc_id" — that's a different contract.
#
# Capture the existing generation's handles, run the insert, and
# only then drop the old vectors — so a failed insert (dim
# mismatch, non-finite embeddings) never destroys the data being
# replaced (issue #89). We delete by captured handle rather than
# re-querying by content_hash, because insert() re-derives ids
# under the SAME content_hash and would otherwise clobber the
# just-inserted rows.
old_handles = self._handles_for_content_hash(content_hash)
self.insert(content_hash, documents, filters)
for handle in old_handles:
self._remove_handle(handle)
async def async_upsert(
self,
content_hash: str,
documents: List[Document],
filters: Optional[Dict[str, Any]] = None,
) -> None:
old_handles = self._handles_for_content_hash(content_hash)
await self.async_insert(content_hash, documents, filters)
for handle in old_handles:
self._remove_handle(handle)
def _handles_for_content_hash(self, content_hash: str) -> List[int]:
"""Internal handles of every document currently stored under this
content_hash. Used by upsert to defer removal of the previous
generation until the replacement add has succeeded (issue #89)."""
return [
handle
for handle, data in self._u64_to_doc.items()
if data.get("content_hash") == content_hash
]
def _remove_handle(self, handle: int) -> None:
"""Remove a single vector by its internal handle, leaving other
handles intact — including ones that share this document's derived
id (two distinct documents can map to the same doc_id, matching
LanceDb). Cleans the id, name, and content_hash side-indexes only
where no surviving handle still needs them."""
if self._index is None:
return
data = self._u64_to_doc.pop(handle, None)
if data is None:
return
self._index.remove(handle)
doc_id = data.get("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 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]
# Drop the content_hash only if no surviving doc carries it.
ch = 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)
# ---- VectorDb protocol: search ----------------------------------------
def _resolve_filter_to_handles(
self, filters: Optional[Union[Dict[str, Any], List[Any]]]
) -> Optional[List[int]]:
"""Convert a dict filter into the list of internal u64 handles
whose document's ``meta_data`` matches every key/value pair (AND).
Returns ``None`` when no filter was supplied — caller should run
an unfiltered search. Returns ``[]`` to mean "no matches".
Matches LanceDb's dict-filter semantics (exact equality, AND of
keys). ``FilterExpr``-style list filters are not yet supported
in LanceDb itself, so we silently ignore them here too with a
debug log.
"""
if filters is None:
return None
if isinstance(filters, list):
# LanceDb logs a warning and ignores. Mirror that — the
# alternative is to error and break callers that pass an
# accidental list.
return None
if not isinstance(filters, dict) or not filters:
return None
items = list(filters.items())
return [
handle
for handle, data in self._u64_to_doc.items()
if all((data.get("meta_data") or {}).get(k) == v for k, v in items)
]
def _scaled_similarity(self, raw: float) -> float:
"""Map cosine similarity in ``[-1, 1]`` to ``[0, 1]``. Clamped to
absorb the small overshoot caused by quantization noise."""
return max(0.0, min(1.0, (raw + 1.0) / 2.0))
def _build_results(
self, scores: np.ndarray, handles: np.ndarray
) -> List[Document]:
results: List[Document] = []
threshold = self.similarity_threshold
for raw_score, handle in zip(scores[0], handles[0]):
doc_data = self._u64_to_doc.get(int(handle))
if doc_data is None:
continue
similarity = self._scaled_similarity(float(raw_score))
if threshold is not None and similarity < threshold:
continue
results.append(
Document(
id=doc_data["id"],
name=doc_data.get("name"),
content=doc_data.get("content", ""),
meta_data=dict(doc_data.get("meta_data") or {}),
usage=doc_data.get("usage"),
content_id=doc_data.get("content_id"),
# Match LanceDb._build_search_results: thread the
# store's embedder through so downstream code can call
# `doc.embed()` / `doc.async_embed()` on a retrieved
# hit without explicitly passing the embedder back in.
# Without this, `doc.embed()` raises
# "No embedder provided".
embedder=self.embedder,
)
)
return results
def search(
self,
query: str,
limit: int = 5,
filters: Optional[Union[Dict[str, Any], List[Any]]] = None,
) -> List[Document]:
# An empty query string usually indicates an upstream bug
# (uninitialised variable, failed prompt construction). LanceDb
# short-circuits this to [] rather than searching with a hash-
# derived embedding of "", which would return arbitrary garbage.
if not query:
return []
if self._index is None or len(self._index) == 0:
return []
query_embedding = self.embedder.get_embedding(query)
if query_embedding is None:
return []
qvec = np.asarray(query_embedding, dtype=np.float32)
if qvec.ndim == 1:
qvec = qvec[None, :]
if not qvec.flags["C_CONTIGUOUS"]:
qvec = np.ascontiguousarray(qvec)
allowed_handles = self._resolve_filter_to_handles(filters)
if allowed_handles is None:
# Unfiltered.
k = min(limit, len(self._index))
scores, handles = self._index.search(qvec, k)
else:
if not allowed_handles:
return []
allowlist = np.asarray(allowed_handles, dtype=np.uint64)
scores, handles = self._index.search(qvec, limit, allowlist=allowlist)
results = self._build_results(scores, handles)
if self.reranker is not None and results:
results = self.reranker.rerank(query=query, documents=results)
return results
async def async_search(
self,
query: str,
limit: int = 5,
filters: Optional[Union[Dict[str, Any], List[Any]]] = None,
) -> List[Document]:
if not query:
return []
if self._index is None or len(self._index) == 0:
return []
if hasattr(self.embedder, "async_get_embedding"):
query_embedding = await self.embedder.async_get_embedding(query)
else:
query_embedding = self.embedder.get_embedding(query)
if query_embedding is None:
return []
qvec = np.asarray(query_embedding, dtype=np.float32)
if qvec.ndim == 1:
qvec = qvec[None, :]
if not qvec.flags["C_CONTIGUOUS"]:
qvec = np.ascontiguousarray(qvec)
allowed_handles = self._resolve_filter_to_handles(filters)
if allowed_handles is None:
k = min(limit, len(self._index))
scores, handles = self._index.search(qvec, k)
else:
if not allowed_handles:
return []
allowlist = np.asarray(allowed_handles, dtype=np.uint64)
scores, handles = self._index.search(qvec, limit, allowlist=allowlist)
results = self._build_results(scores, handles)
if self.reranker is not None and results:
results = self.reranker.rerank(query=query, documents=results)
return results
def get_supported_search_types(self) -> List[SearchType]:
# Only vector. Keyword and hybrid would require an external BM25
# / lexical index that turbovec doesn't ship. Return shape
# mirrors LanceDb: a list of SearchType enum members (not their
# `.value` strings).
return [SearchType.vector]
# ---- VectorDb protocol: delete ----------------------------------------
def delete_by_id(self, id: str) -> bool:
if self._index is None:
return False
handles = self._str_to_u64.get(id)
if not handles:
return False
# 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
# 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())
# 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 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
# 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 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``
payload field of every document whose ``content_id`` matches.
Mirrors LanceDb's update_metadata semantic which writes to both
fields (used by callers that pass filter-style restrictions at
retrieval time)."""
if self._index is None:
return
for data in self._u64_to_doc.values():
if data.get("content_id") == content_id:
meta = dict(data.get("meta_data") or {})
meta.update(metadata)
data["meta_data"] = meta
filters = data.get("filters")
if isinstance(filters, dict):
filters = dict(filters)
filters.update(metadata)
data["filters"] = filters
else:
data["filters"] = dict(metadata)
# ---- Persistence (JSON side-car) --------------------------------------
def save(self, folder_path: Optional[str] = None) -> None:
"""Persist the quantized index plus a JSON side-car to disk. Pass
``folder_path`` to override the constructor's ``path=``.
Writes two files under ``folder_path``:
- ``index.tvim`` — the :class:`IdMapIndex` payload.
- ``docstore.json`` — JSON-encoded document text, metadata, and
id maps. Side-car carries a ``schema_version`` field; loaders
reject unknown versions rather than silently misinterpreting
bytes.
"""
path = folder_path if folder_path is not None else self.path
if path is None:
raise ValueError(
"No path to save to. Pass `folder_path=` here or set "
"`path=` on the constructor."
)
if self._index is None:
raise RuntimeError(
"TurboQuantVectorDb has no index to save — call create() first."
)
folder = Path(path)
folder.mkdir(parents=True, exist_ok=True)
self._index.write(str(folder / _INDEX_FILENAME))
payload = {
"schema_version": _DOCSTORE_SCHEMA_VERSION,
# Round-trip int handles via list-of-pairs (JSON keys must be
# strings, but our handles are ints).
"u64_to_doc": [[h, d] for h, d in self._u64_to_doc.items()],
"next_u64": self._next_u64,
"bit_width": self.bit_width,
"dimensions": self.dimensions,
}
with open(folder / _STORE_FILENAME, "w") as f:
json.dump(payload, f)
def _load_from(self, folder: Path) -> None:
side_car = folder / _STORE_FILENAME
index_file = folder / _INDEX_FILENAME
if not side_car.exists() or not index_file.exists():
raise FileNotFoundError(
f"missing one of {_STORE_FILENAME}/{_INDEX_FILENAME} under {folder}"
)
with open(side_car) as f:
state = json.load(f)
version = state.get("schema_version", 0)
if version != _DOCSTORE_SCHEMA_VERSION:
raise ValueError(
f"{_STORE_FILENAME} has schema_version {version}; this "
f"turbovec expects {_DOCSTORE_SCHEMA_VERSION}"
)
if state.get("dimensions") != self.dimensions:
raise ValueError(
f"persisted dimensions={state.get('dimensions')} does not "
f"match this store's embedder dimensions={self.dimensions}"
)
self._index = IdMapIndex.load(str(index_file))
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. 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():
ch = data.get("content_hash")
if ch:
self._content_hashes.add(ch)
name = data.get("name")
if name:
self._name_to_ids.setdefault(name, set()).add(data["id"])
__all__ = ["TurboQuantVectorDb"]
turbovec-python/python/turbovec/haystack.py
+747
View File
@@ -0,0 +1,747 @@
"""Haystack DocumentStore backed by turbovec's quantized index.
Install with: ``pip install turbovec[haystack]``.
Implements the Haystack 2.x ``DocumentStore`` protocol and mirrors most
of ``InMemoryDocumentStore``'s public surface (write/filter/delete,
``embedding_retrieval``, ``save_to_disk``/``load_from_disk``, pipeline
``to_dict``/``from_dict``). BM25 (sparse-text) retrieval is not
implemented — wire an ``InMemoryBM25Retriever`` against a separate
store if you need keyword search alongside vector search. The
quantized index discards full-precision embeddings after compression —
callers that rely on ``Document.embedding`` after retrieval will see
``None``.
"""
from __future__ import annotations
import asyncio
import json
import math
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple
import numpy as np
from ._persist import check_persisted_handles
from ._turbovec import IdMapIndex
try:
from haystack import Document
from haystack.dataclasses import ByteStream
from haystack.dataclasses.sparse_embedding import SparseEmbedding
from haystack.document_stores.errors import DuplicateDocumentError
from haystack.document_stores.types import DuplicatePolicy
from haystack.utils.filters import document_matches_filter
except ImportError as exc:
raise ImportError(
"haystack-ai is required to use turbovec.haystack. "
"Install with: pip install turbovec[haystack]"
) from exc
class TurboQuantDocumentStore:
"""Haystack DocumentStore backed by a :class:`~turbovec.IdMapIndex`.
Vectors are quantized to 2–4 bits per dimension. Full-precision
embeddings are dropped after quantization — callers requesting
``return_embedding=True`` on retrieval will see ``None`` on the
returned documents' ``embedding`` field regardless of the flag.
Example::
from turbovec.haystack import TurboQuantDocumentStore
from haystack import Document
store = TurboQuantDocumentStore(dim=1536, bit_width=4)
store.write_documents([
Document(content="...", embedding=[...], meta={"source": "a"}),
...
])
results = store.embedding_retrieval(query_embedding=[...], top_k=5)
"""
def __init__(
self,
dim: Optional[int] = None,
bit_width: int = 4,
*,
embedding_similarity_function: Literal["dot_product", "cosine"] = "cosine",
async_executor: Optional[ThreadPoolExecutor] = None,
return_embedding: bool = False,
) -> None:
"""
:param dim: Vector dimensionality. When omitted, the underlying
quantized index is created lazily by ``IdMapIndex`` itself on
the first ``write_documents`` call — matches the no-``dim``
ergonomics of ``InMemoryDocumentStore``.
:param bit_width: Quantization width per coordinate (2 or 4).
:param embedding_similarity_function: ``"cosine"`` (default) or
``"dot_product"``. Used to choose the ``scale_score`` formula
during retrieval. Defaults to ``"cosine"`` because turbovec
stores unit-normalized vectors.
:param async_executor: Optional executor for the ``*_async``
methods. If omitted, a single-threaded executor is created
and cleaned up on instance destruction.
:param return_embedding: Whether retrieval methods should leave
the ``embedding`` field populated on returned Documents.
turbovec never has the full-precision embedding available, so
this is always ``None`` either way; the flag is accepted for
API parity with ``InMemoryDocumentStore``.
"""
self._bit_width = bit_width
self.embedding_similarity_function = embedding_similarity_function
self.return_embedding = return_embedding
# IdMapIndex itself supports lazy construction — pass dim through
# and let it handle eager vs lazy. No per-store lazy wrapping.
self._index = IdMapIndex(dim, bit_width)
# Haystack doc_id (str) -> u64 handle
self._str_to_u64: Dict[str, int] = {}
# u64 handle -> stored doc data {id, content, meta}
self._u64_to_doc: Dict[int, Dict[str, Any]] = {}
# Counter for assigning u64 handles. Starts at 0; each new
# handle is `_next_u64 + 1`, then we bump. Plain int so pickle
# can round-trip it directly.
self._next_u64: int = 0
# Executor lifecycle mirrors InMemoryDocumentStore: own one when
# the caller didn't pass one in, and shut it down in __del__.
self._owns_executor = async_executor is None
self.executor = async_executor or ThreadPoolExecutor(
thread_name_prefix=f"async-turbovec-docstore-executor-{id(self)}",
max_workers=1,
)
def __del__(self) -> None:
if (
hasattr(self, "_owns_executor")
and self._owns_executor
and hasattr(self, "executor")
):
self.executor.shutdown(wait=True)
def shutdown(self) -> None:
"""Explicitly shut down the async executor if this store owns it."""
if self._owns_executor:
self.executor.shutdown(wait=True)
def _issue_handle(self) -> int:
self._next_u64 += 1
return self._next_u64
@property
def storage(self) -> Dict[str, Document]:
"""Map of ``doc_id -> Document`` for the currently stored documents.
Documents are reconstructed on every access; the
``embedding`` field is always ``None``.
"""
return {data["id"]: self._reconstruct(data) for data in self._u64_to_doc.values()}
# ---- DocumentStore protocol ---------------------------------------
def count_documents(self) -> int:
return len(self._str_to_u64)
def filter_documents(
self, filters: Optional[Dict[str, Any]] = None
) -> List[Document]:
if filters:
self._validate_filters(filters)
docs = [
self._reconstruct(data)
for data in self._u64_to_doc.values()
if document_matches_filter(filters=filters, document=self._reconstruct(data))
]
else:
docs = [self._reconstruct(data) for data in self._u64_to_doc.values()]
# `return_embedding` is informational here — we never have the
# full-precision embedding to begin with. Kept for parity.
return docs
def write_documents(
self,
documents: List[Document],
policy: DuplicatePolicy = DuplicatePolicy.NONE,
) -> int:
# Match InMemoryDocumentStore's input-shape validation rather
# than letting a bad input AttributeError on `.embedding`.
if (
not isinstance(documents, Iterable)
or isinstance(documents, str)
or any(not isinstance(doc, Document) for doc in documents)
):
raise ValueError("Please provide a list of Documents.")
if policy == DuplicatePolicy.NONE:
policy = DuplicatePolicy.FAIL
# First pass: validate and resolve duplicates according to policy.
# Duplicates are resolved against the batch-so-far as well as the
# existing store: InMemoryDocumentStore writes into its dict as it
# iterates, so a repeated id *within a single call* is resolved the
# same way a cross-call repeat would be. Without tracking the batch,
# every duplicate row still gets its own vector while _str_to_u64
# keeps only the last handle, orphaning the earlier vectors.
to_write: List[Document] = []
batch_pos: Dict[str, int] = {} # doc.id -> index into to_write
to_remove: List[str] = [] # existing ids to drop, deferred past add
written = len(documents)
for doc in documents:
if doc.embedding is None:
raise ValueError(
f"Document {doc.id!r} has no embedding. "
"TurboQuantDocumentStore only stores documents with precomputed "
"embeddings — run an embedder component before writing."
)
present = doc.id in self._str_to_u64 or doc.id in batch_pos
if policy != DuplicatePolicy.OVERWRITE and present:
if policy == DuplicatePolicy.FAIL:
raise DuplicateDocumentError(
f"ID '{doc.id}' already exists in the document store."
)
if policy == DuplicatePolicy.SKIP:
written -= 1
continue
if policy == DuplicatePolicy.OVERWRITE:
if doc.id in self._str_to_u64:
# Defer the removal until after the add succeeds so a
# failed validation/add never destroys existing data
# (issue #89).
to_remove.append(doc.id)
if doc.id in batch_pos:
# Last write wins: replace the earlier queued document
# in place rather than appending a second vector.
to_write[batch_pos[doc.id]] = doc
continue
batch_pos[doc.id] = len(to_write)
to_write.append(doc)
if not to_write:
return written
vectors = np.asarray(
[doc.embedding for doc in to_write], dtype=np.float32
)
if vectors.ndim != 2:
raise ValueError(
f"expected 2D embedding batch, got {vectors.ndim}D"
)
# IdMapIndex.add_with_ids handles both eager (dim must match) and
# lazy (locks dim on first call) cases. Surface its mismatch
# panic as a clean ValueError for parity with previous behaviour.
existing_dim = self._index.dim
if existing_dim is not None and vectors.shape[1] != existing_dim:
raise ValueError(
f"embedding dim {vectors.shape[1]} does not match store dim {existing_dim}"
)
if not vectors.flags["C_CONTIGUOUS"]:
vectors = np.ascontiguousarray(vectors)
handles = np.array(
[self._issue_handle() for _ in to_write], dtype=np.uint64
)
self._index.add_with_ids(vectors, handles)
# The add succeeded — now it's safe to drop the old vectors for any
# overwritten ids. Done before the mapping loop below so _remove_one
# resolves the old handle, not the one we're about to assign.
for doc_id in to_remove:
self._remove_one(doc_id)
for doc, handle in zip(to_write, handles):
h = int(handle)
self._str_to_u64[doc.id] = h
self._u64_to_doc[h] = {
"id": doc.id,
"content": doc.content,
"meta": dict(doc.meta),
"blob": doc.blob,
"sparse_embedding": doc.sparse_embedding,
}
return written
def delete_documents(self, document_ids: List[str]) -> None:
# Haystack's protocol says silently ignore missing ids.
for doc_id in document_ids:
self._remove_one(doc_id)
# ---- Utility methods (InMemoryDocumentStore parity) ---------------
def delete_all_documents(self) -> None:
"""Delete every document in the store."""
for doc_id in list(self._str_to_u64.keys()):
self._remove_one(doc_id)
def update_by_filter(
self, filters: Dict[str, Any], meta: Dict[str, Any]
) -> int:
"""Update metadata on every document matching ``filters``.
The new ``meta`` is merged into each matching document's existing
metadata. Embeddings are not touched — we never had them at full
precision anyway. Returns the number of documents updated.
"""
self._validate_filters(filters)
updated = 0
for data in self._u64_to_doc.values():
if document_matches_filter(filters=filters, document=self._reconstruct(data)):
data["meta"].update(meta)
updated += 1
return updated
def delete_by_filter(self, filters: Dict[str, Any]) -> int:
"""Delete every document matching ``filters``. Returns the count."""
self._validate_filters(filters)
matching_ids = [
data["id"]
for data in self._u64_to_doc.values()
if document_matches_filter(filters=filters, document=self._reconstruct(data))
]
for doc_id in matching_ids:
self._remove_one(doc_id)
return len(matching_ids)
def count_documents_by_filter(self, filters: Dict[str, Any]) -> int:
if filters:
self._validate_filters(filters)
return sum(
1
for data in self._u64_to_doc.values()
if document_matches_filter(filters=filters, document=self._reconstruct(data))
)
return self.count_documents()
def count_unique_metadata_by_filter(
self, filters: Dict[str, Any], metadata_fields: List[str]
) -> Dict[str, int]:
if filters:
self._validate_filters(filters)
docs_meta = [
data["meta"]
for data in self._u64_to_doc.values()
if document_matches_filter(filters=filters, document=self._reconstruct(data))
]
else:
docs_meta = [data["meta"] for data in self._u64_to_doc.values()]
result: Dict[str, int] = {}
for field in metadata_fields:
key = field.removeprefix("meta.") if field.startswith("meta.") else field
values = {meta.get(key) for meta in docs_meta if key in meta and meta[key] is not None}
result[key] = len(values)
return result
def get_metadata_fields_info(self) -> Dict[str, Dict[str, str]]:
type_map: Dict[str, str] = {}
for data in self._u64_to_doc.values():
for key, value in data["meta"].items():
if value is None:
continue
if isinstance(value, bool):
type_map[key] = "boolean"
elif isinstance(value, int):
type_map[key] = "int"
elif isinstance(value, float):
type_map[key] = "float"
else:
type_map[key] = "keyword"
return {k: {"type": v} for k, v in type_map.items()}
def get_metadata_field_min_max(self, metadata_field: str) -> Dict[str, Any]:
key = (
metadata_field.removeprefix("meta.")
if metadata_field.startswith("meta.")
else metadata_field
)
values = [
data["meta"][key]
for data in self._u64_to_doc.values()
if key in data["meta"]
and data["meta"][key] is not None
and isinstance(data["meta"][key], (int, float, str))
]
if not values:
return {"min": None, "max": None}
try:
return {"min": min(values), "max": max(values)}
except TypeError:
return {"min": None, "max": None}
def get_metadata_field_unique_values(
self, metadata_field: str, search_term: Optional[str] = None
) -> Tuple[List[str], int]:
key = (
metadata_field.removeprefix("meta.")
if metadata_field.startswith("meta.")
else metadata_field
)
if search_term:
docs_data = [
data
for data in self._u64_to_doc.values()
if data["content"] and search_term.lower() in data["content"].lower()
]
else:
docs_data = list(self._u64_to_doc.values())
values = sorted(
{
str(data["meta"][key])
for data in docs_data
if key in data["meta"] and data["meta"][key] is not None
},
key=str,
)
return values, len(values)
@staticmethod
def _validate_filters(filters: Optional[Dict[str, Any]]) -> None:
# Match InMemoryDocumentStore (document_store.py:504-509): a
# filter dict must have a top-level "operator" (simple comparison
# or logical) or "conditions" (compound). A bare "field" without
# an operator is malformed and the reference rejects it; we do too.
if (
filters
and "operator" not in filters
and "conditions" not in filters
):
raise ValueError(
"Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
)
# ---- Retrieval (not in core protocol but expected) ----------------
def embedding_retrieval(
self,
query_embedding: List[float],
filters: Optional[Dict[str, Any]] = None,
top_k: int = 10,
scale_score: bool = False,
return_embedding: Optional[bool] = None,
) -> List[Document]:
"""Return the ``top_k`` documents most similar to ``query_embedding``.
``return_embedding=None`` (default) honours the store-level
``return_embedding`` set in the constructor. turbovec never has
the full-precision embedding either way — the parameter is here
for API parity with ``InMemoryDocumentStore``.
``filters`` are resolved to an allowlist before scoring, so the
kernel never wastes work on non-matching documents and the result
count is always ``min(top_k, n_matches)`` rather than ``< top_k``
when the filter is selective.
"""
# `return_embedding` is accepted but we never have the full
# embedding to populate; left as-is for signature parity.
_ = return_embedding # noqa: F841
if self.count_documents() == 0:
return []
qvec = np.asarray(query_embedding, dtype=np.float32)
if qvec.ndim == 1:
qvec = qvec[None, :]
# By this point n_documents > 0, so the index has a committed dim.
expected_dim = self._index.dim
if qvec.shape[1] != expected_dim:
raise ValueError(
f"query_embedding dim {qvec.shape[1]} does not match store dim {expected_dim}"
)
if not qvec.flags["C_CONTIGUOUS"]:
qvec = np.ascontiguousarray(qvec)
if filters is None:
fetch_k = min(top_k, self.count_documents())
scores, handles = self._index.search(qvec, fetch_k)
else:
self._validate_filters(filters)
# Resolve filter → handle allowlist by walking the in-memory
# doc table once. This is the same O(N) cost as the old
# post-filter pass, just moved upfront so the kernel can score
# only matching vectors.
allowed_handles = [
handle
for handle, data in self._u64_to_doc.items()
if document_matches_filter(filters, self._reconstruct(data))
]
if not allowed_handles:
return []
allowlist = np.asarray(allowed_handles, dtype=np.uint64)
scores, handles = self._index.search(qvec, top_k, allowlist=allowlist)
out: List[Document] = []
for score, handle in zip(scores[0], handles[0]):
data = self._u64_to_doc[int(handle)]
out.append(self._reconstruct(data, score=float(score), scale_score=scale_score))
return out
# ---- Async variants ----------------------------------------------
async def count_documents_async(self) -> int:
return self.count_documents()
async def filter_documents_async(
self, filters: Optional[Dict[str, Any]] = None
) -> List[Document]:
return await asyncio.get_running_loop().run_in_executor(
self.executor, lambda: self.filter_documents(filters=filters)
)
async def write_documents_async(
self,
documents: List[Document],
policy: DuplicatePolicy = DuplicatePolicy.NONE,
) -> int:
return await asyncio.get_running_loop().run_in_executor(
self.executor, lambda: self.write_documents(documents=documents, policy=policy)
)
async def delete_documents_async(self, document_ids: List[str]) -> None:
await asyncio.get_running_loop().run_in_executor(
self.executor, lambda: self.delete_documents(document_ids=document_ids)
)
async def delete_all_documents_async(self) -> None:
await asyncio.get_running_loop().run_in_executor(
self.executor, self.delete_all_documents
)
async def update_by_filter_async(
self, filters: Dict[str, Any], meta: Dict[str, Any]
) -> int:
return await asyncio.get_running_loop().run_in_executor(
self.executor, lambda: self.update_by_filter(filters=filters, meta=meta)
)
async def count_documents_by_filter_async(self, filters: Dict[str, Any]) -> int:
return await asyncio.get_running_loop().run_in_executor(
self.executor, lambda: self.count_documents_by_filter(filters=filters)
)
async def count_unique_metadata_by_filter_async(
self, filters: Dict[str, Any], metadata_fields: List[str]
) -> Dict[str, int]:
return await asyncio.get_running_loop().run_in_executor(
self.executor,
lambda: self.count_unique_metadata_by_filter(
filters=filters, metadata_fields=metadata_fields
),
)
async def get_metadata_fields_info_async(self) -> Dict[str, Dict[str, str]]:
return await asyncio.get_running_loop().run_in_executor(
self.executor, self.get_metadata_fields_info
)
async def get_metadata_field_min_max_async(
self, metadata_field: str
) -> Dict[str, Any]:
return await asyncio.get_running_loop().run_in_executor(
self.executor,
lambda: self.get_metadata_field_min_max(metadata_field=metadata_field),
)
async def get_metadata_field_unique_values_async(
self, metadata_field: str, search_term: Optional[str] = None
) -> Tuple[List[str], int]:
return await asyncio.get_running_loop().run_in_executor(
self.executor,
lambda: self.get_metadata_field_unique_values(
metadata_field=metadata_field, search_term=search_term
),
)
async def embedding_retrieval_async(
self,
query_embedding: List[float],
filters: Optional[Dict[str, Any]] = None,
top_k: int = 10,
scale_score: bool = False,
return_embedding: Optional[bool] = None,
) -> List[Document]:
return await asyncio.get_running_loop().run_in_executor(
self.executor,
lambda: self.embedding_retrieval(
query_embedding=query_embedding,
filters=filters,
top_k=top_k,
scale_score=scale_score,
return_embedding=return_embedding,
),
)
# ---- Serialization (Pipeline to_dict / from_dict) -----------------
def to_dict(self) -> Dict[str, Any]:
return {
"type": f"{self.__class__.__module__}.{self.__class__.__name__}",
"init_parameters": {
# `_index.dim` is None on a lazy uncommitted store and an
# int once an add has locked the dim — both round-trip cleanly.
"dim": self._index.dim,
"bit_width": self._bit_width,
"embedding_similarity_function": self.embedding_similarity_function,
"return_embedding": self.return_embedding,
},
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "TurboQuantDocumentStore":
params = data.get("init_parameters", {})
return cls(**params)
# ---- Persistence -------------------------------------------------
# Side-car schema. Bump when the on-disk shape changes; loader
# accepts the current version plus any older versions whose missing
# fields we know how to reconstruct (currently v1, written before
# blob / sparse_embedding round-trip was added — both default to None
# on load).
_DOCSTORE_SCHEMA_VERSION = 2
_DOCSTORE_SCHEMA_COMPAT = (1, 2)
def save_to_disk(self, folder_path: str | Path) -> None:
"""Persist the quantized index plus the Haystack side-car to disk.
Writes into ``folder_path``:
- ``index.tvim`` — the :class:`IdMapIndex` payload. On a lazy
store that has never seen a write the file encodes the
uncommitted state via a ``dim=0`` sentinel.
- ``docstore.json`` — the str-id ↔ Document mapping and store
init parameters, JSON-encoded. Document metadata must be
JSON-serializable (the same constraint as
``InMemoryDocumentStore.save_to_disk``).
"""
folder = Path(folder_path)
folder.mkdir(parents=True, exist_ok=True)
self._index.write(str(folder / "index.tvim"))
# Keys in `_u64_to_doc` are ints (u64 handles); JSON object keys
# must be strings. Serialize as a list of [handle, data] pairs
# so we don't lose type fidelity on the round-trip.
payload = {
"schema_version": self._DOCSTORE_SCHEMA_VERSION,
"u64_to_doc": [
[h, self._serialize_doc_data(d)] for h, d in self._u64_to_doc.items()
],
"next_u64": self._next_u64,
"bit_width": self._bit_width,
"embedding_similarity_function": self.embedding_similarity_function,
"return_embedding": self.return_embedding,
}
with open(folder / "docstore.json", "w") as f:
json.dump(payload, f)
@classmethod
def load_from_disk(
cls,
folder_path: str | Path,
) -> "TurboQuantDocumentStore":
"""Reload a store from a folder previously written by
:meth:`save_to_disk`. Safe to call on any path — the side-car is
plain JSON, never pickle, so there's no deserialization-of-code
risk."""
folder = Path(folder_path)
with open(folder / "docstore.json") as f:
state = json.load(f)
version = state.get("schema_version", 0)
if version not in cls._DOCSTORE_SCHEMA_COMPAT:
raise ValueError(
f"docstore.json has schema version {version}; "
f"this turbovec accepts versions {list(cls._DOCSTORE_SCHEMA_COMPAT)}"
)
store = cls(
bit_width=state["bit_width"],
embedding_similarity_function=state.get(
"embedding_similarity_function", "cosine"
),
return_embedding=state.get("return_embedding", False),
)
# Reload the index — it carries dim internally (None for lazy
# uncommitted, int otherwise).
store._index = IdMapIndex.load(str(folder / "index.tvim"))
# Reconstruct {int handle: doc data} from the list-of-pairs form.
# `_deserialize_doc_data` is shape-tolerant: v1 entries lack the
# `blob` / `sparse_embedding` keys and come back with both set to
# None, which matches their original on-write state.
store._u64_to_doc = {
int(h): cls._deserialize_doc_data(d) for h, d in state["u64_to_doc"]
}
store._next_u64 = state["next_u64"]
# Rebuild str_to_u64 from the reloaded doc table.
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 ----------------------------------------------------
def _remove_one(self, doc_id: str) -> bool:
handle = self._str_to_u64.pop(doc_id, None)
if handle is None:
return False
del self._u64_to_doc[handle]
self._index.remove(handle)
return True
def _reconstruct(
self,
data: Dict[str, Any],
score: Optional[float] = None,
scale_score: bool = False,
) -> Document:
if score is not None and scale_score:
# Match Haystack's InMemoryDocumentStore._compute_query_embedding_similarity_scores
# (document_store.py:818-822): different formula per similarity
# function. turbovec uses unit-normalized vectors so the cosine
# branch is the natural default.
if self.embedding_similarity_function == "dot_product":
score = 1.0 / (1.0 + math.exp(-score / 100.0))
elif self.embedding_similarity_function == "cosine":
# Clamp to the exact cosine range before rescaling. Cauchy–Schwarz
# bounds the true cosine in [-1, 1], but the LUT scoring kernel's
# float-precision noise can land slightly outside that range on
# near-identical document/query pairs (e.g. a self-query under the
# length-renormalized estimator produces ~1.00016). Clamping
# preserves the [0, 1] contract for ``scale_score=True`` consumers.
score = (max(-1.0, min(1.0, score)) + 1.0) / 2.0
return Document(
id=data["id"],
content=data["content"],
meta=dict(data["meta"]),
blob=data.get("blob"),
sparse_embedding=data.get("sparse_embedding"),
score=score,
)
@staticmethod
def _serialize_doc_data(data: Dict[str, Any]) -> Dict[str, Any]:
# blob is a ByteStream; sparse_embedding is a SparseEmbedding.
# Both have a JSON-safe to_dict() form.
blob = data.get("blob")
sparse = data.get("sparse_embedding")
return {
"id": data["id"],
"content": data["content"],
"meta": data["meta"],
"blob": blob.to_dict() if blob is not None else None,
"sparse_embedding": sparse.to_dict() if sparse is not None else None,
}
@staticmethod
def _deserialize_doc_data(d: Dict[str, Any]) -> Dict[str, Any]:
blob = d.get("blob")
sparse = d.get("sparse_embedding")
return {
"id": d["id"],
"content": d["content"],
"meta": d["meta"],
"blob": ByteStream.from_dict(blob) if blob is not None else None,
"sparse_embedding": (
SparseEmbedding.from_dict(sparse) if sparse is not None else None
),
}
__all__ = ["TurboQuantDocumentStore"]
turbovec-python/python/turbovec/langchain.py
+558
View File
@@ -0,0 +1,558 @@
"""LangChain VectorStore backed by turbovec's quantized index.
Install with: ``pip install turbovec[langchain]``.
The public surface mirrors langchain_core's in-tree ``InMemoryVectorStore``
so this store can be swapped in wherever the in-memory store is used.
"""
from __future__ import annotations
import json
import uuid
from pathlib import Path
from typing import Any, Callable, Iterable, Sequence
import numpy as np
from ._dedup import DuplicatePolicy, resolve_duplicates
from ._persist import check_persisted_handles
from ._turbovec import IdMapIndex
try:
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore
except ImportError as exc:
raise ImportError(
"langchain-core is required to use turbovec.langchain. "
"Install with: pip install turbovec[langchain]"
) from exc
_INDEX_FILENAME = "index.tvim"
_STORE_FILENAME = "docstore.json"
# Bump when the docstore.json shape changes; loader refuses to deserialize
# unknown versions.
_DOCSTORE_SCHEMA_VERSION = 1
class TurboQuantVectorStore(VectorStore):
"""LangChain VectorStore backed by a :class:`IdMapIndex`.
Vectors are quantized to 2–4 bits per dimension. A side-car dictionary
holds the original text and metadata keyed by document id. Deletion
is supported in O(1) per id via the underlying :class:`IdMapIndex`.
"""
def __init__(
self,
embedding: Embeddings,
index: IdMapIndex | None = None,
*,
bit_width: int = 4,
docs: dict[str, tuple[str, dict[str, Any]]] | None = None,
str_to_u64: dict[str, int] | None = None,
next_u64: int = 0,
) -> None:
"""
:param embedding: LangChain ``Embeddings`` instance used to encode
documents and queries.
:param index: Optional pre-built :class:`IdMapIndex`. When omitted,
a lazy ``IdMapIndex`` is created — it commits to a dim on the
first add and lets us match the no-arg constructor pattern of
langchain_core's ``InMemoryVectorStore``.
:param bit_width: Quantization width (2 or 4) used when the index
is created from scratch. Ignored if ``index`` is supplied.
"""
self._embedding = embedding
# IdMapIndex itself supports lazy construction now — no per-store
# lazy wrapping needed. When `index` is None we create a lazy
# IdMapIndex(dim=None, bit_width) and let it handle the rest.
self._index = index if index is not None else IdMapIndex(bit_width=bit_width)
self._docs: dict[str, tuple[str, dict[str, Any]]] = docs if docs is not None else {}
self._str_to_u64: dict[str, int] = str_to_u64 if str_to_u64 is not None else {}
# Reverse map (u64 handle → str id) kept in sync so search results
# can translate handles back to LangChain document ids.
self._u64_to_str: dict[int, str] = {
handle: sid for sid, handle in self._str_to_u64.items()
}
self._next_u64: int = next_u64
def _issue_handle(self) -> int:
self._next_u64 += 1
return self._next_u64
@property
def embeddings(self) -> Embeddings:
return self._embedding
# ---- Relevance score normalization --------------------------------
def _select_relevance_score_fn(self) -> Callable[[float], float]:
# turbovec returns the raw inner product of unit-normalized vectors —
# ideally cosine similarity in [-1, 1]. Quantization noise can
# push that very slightly outside the bounds, so clamp after
# mapping to LangChain's [0, 1] relevance scale via (sim + 1) / 2.
return lambda sim: max(0.0, min(1.0, (sim + 1.0) / 2.0))
# ---- Write path ---------------------------------------------------
def add_texts(
self,
texts: Iterable[str],
metadatas: list[dict] | None = None,
ids: list[str] | None = None,
**_: Any,
) -> list[str]:
texts_list = list(texts)
if not texts_list:
return []
if metadatas is None:
metadatas = [{} for _ in texts_list]
if ids is None:
ids = [str(uuid.uuid4()) for _ in texts_list]
if len(metadatas) != len(texts_list) or len(ids) != len(texts_list):
raise ValueError("texts, metadatas, and ids must all have the same length")
vectors = np.asarray(self._embedding.embed_documents(texts_list), dtype=np.float32)
return self._store_texts_and_vectors(texts_list, vectors, metadatas, ids)
async def aadd_texts(
self,
texts: Iterable[str],
metadatas: list[dict] | None = None,
ids: list[str] | None = None,
**_: Any,
) -> list[str]:
texts_list = list(texts)
if not texts_list:
return []
if metadatas is None:
metadatas = [{} for _ in texts_list]
if ids is None:
ids = [str(uuid.uuid4()) for _ in texts_list]
if len(metadatas) != len(texts_list) or len(ids) != len(texts_list):
raise ValueError("texts, metadatas, and ids must all have the same length")
vectors = np.asarray(
await self._embedding.aembed_documents(texts_list), dtype=np.float32
)
return self._store_texts_and_vectors(texts_list, vectors, metadatas, ids)
def add_documents(
self,
documents: list[Document],
ids: list[str] | None = None,
**kwargs: Any,
) -> list[str]:
# Override the base class default which drops the entire `ids` array
# if any Document has a None id. The reference InMemoryVectorStore
# falls back per-document so partial ids are honoured.
texts = [doc.page_content for doc in documents]
metadatas = [doc.metadata for doc in documents]
if ids is None:
ids = [doc.id or str(uuid.uuid4()) for doc in documents]
return self.add_texts(texts=texts, metadatas=metadatas, ids=ids, **kwargs)
async def aadd_documents(
self,
documents: list[Document],
ids: list[str] | None = None,
**kwargs: Any,
) -> list[str]:
texts = [doc.page_content for doc in documents]
metadatas = [doc.metadata for doc in documents]
if ids is None:
ids = [doc.id or str(uuid.uuid4()) for doc in documents]
return await self.aadd_texts(
texts=texts, metadatas=metadatas, ids=ids, **kwargs
)
def _store_texts_and_vectors(
self,
texts_list: list[str],
vectors: np.ndarray,
metadatas: list[dict],
ids: list[str],
) -> list[str]:
if vectors.ndim != 2:
raise ValueError(f"expected 2D embedding batch, got {vectors.ndim}D")
# Dedup intra-batch duplicate ids, keeping the last occurrence —
# matches InMemoryVectorStore, whose dict store silently overwrites
# on a repeated id. Without this every row is added to the index but
# _str_to_u64 keeps only the last handle per id, orphaning the
# earlier vectors. The returned id list still mirrors the input
# (one entry per input text), as the reference does.
result_ids = ids
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]
vectors = vectors[keep]
# Validate before mutating any existing data. IdMapIndex.add_with_ids
# handles both eager (dim must match) and lazy (locks dim on first
# call) cases. Pre-check the eager case so we surface a clean
# ValueError rather than a Rust panic.
existing_dim = self._index.dim
if existing_dim is not None and vectors.shape[1] != existing_dim:
raise ValueError(
f"embedding dimension {vectors.shape[1]} does not match index dim {existing_dim}"
)
if not vectors.flags["C_CONTIGUOUS"]:
vectors = np.ascontiguousarray(vectors)
handles = np.array(
[self._issue_handle() for _ in texts_list], dtype=np.uint64
)
# Add first; if encoding rejects the batch (e.g. non-finite values)
# this raises before any existing data is touched. Only once the add
# has succeeded do we remove the old vectors for colliding ids, so a
# failed upsert never destroys existing data (issue #89). Handles are
# freshly issued, so the old and new vectors coexist until the delete.
self._index.add_with_ids(vectors, handles)
# Upsert: any id that already existed is removed so the re-added
# vector wins. Matches LangChain user expectation that `add_texts`
# with an existing id updates in place.
duplicates = [i for i in ids if i in self._str_to_u64]
if duplicates:
self.delete(duplicates)
for id_, text, meta, handle in zip(ids, texts_list, metadatas, handles):
h = int(handle)
self._str_to_u64[id_] = h
self._u64_to_str[h] = id_
self._docs[id_] = (text, dict(meta))
return result_ids
# ---- Read path (similarity search) --------------------------------
def similarity_search(
self,
query: str,
k: int = 4,
filter: dict[str, Any] | Callable[[Document], bool] | None = None,
**_: Any,
) -> list[Document]:
return [
doc
for doc, _score in self.similarity_search_with_score(query, k=k, filter=filter)
]
async def asimilarity_search(
self,
query: str,
k: int = 4,
filter: dict[str, Any] | Callable[[Document], bool] | None = None,
**_: Any,
) -> list[Document]:
return [
doc
for doc, _score in await self.asimilarity_search_with_score(
query, k=k, filter=filter
)
]
def similarity_search_with_score(
self,
query: str,
k: int = 4,
filter: dict[str, Any] | Callable[[Document], bool] | None = None,
**_: Any,
) -> list[tuple[Document, float]]:
qvec = np.asarray(self._embedding.embed_query(query), dtype=np.float32)
return self._search_vector(qvec, k, filter=filter)
async def asimilarity_search_with_score(
self,
query: str,
k: int = 4,
filter: dict[str, Any] | Callable[[Document], bool] | None = None,
**_: Any,
) -> list[tuple[Document, float]]:
qvec = np.asarray(
await self._embedding.aembed_query(query), dtype=np.float32
)
return self._search_vector(qvec, k, filter=filter)
def similarity_search_by_vector(
self,
embedding: list[float],
k: int = 4,
filter: dict[str, Any] | Callable[[Document], bool] | None = None,
**_: Any,
) -> list[Document]:
qvec = np.asarray(embedding, dtype=np.float32)
return [doc for doc, _score in self._search_vector(qvec, k, filter=filter)]
async def asimilarity_search_by_vector(
self,
embedding: list[float],
k: int = 4,
filter: dict[str, Any] | Callable[[Document], bool] | None = None,
**_: Any,
) -> list[Document]:
# The search itself is sync (no embedding step). Delegate.
return self.similarity_search_by_vector(embedding, k=k, filter=filter)
def _search_vector(
self,
qvec: np.ndarray,
k: int,
filter: dict[str, Any] | Callable[[Document], bool] | None = None,
) -> list[tuple[Document, float]]:
if qvec.ndim == 1:
qvec = qvec[None, :]
if not qvec.flags["C_CONTIGUOUS"]:
qvec = np.ascontiguousarray(qvec)
# IdMapIndex handles the lazy-uncommitted case internally (returns
# empty search results). A len-zero check covers both that and
# the eager-but-empty case.
if len(self._index) == 0:
return []
if filter is None:
search_k = min(k, len(self._index))
scores, handles = self._index.search(qvec, search_k)
else:
predicate = self._compile_filter(filter)
allowed_handles = [
self._str_to_u64[sid]
for sid, (text, meta) in self._docs.items()
if predicate(Document(id=sid, page_content=text, metadata=dict(meta)))
]
if not allowed_handles:
return []
allowlist = np.asarray(allowed_handles, dtype=np.uint64)
scores, handles = self._index.search(qvec, k, allowlist=allowlist)
results: list[tuple[Document, float]] = []
for score, handle in zip(scores[0], handles[0]):
sid = self._u64_to_str[int(handle)]
text, meta = self._docs[sid]
results.append(
(Document(id=sid, page_content=text, metadata=dict(meta)), float(score))
)
return results
@staticmethod
def _compile_filter(
filter: dict[str, Any] | Callable[[Document], bool],
) -> Callable[[Document], bool]:
# Match the in-tree InMemoryVectorStore convention: callable filters
# receive a Document, not a metadata dict
# (langchain_core/vectorstores/in_memory.py).
if callable(filter):
return filter
if isinstance(filter, dict):
items = list(filter.items())
return lambda doc: all(doc.metadata.get(k) == v for k, v in items)
raise TypeError(
"filter must be a dict of metadata key/value pairs or a callable "
f"taking a Document, got {type(filter).__name__}"
)
# ---- Max marginal relevance ---------------------------------------
#
# MMR requires the full-precision vector of every candidate to compute
# pairwise diversity scores. turbovec discards full vectors after
# quantization (that's the point), so we can't faithfully implement
# MMR. Raise loudly with a useful message rather than silently fall
# back to the base class's bare NotImplementedError.
_MMR_MSG = (
"TurboQuantVectorStore does not support max-marginal-relevance "
"search because the underlying quantized index discards "
"full-precision vectors after compression. MMR requires the "
"original embedding for every candidate to compute pairwise "
"diversity. Use `similarity_search` / `similarity_search_with_score` "
"instead, or maintain a parallel store with full-precision "
"embeddings if you need MMR specifically."
)
def max_marginal_relevance_search(
self,
query: str,
k: int = 4,
fetch_k: int = 20,
lambda_mult: float = 0.5,
**kwargs: Any,
) -> list[Document]:
raise NotImplementedError(self._MMR_MSG)
def max_marginal_relevance_search_by_vector(
self,
embedding: list[float],
k: int = 4,
fetch_k: int = 20,
lambda_mult: float = 0.5,
*,
filter: Callable[[Document], bool] | None = None,
**kwargs: Any,
) -> list[Document]:
raise NotImplementedError(self._MMR_MSG)
async def amax_marginal_relevance_search(
self,
query: str,
k: int = 4,
fetch_k: int = 20,
lambda_mult: float = 0.5,
**kwargs: Any,
) -> list[Document]:
raise NotImplementedError(self._MMR_MSG)
# ---- Get / delete -------------------------------------------------
def get_by_ids(self, ids: Sequence[str], /) -> list[Document]:
"""Return Documents for the given ids. Missing ids are silently skipped
(matches the InMemoryVectorStore reference)."""
out: list[Document] = []
for sid in ids:
if sid in self._docs:
text, meta = self._docs[sid]
out.append(Document(id=sid, page_content=text, metadata=dict(meta)))
return out
async def aget_by_ids(self, ids: Sequence[str], /) -> list[Document]:
return self.get_by_ids(ids)
def delete(self, ids: list[str] | None = None, **_: Any) -> None:
"""Remove documents by id. Missing ids are silently skipped — matches
the InMemoryVectorStore reference (which also accepts ``ids=None``
as a no-op)."""
if not ids:
return
for sid in ids:
handle = self._str_to_u64.pop(sid, None)
if handle is None:
continue
self._u64_to_str.pop(handle, None)
self._docs.pop(sid, None)
self._index.remove(handle)
async def adelete(self, ids: list[str] | None = None, **_: Any) -> None:
self.delete(ids)
# ---- Construction helpers -----------------------------------------
@classmethod
def from_texts(
cls,
texts: list[str],
embedding: Embeddings,
metadatas: list[dict] | None = None,
*,
bit_width: int = 4,
ids: list[str] | None = None,
**_: Any,
) -> "TurboQuantVectorStore":
# The underlying index is created lazily on the first `add_texts`
# call, picking up `dim` from the first batch of embeddings — same
# no-`dim` ergonomics as InMemoryVectorStore.
store = cls(embedding=embedding, bit_width=bit_width)
if texts:
store.add_texts(texts, metadatas=metadatas, ids=ids)
return store
@classmethod
async def afrom_texts(
cls,
texts: list[str],
embedding: Embeddings,
metadatas: list[dict] | None = None,
*,
bit_width: int = 4,
ids: list[str] | None = None,
**_: Any,
) -> "TurboQuantVectorStore":
store = cls(embedding=embedding, bit_width=bit_width)
if texts:
await store.aadd_texts(texts, metadatas=metadatas, ids=ids)
return store
# ---- Persistence --------------------------------------------------
#
# Method names match the InMemoryVectorStore reference (`dump`/`load`),
# but the on-disk layout is a folder containing the binary index file
# plus a JSON side-car (we can't embed the binary Rust index in a
# single JSON file the way the reference does with its raw-vector
# store).
def dump(self, folder_path: str | Path) -> None:
"""Persist the quantized index plus the side-car to disk.
``folder_path`` is a directory; turbovec writes ``index.tvim``
and ``docstore.json`` inside it. Document metadata must be
JSON-serializable (same constraint as ``InMemoryVectorStore``).
A lazy uncommitted index encodes its state via the index file's
own ``dim = 0`` sentinel; no special-case handling needed here.
"""
folder = Path(folder_path)
folder.mkdir(parents=True, exist_ok=True)
self._index.write(str(folder / _INDEX_FILENAME))
# `_docs` stores tuples `(text, metadata)` — JSON would drop the
# tuple-ness on round-trip, so serialize each entry as an explicit
# `{"text": ..., "metadata": ...}` dict.
docs_payload = {
sid: {"text": text, "metadata": meta}
for sid, (text, meta) in self._docs.items()
}
payload = {
"schema_version": _DOCSTORE_SCHEMA_VERSION,
"docs": docs_payload,
"str_to_u64": self._str_to_u64,
"next_u64": self._next_u64,
# Pull bit_width off the live index — same value whether
# the index was constructed eagerly or lazily.
"bit_width": self._index.bit_width,
}
with open(folder / _STORE_FILENAME, "w") as f:
json.dump(payload, f)
@classmethod
def load(
cls,
folder_path: str | Path,
embedding: Embeddings,
) -> "TurboQuantVectorStore":
"""Reload a store from a folder previously written by :meth:`dump`.
Safe to call on any path — the side-car is plain JSON, never
pickle, so there's no deserialization-of-code risk."""
folder = Path(folder_path)
with open(folder / _STORE_FILENAME) as f:
state = json.load(f)
version = state.get("schema_version", 0)
if version != _DOCSTORE_SCHEMA_VERSION:
raise ValueError(
f"docstore.json has schema version {version}; "
f"this turbovec expects version {_DOCSTORE_SCHEMA_VERSION}"
)
# IdMapIndex.load handles the dim=0 (lazy-uncommitted) sentinel
# internally and reconstructs the index in the right state.
index = IdMapIndex.load(str(folder / _INDEX_FILENAME))
# Rehydrate `_docs` from the explicit `{"text", "metadata"}` form
# back into the internal tuple representation.
docs = {
sid: (entry["text"], entry["metadata"])
for sid, entry in state["docs"].items()
}
# 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,
bit_width=state.get("bit_width", 4),
docs=docs,
str_to_u64=str_to_u64,
next_u64=int(state["next_u64"]),
)
__all__ = ["TurboQuantVectorStore"]
turbovec-python/python/turbovec/llama_index.py
+678
View File
@@ -0,0 +1,678 @@
"""LlamaIndex VectorStore backed by turbovec's quantized index.
Install with: ``pip install turbovec[llama-index]``.
"""
from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Any, List, Optional, Sequence
import numpy as np
from ._dedup import DuplicatePolicy, resolve_duplicates
from ._persist import check_persisted_handles
from ._turbovec import IdMapIndex
try:
from llama_index.core.bridge.pydantic import PrivateAttr
from llama_index.core.schema import (
BaseNode,
NodeRelationship,
RelatedNodeInfo,
TextNode,
)
from llama_index.core.vector_stores.types import (
BasePydanticVectorStore,
FilterCondition,
FilterOperator,
MetadataFilter,
MetadataFilters,
VectorStoreQuery,
VectorStoreQueryMode,
VectorStoreQueryResult,
)
from llama_index.core.vector_stores.utils import (
metadata_dict_to_node,
node_to_metadata_dict,
)
except ImportError as exc:
raise ImportError(
"llama-index-core is required to use turbovec.llama_index. "
"Install with: pip install turbovec[llama-index]"
) from exc
# Persistence layout: a single user-facing ``persist_path`` (the file the
# framework asks us to write) is split into two files by extension —
# ``{base}.tvim`` for the binary IdMapIndex and ``{base}.nodes.json`` for
# the node side-car. Both live next to each other so the layout fits the
# directory-of-namespaced-files pattern used by StorageContext.
_INDEX_EXT = ".tvim"
_STORE_EXT = ".nodes.json"
# Filename template used by SimpleVectorStore for namespace lookup —
# mirrored here so `from_persist_dir` works the same way.
_NAMESPACE_SEP = "__"
_DEFAULT_PERSIST_FNAME = "vector_store.json"
_DEFAULT_VECTOR_STORE = "default"
# Bump when the nodes.json shape changes; loader accepts the current
# version plus any older versions whose missing fields we know how to
# reconstruct (currently v1, written before full-node round-trip was
# added — v1 entries are reconstructed as bare TextNodes with only
# text + metadata + SOURCE relationship, matching the original
# lossy behaviour rather than failing to load).
_NODES_SCHEMA_VERSION = 2
_NODES_SCHEMA_COMPAT = (1, 2)
def _split_persist_base(persist_path: str | Path) -> Path:
"""Strip the framework-provided extension off `persist_path` so the
binary index and JSON side-car can sit next to each other under a
shared base. We then append our own extensions in persist / load."""
p = Path(persist_path)
# Use the path without its suffix so both .tvim and .nodes.json share
# a base. If the input has no suffix (e.g. a bare folder-like name),
# use it as-is.
return p.with_suffix("") if p.suffix else p
class TurboQuantVectorStore(BasePydanticVectorStore):
"""LlamaIndex VectorStore backed by a :class:`IdMapIndex`.
Vectors are quantized to 2–4 bits per dimension. A side-car dictionary
holds node text and metadata keyed by ``node_id``. Supports ``delete``
(by ``ref_doc_id``, removing every node with that ref) and
``delete_nodes`` (by ``node_id``) — both O(1) per node.
"""
stores_text: bool = True
is_embedding_query: bool = True
flat_metadata: bool = False
_index: Any = PrivateAttr()
_nodes: dict[str, dict[str, Any]] = PrivateAttr()
_node_id_to_u64: dict[str, int] = PrivateAttr()
_u64_to_node_id: dict[int, str] = PrivateAttr()
_next_u64: int = PrivateAttr()
def __init__(self, index: IdMapIndex | None = None, *, bit_width: int = 4, **kwargs: Any) -> None:
"""Construct the vector store.
:param index: Optional pre-built :class:`IdMapIndex`. When omitted,
a lazy ``IdMapIndex`` is created — it commits to a dim on the
first add and lets callers use the no-arg construction pattern
common to LlamaIndex's other vector stores (e.g. via
``StorageContext.from_defaults(vector_store=TurboQuantVectorStore())``).
:param bit_width: Quantization width used when constructing the
lazy index. Ignored if ``index`` is supplied.
"""
super().__init__(**kwargs)
# IdMapIndex itself supports lazy construction now — no per-store
# lazy wrapping needed.
self._index = index if index is not None else IdMapIndex(bit_width=bit_width)
self._nodes = {}
self._node_id_to_u64 = {}
self._u64_to_node_id = {}
self._next_u64 = 0
def _issue_handle(self) -> int:
self._next_u64 += 1
return self._next_u64
@classmethod
def class_name(cls) -> str:
return "TurboQuantVectorStore"
@classmethod
def from_params(cls, dim: int | None = None, bit_width: int = 4) -> "TurboQuantVectorStore":
"""Build a store with a known ``dim`` (eager) or lazy when ``dim``
is omitted."""
return cls(index=IdMapIndex(dim, bit_width))
@property
def client(self) -> IdMapIndex:
return self._index
def add(self, nodes: list[BaseNode], **_: Any) -> list[str]:
if not nodes:
return []
# Reject intra-batch duplicates loudly. Letting them through would
# leave the index with N vectors but only the last node_id mapped
# back to one of them — the earlier handles become orphans that
# `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.
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)
if vectors.ndim != 2:
raise ValueError(
f"expected 2D embedding batch, got {vectors.ndim}D"
)
# IdMapIndex.add_with_ids handles eager (dim must match) and lazy
# (locks dim on first add) — pre-check the eager case so we
# surface a clean ValueError rather than a Rust panic.
existing_dim = self._index.dim
if existing_dim is not None and vectors.shape[1] != existing_dim:
raise ValueError(
f"node embedding dim {vectors.shape[1]} does not match index dim {existing_dim}"
)
if not vectors.flags["C_CONTIGUOUS"]:
vectors = np.ascontiguousarray(vectors)
handles = np.array([self._issue_handle() for _ in nodes], dtype=np.uint64)
# Add first; if validation above or encoding here (e.g. non-finite
# values) rejects the batch, it raises before any existing data is
# touched. Only after the add succeeds do we remove the old entries
# for colliding node_ids, so a failed upsert never destroys existing
# data (issue #89). Handles are freshly issued, so the old and new
# vectors coexist until the delete.
self._index.add_with_ids(vectors, handles)
# Upsert-like: if a node_id is already present in the STORE, remove
# the old entry so the new embedding wins.
duplicates = [n.node_id for n in nodes if n.node_id in self._node_id_to_u64]
for node_id in duplicates:
self._remove_node_by_id(node_id)
ids: list[str] = []
for node, handle in zip(nodes, handles):
h = int(handle)
nid = node.node_id
self._node_id_to_u64[nid] = h
self._u64_to_node_id[h] = nid
# `metadata` and `ref_doc_id` are kept at top level for fast
# filter / doc-id lookup (queries hit these on every hit;
# parsing _node_content per hit would be wasteful). `node_dict`
# is the framework's canonical metadata representation
# (`_node_content` + `_node_type` + original metadata keys),
# which `metadata_dict_to_node` reconstructs into a full
# BaseNode — preserving relationships (PREVIOUS / NEXT /
# PARENT / CHILD), excluded_*_metadata_keys, template fields,
# start/end_char_idx and mimetype on retrieval. The narrow
# `{text, metadata, ref_doc_id}` schema we used to keep lost
# all of those silently.
self._nodes[nid] = {
"metadata": dict(node.metadata),
"ref_doc_id": node.ref_doc_id,
"node_dict": node_to_metadata_dict(
node, remove_text=False, flat_metadata=False
),
}
ids.append(nid)
return ids
def delete(self, ref_doc_id: str, **_: Any) -> None:
"""Delete every node whose ``ref_doc_id`` matches."""
matching = [
nid for nid, data in self._nodes.items() if data.get("ref_doc_id") == ref_doc_id
]
for nid in matching:
self._remove_node_by_id(nid)
def delete_nodes(
self,
node_ids: Optional[List[str]] = None,
filters: Optional[MetadataFilters] = None,
**_: Any,
) -> None:
"""Delete every node matching ``node_ids`` and/or ``filters``. Both
constraints intersect when supplied. Missing node_ids are ignored.
Matches the signature and semantics of ``SimpleVectorStore.delete_nodes``.
"""
if not node_ids and filters is None:
return
candidates = list(self._nodes.items())
if node_ids is not None:
node_id_set = set(node_ids)
candidates = [(nid, data) for nid, data in candidates if nid in node_id_set]
if filters is not None:
candidates = [
(nid, data)
for nid, data in candidates
if self._filters_match(data["metadata"], filters)
]
for nid, _data in candidates:
self._remove_node_by_id(nid)
def clear(self) -> None:
"""Drop every node from the store and reset to a fresh lazy index.
The new index keeps the same ``bit_width`` so subsequent adds
commit a new ``dim`` lazily.
"""
bw = self._index.bit_width
self._index = IdMapIndex(bit_width=bw)
self._nodes = {}
self._node_id_to_u64 = {}
self._u64_to_node_id = {}
self._next_u64 = 0
def get(self, text_id: str) -> List[float]:
"""LlamaIndex's protocol expects this to return the full-precision
embedding for a given node id. turbovec discards full-precision
embeddings after quantization, so we raise loudly with an
explanation rather than return a lossy reconstruction or zeroes.
"""
raise NotImplementedError(
"TurboQuantVectorStore.get(text_id) cannot return the original "
"embedding because turbovec quantizes vectors to 2-4 bits per "
"dimension and discards full precision after encoding. Keep a "
"parallel docstore if you need the raw embedding."
)
def get_nodes(
self,
node_ids: Optional[List[str]] = None,
filters: Optional[MetadataFilters] = None,
) -> List[BaseNode]:
"""Return the nodes matching ``node_ids`` and/or ``filters``. Both
constraints intersect when supplied; missing node_ids are
silently skipped.
Unlike ``SimpleVectorStore`` (which raises NotImplementedError
here because it doesn't store nodes), turbovec keeps node text
and metadata in a side-car so this can return populated
``TextNode`` objects directly.
"""
candidates = list(self._nodes.items())
if node_ids is not None:
node_id_set = set(node_ids)
candidates = [(nid, data) for nid, data in candidates if nid in node_id_set]
if filters is not None:
candidates = [
(nid, data)
for nid, data in candidates
if self._filters_match(data["metadata"], filters)
]
return [self._reconstruct_node(nid, data) for nid, data in candidates]
@staticmethod
def _reconstruct_node(nid: str, data: dict[str, Any]) -> BaseNode:
# v2 entries carry `node_dict` — round-trip via the framework's
# own helper so we get the full BaseNode subclass back
# (TextNode / IndexNode / ImageNode) with every field populated.
if "node_dict" in data:
return metadata_dict_to_node(data["node_dict"])
# v1 fallback: stores that were persisted before the full-node
# round-trip landed only have {text, metadata, ref_doc_id}.
# Reconstruct the minimum-fidelity TextNode they used to produce
# so old on-disk stores keep loading without manual migration.
node = TextNode(
id_=nid,
text=data["text"],
metadata=dict(data["metadata"]),
)
if data.get("ref_doc_id") is not None:
node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(
node_id=data["ref_doc_id"]
)
return node
def _remove_node_by_id(self, node_id: str) -> bool:
handle = self._node_id_to_u64.pop(node_id, None)
if handle is None:
return False
self._u64_to_node_id.pop(handle, None)
self._nodes.pop(node_id, None)
self._index.remove(handle)
return True
def _resolve_allowed_handles(
self,
filters: MetadataFilters | None,
node_ids: list[str] | None,
doc_ids: list[str] | None,
) -> list[int]:
"""Resolve ``query.filters``, ``query.node_ids`` and ``query.doc_ids``
to the list of internal u64 handles that satisfy the filter. Empty
list means no node matches.
Semantics (matching the SimpleVectorStore reference where applicable):
- ``node_ids``: filter by node_id (set membership).
- ``doc_ids``: filter by ``ref_doc_id`` only (source document id).
- ``filters``: apply metadata filters.
All three intersect when more than one is supplied.
"""
candidates = self._nodes.items()
if node_ids:
node_id_set = set(node_ids)
candidates = [(nid, data) for nid, data in candidates if nid in node_id_set]
if doc_ids:
doc_id_set = set(doc_ids)
candidates = [
(nid, data)
for nid, data in candidates
if data.get("ref_doc_id") in doc_id_set
]
if filters is None:
return [self._node_id_to_u64[nid] for nid, _ in candidates]
return [
self._node_id_to_u64[nid]
for nid, data in candidates
if self._filters_match(data["metadata"], filters)
]
@classmethod
def _filters_match(
cls, metadata: dict[str, Any], filters: MetadataFilters
) -> bool:
condition = getattr(filters, "condition", None) or FilterCondition.AND
results: list[bool] = []
for f in filters.filters:
if isinstance(f, MetadataFilters):
results.append(cls._filters_match(metadata, f))
else:
results.append(cls._single_filter_match(metadata, f))
if condition == FilterCondition.AND:
return all(results) if results else True
if condition == FilterCondition.OR:
return any(results) if results else True
if condition == FilterCondition.NOT:
# Reference semantics (`build_metadata_filter_fn`,
# `utils.py:187-189`): NOT matches when none of the inner
# filters match. Empty inner list trivially satisfies NOT.
return not any(results)
raise NotImplementedError(
f"filter condition {condition!r} not supported by TurboQuantVectorStore"
)
@staticmethod
def _single_filter_match(metadata: dict[str, Any], f: MetadataFilter) -> bool:
# Semantics mirror SimpleVectorStore's _build_metadata_filter_fn
# (llama_index/core/vector_stores/simple.py) so that filtered
# results agree with the in-tree reference store.
op = f.operator
target = f.value
value = metadata.get(f.key)
# IS_EMPTY is the only operator that treats a missing key as a hit.
if op == FilterOperator.IS_EMPTY:
return value is None or value == "" or value == []
# Every other operator returns False when the key is absent — this
# matches the reference implementation (notably NE returns False on
# missing, not True).
if value is None:
return False
if op == FilterOperator.EQ:
return value == target
if op == FilterOperator.NE:
return value != target
if op == FilterOperator.GT:
return value > target
if op == FilterOperator.LT:
return value < target
if op == FilterOperator.GTE:
return value >= target
if op == FilterOperator.LTE:
return value <= target
if op == FilterOperator.IN:
return value in target
if op == FilterOperator.NIN:
return value not in target
if op == FilterOperator.CONTAINS:
return target in value
if op == FilterOperator.TEXT_MATCH:
# Reference (`utils.py:138-144`): case-SENSITIVE substring,
# both sides must be strings. Previous turbovec impl
# lowercased both sides — a silent semantic divergence that
# caused our results to disagree with SimpleVectorStore on
# mixed-case keys.
if isinstance(target, str) and isinstance(value, str):
return target in value
raise TypeError(
"Both metadata value and filter value must be strings "
"for the TEXT_MATCH operator"
)
if op == FilterOperator.TEXT_MATCH_INSENSITIVE:
if isinstance(target, str) and isinstance(value, str):
return target.lower() in value.lower()
raise TypeError(
"Both metadata value and filter value must be strings "
"for the TEXT_MATCH_INSENSITIVE operator"
)
if op == FilterOperator.ALL:
# Reference (`utils.py:152-153`): every element of `target`
# must be present in the metadata value (which is typically
# a list — tag-set matching).
return all(t in value for t in target)
if op == FilterOperator.ANY:
return any(t in value for t in target)
raise NotImplementedError(
f"filter operator {op!r} not supported by TurboQuantVectorStore"
)
def query(self, query: VectorStoreQuery, **_: Any) -> VectorStoreQueryResult:
# MMR / SVM / LINEAR_REGRESSION / HYBRID etc. all need access to
# full-precision vectors (for pairwise diversity, learned scoring,
# or sparse-dense fusion). turbovec discards full precision after
# quantization, so any non-DEFAULT mode is unsupportable here.
# Raise loudly instead of silently treating it as DEFAULT, which
# the previous impl did and which let callers think they were
# getting e.g. MMR diversity when they were not.
if query.mode != VectorStoreQueryMode.DEFAULT:
raise NotImplementedError(
f"TurboQuantVectorStore does not support query mode "
f"{query.mode!r}. Only VectorStoreQueryMode.DEFAULT is "
"supported — MMR / SVM / hybrid modes need access to "
"full-precision vectors which turbovec discards after "
"quantization. Maintain a parallel store with full vectors "
"if you need a non-default scoring mode."
)
if query.query_embedding is None:
raise ValueError(
"TurboQuantVectorStore requires a pre-computed query_embedding "
"(is_embedding_query=True)."
)
qvec = np.asarray(query.query_embedding, dtype=np.float32)
if qvec.ndim == 1:
qvec = qvec[None, :]
if not qvec.flags["C_CONTIGUOUS"]:
qvec = np.ascontiguousarray(qvec)
if len(self._index) == 0:
return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
has_filters = (
query.filters is not None
or bool(query.node_ids)
or bool(query.doc_ids)
)
if not has_filters:
k = min(query.similarity_top_k, len(self._index))
scores, handles = self._index.search(qvec, k)
else:
allowed_handles = self._resolve_allowed_handles(
query.filters, query.node_ids, query.doc_ids
)
if not allowed_handles:
return VectorStoreQueryResult(nodes=[], similarities=[], ids=[])
allowlist = np.asarray(allowed_handles, dtype=np.uint64)
scores, handles = self._index.search(
qvec, query.similarity_top_k, allowlist=allowlist
)
result_nodes: list[TextNode] = []
similarities: list[float] = []
ids: list[str] = []
for score, handle in zip(scores[0], handles[0]):
nid = self._u64_to_node_id[int(handle)]
data = self._nodes[nid]
result_nodes.append(self._reconstruct_node(nid, data))
similarities.append(float(score))
ids.append(nid)
return VectorStoreQueryResult(nodes=result_nodes, similarities=similarities, ids=ids)
# ---- Async overrides --------------------------------------------------
#
# The base class provides default async impls that delegate to sync via
# `return self.<sync>(...)`. We override them explicitly so the signature
# is visible on the class and an autodoc tool / IDE doesn't make
# callers chase the abstract base class for the documentation.
async def async_add(
self, nodes: Sequence[BaseNode], **kwargs: Any
) -> List[str]:
return self.add(list(nodes), **kwargs)
async def adelete(self, ref_doc_id: str, **kwargs: Any) -> None:
self.delete(ref_doc_id, **kwargs)
async def adelete_nodes(
self,
node_ids: Optional[List[str]] = None,
filters: Optional[MetadataFilters] = None,
**kwargs: Any,
) -> None:
self.delete_nodes(node_ids=node_ids, filters=filters, **kwargs)
async def aclear(self) -> None:
self.clear()
async def aquery(
self, query: VectorStoreQuery, **kwargs: Any
) -> VectorStoreQueryResult:
return self.query(query, **kwargs)
async def aget_nodes(
self,
node_ids: Optional[List[str]] = None,
filters: Optional[MetadataFilters] = None,
) -> List[BaseNode]:
return self.get_nodes(node_ids=node_ids, filters=filters)
# ---- Config serialization ---------------------------------------------
def to_dict(self, **_: Any) -> dict[str, Any]:
"""Serialize the store's *configuration* (not its data) so a
fresh instance can be reconstructed via ``from_dict``. Mirrors
the contract of ``SimpleVectorStore.to_dict`` — config-only;
node data round-trips through ``persist`` / ``from_persist_path``.
"""
return {
"bit_width": self._index.bit_width,
"dim": self._index.dim, # may be None (lazy uncommitted)
}
@classmethod
def from_dict(cls, data: dict[str, Any], **_: Any) -> "TurboQuantVectorStore":
"""Construct an empty store from a config dict produced by
``to_dict``. To restore data, use ``from_persist_path``."""
dim = data.get("dim")
bit_width = data.get("bit_width", 4)
return cls(index=IdMapIndex(dim, bit_width))
def persist(self, persist_path: str, fs: Any = None) -> None:
"""Persist the store. ``persist_path`` is treated as a path *stem*:
the binary index goes to ``{stem}.tvim`` and the node side-car to
``{stem}.nodes.json``. Any extension on ``persist_path`` (e.g.
``.json`` from a StorageContext default) is replaced.
This matches the layout assumed by ``StorageContext.persist`` —
which calls us with ``persist_path = {persist_dir}/{namespace}__vector_store.json`` —
and lets multiple namespaced stores coexist in the same directory.
Node metadata must be JSON-serializable (same constraint as
``SimpleVectorStore``). ``fs`` (fsspec) is not yet supported;
pass a local path.
"""
if fs is not None:
raise NotImplementedError(
"fsspec filesystems are not supported yet; pass a local path."
)
base = _split_persist_base(persist_path)
base.parent.mkdir(parents=True, exist_ok=True)
self._index.write(str(base.with_suffix(_INDEX_EXT)))
payload = {
"schema_version": _NODES_SCHEMA_VERSION,
"nodes": self._nodes,
# JSON object keys must be strings; round-trip int keys via
# an explicit list of [node_id, handle] pairs to preserve
# type fidelity.
"node_id_to_u64": list(self._node_id_to_u64.items()),
"next_u64": self._next_u64,
}
with open(base.with_suffix(_STORE_EXT), "w") as f:
json.dump(payload, f)
@classmethod
def from_persist_path(
cls,
persist_path: str,
fs: Any = None,
) -> "TurboQuantVectorStore":
"""Load a previously-persisted store. ``persist_path`` is the same
path that was passed to :meth:`persist` (extension is ignored;
``{stem}.tvim`` and ``{stem}.nodes.json`` are read).
Safe to call on any path — the side-car is plain JSON, never
pickle, so there's no deserialization-of-code risk.
"""
if fs is not None:
raise NotImplementedError(
"fsspec filesystems are not supported yet; pass a local path."
)
base = _split_persist_base(persist_path)
index = IdMapIndex.load(str(base.with_suffix(_INDEX_EXT)))
with open(base.with_suffix(_STORE_EXT)) as f:
state = json.load(f)
version = state.get("schema_version", 0)
if version not in _NODES_SCHEMA_COMPAT:
raise ValueError(
f"{_STORE_EXT.lstrip('.')} has schema version {version}; "
f"this turbovec accepts versions {list(_NODES_SCHEMA_COMPAT)}"
)
store = cls(index=index)
# v1 entries lack `node_dict` and reconstruct as narrow TextNodes;
# v2 entries carry it and reconstruct with full BaseNode fidelity.
# `_reconstruct_node` dispatches on shape, so we just load the
# dict as-is.
store._nodes = state["nodes"]
# Reconstruct {node_id: int handle} from the list-of-pairs form.
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
def from_persist_dir(
cls,
persist_dir: str,
namespace: str = _DEFAULT_VECTOR_STORE,
fs: Any = None,
) -> "TurboQuantVectorStore":
"""Load a store from a ``StorageContext``-style persist directory.
Builds the namespaced filename
``{persist_dir}/{namespace}__vector_store.json`` and forwards to
:meth:`from_persist_path`. The ``.json`` suffix is conventional —
our actual on-disk files use ``.tvim`` and ``.nodes.json``
extensions derived from the same stem.
"""
persist_fname = f"{namespace}{_NAMESPACE_SEP}{_DEFAULT_PERSIST_FNAME}"
persist_path = os.path.join(persist_dir, persist_fname)
return cls.from_persist_path(persist_path, fs=fs)
__all__ = ["TurboQuantVectorStore"]
turbovec-python/src/lib.rs
+406
View File
@@ -0,0 +1,406 @@
use numpy::{IntoPyArray, PyArray2, PyReadonlyArray1, PyReadonlyArray2};
use pyo3::prelude::*;
use pyo3::types::PyType;
fn not_contiguous_err(kind: &str) -> PyErr {
pyo3::exceptions::PyValueError::new_err(format!(
"{kind} must be C-contiguous; call np.ascontiguousarray(...) first",
))
}
/// 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,
}
#[pymethods]
impl TurboQuantIndex {
/// Construct an index. `dim` is optional: when omitted, the
/// underlying quantized index is created lazily on the first
/// `add` call, picking up the dimensionality from the input
/// array's shape.
#[new]
#[pyo3(signature = (dim=None, bit_width=4))]
fn new(dim: Option<usize>, bit_width: usize) -> PyResult<Self> {
let inner = match dim {
Some(d) => turbovec_core::TurboQuantIndex::new(d, bit_width),
None => turbovec_core::TurboQuantIndex::new_lazy(bit_width),
}
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
Ok(Self { inner })
}
fn add(&mut self, vectors: PyReadonlyArray2<f32>) -> PyResult<()> {
let arr = vectors.as_array();
let dim = arr.ncols();
let slice = arr.as_slice().ok_or_else(|| not_contiguous_err("vectors"))?;
// `add_2d` handles both eager (dim must match) and lazy (locks
// dim on first call) cases.
self.inner
.add_2d(slice, dim)
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}
/// Run a top-`k` search against the index.
///
/// `mask`, when given, is a bool array of length `len(self)`. Only slots
/// with `mask[i] == True` contribute to the returned top-`k`. The
/// returned result count per query is `min(k, mask.sum())`.
#[pyo3(signature = (queries, k, *, mask=None))]
fn search<'py>(
&self,
py: Python<'py>,
queries: PyReadonlyArray2<f32>,
k: usize,
mask: Option<PyReadonlyArray1<bool>>,
) -> PyResult<(Bound<'py, PyArray2<f32>>, Bound<'py, PyArray2<i64>>)> {
let arr = queries.as_array();
let nq = arr.nrows();
let q_slice = arr.as_slice().ok_or_else(|| not_contiguous_err("queries"))?;
// Reject wrong-dim queries cleanly. Previously the inner
// `assert_eq!(queries.len(), nq * dim)` would fire as a Rust
// panic and surface to Python as a PanicException, not the
// ValueError users expect for input-shape mismatch.
if let Some(idx_dim) = self.inner.dim_opt() {
if arr.ncols() != idx_dim {
return Err(pyo3::exceptions::PyValueError::new_err(format!(
"query dim {} does not match index dim {}",
arr.ncols(),
idx_dim,
)));
}
}
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() {
Some(m_arr) => {
let expected = self.inner.len();
if m_arr.len() != expected {
return Err(pyo3::exceptions::PyValueError::new_err(format!(
"mask length {} does not match index size {}",
m_arr.len(),
expected,
)));
}
Some(m_arr.as_slice().ok_or_else(|| not_contiguous_err("mask"))?)
}
None => None,
};
let results = self.inner.search_with_mask(q_slice, k, mask_slice);
let effective_k = results.k;
let scores = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), results.scores)
.map_err(shape_err)?
.into_pyarray(py);
let indices = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), results.indices)
.map_err(shape_err)?
.into_pyarray(py);
Ok((scores, indices))
}
fn write(&self, path: &str) -> PyResult<()> {
self.inner.write(path).map_err(|e| {
pyo3::exceptions::PyIOError::new_err(format!("{}", e))
})
}
#[classmethod]
fn load(_cls: &Bound<PyType>, path: &str) -> PyResult<Self> {
let inner = turbovec_core::TurboQuantIndex::load(path).map_err(|e| {
pyo3::exceptions::PyIOError::new_err(format!("{}", e))
})?;
Ok(Self { inner })
}
/// Warm up the search caches (rotation matrix, Lloyd-Max centroids,
/// SIMD-blocked code layout) so the first `search` call does not pay
/// the one-time initialisation cost.
fn prepare(&self) {
self.inner.prepare();
}
/// Remove the vector at `idx` in O(1) by swapping with the last vector.
///
/// The last vector moves into the deleted slot — order is not
/// preserved. Returns the old index of the moved vector; equals `idx`
/// when `idx` was already the last element.
///
/// Raises ``IndexError`` if ``idx`` is out of range.
fn swap_remove(&mut self, idx: usize) -> PyResult<usize> {
let len = self.inner.len();
if idx >= len {
return Err(pyo3::exceptions::PyIndexError::new_err(format!(
"index {idx} out of range for index of length {len}",
)));
}
Ok(self.inner.swap_remove(idx))
}
fn __len__(&self) -> usize {
self.inner.len()
}
fn __repr__(&self) -> String {
let dim = self
.inner
.dim_opt()
.map_or_else(|| "None".to_string(), |d| d.to_string());
format!(
"turbovec.TurboQuantIndex(dim={}, bit_width={}, n_vectors={})",
dim,
self.inner.bit_width(),
self.inner.len()
)
}
/// Vector dimensionality. Returns ``None`` when the index was
/// constructed lazily (no ``dim=``) and hasn't seen an add yet;
/// otherwise an ``int``.
#[getter]
fn dim(&self) -> Option<usize> {
self.inner.dim_opt()
}
#[getter]
fn bit_width(&self) -> usize {
self.inner.bit_width()
}
}
#[pyclass]
struct IdMapIndex {
inner: turbovec_core::IdMapIndex,
}
#[pymethods]
impl IdMapIndex {
/// Construct an id-mapped index. `dim` is optional: when omitted,
/// the underlying quantized index is created lazily on the first
/// `add_with_ids` call, picking up dim from the input array shape.
#[new]
#[pyo3(signature = (dim=None, bit_width=4))]
fn new(dim: Option<usize>, bit_width: usize) -> PyResult<Self> {
let inner = match dim {
Some(d) => turbovec_core::IdMapIndex::new(d, bit_width),
None => turbovec_core::IdMapIndex::new_lazy(bit_width),
}
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
Ok(Self { inner })
}
/// Add `n = vectors.shape[0]` vectors with the given external `ids`.
///
/// `ids` must be a 1-D array of `uint64` with length equal to
/// `vectors.shape[0]`. Raises `ValueError` if any id is already
/// present or if the lengths don't match. On a lazy index, this
/// call commits the dimensionality from `vectors.shape[1]`.
fn add_with_ids(
&mut self,
vectors: PyReadonlyArray2<f32>,
ids: PyReadonlyArray1<u64>,
) -> PyResult<()> {
let v = vectors.as_array();
let dim = v.ncols();
let v_slice = v.as_slice().ok_or_else(|| not_contiguous_err("vectors"))?;
let i = ids.as_array();
let i_slice = i.as_slice().ok_or_else(|| not_contiguous_err("ids"))?;
self.inner
.add_with_ids_2d(v_slice, dim, i_slice)
.map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}
/// Remove the vector with external id `id`. Returns `True` if it was
/// present, `False` otherwise.
fn remove(&mut self, id: u64) -> bool {
self.inner.remove(id)
}
/// Search for the top-`k` nearest external ids for each query.
///
/// `allowlist`, when given, is a `uint64` array of external ids; the
/// returned top-`k` is restricted to ids in this list. The returned
/// result count per query is `min(k, len(allowlist))` (after
/// de-duplication).
///
/// Returns `(scores, ids)` as `(nq, effective_k)` arrays, `ids` typed
/// `uint64`. Raises `ValueError` for an empty allowlist and `KeyError`
/// if any allowlist id is not present in the index.
#[pyo3(signature = (queries, k, *, allowlist=None))]
fn search<'py>(
&self,
py: Python<'py>,
queries: PyReadonlyArray2<f32>,
k: usize,
allowlist: Option<PyReadonlyArray1<u64>>,
) -> PyResult<(Bound<'py, PyArray2<f32>>, Bound<'py, PyArray2<u64>>)> {
let arr = queries.as_array();
let nq = arr.nrows();
let q_slice = arr.as_slice().ok_or_else(|| not_contiguous_err("queries"))?;
if let Some(idx_dim) = self.inner.dim_opt() {
if arr.ncols() != idx_dim {
return Err(pyo3::exceptions::PyValueError::new_err(format!(
"query dim {} does not match index dim {}",
arr.ncols(),
idx_dim,
)));
}
}
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() {
Some(a_arr) => {
if a_arr.is_empty() {
return Err(pyo3::exceptions::PyValueError::new_err(
"allowlist is empty",
));
}
let slice = a_arr.as_slice().ok_or_else(|| not_contiguous_err("allowlist"))?;
let mut unknown: Vec<u64> = Vec::new();
for &id in slice {
if !self.inner.contains(id) {
if unknown.len() < 5 {
unknown.push(id);
} else {
unknown.push(id);
break;
}
}
}
if !unknown.is_empty() {
let preview: Vec<u64> = unknown.iter().take(5).copied().collect();
return Err(pyo3::exceptions::PyKeyError::new_err(format!(
"allowlist contains id(s) not present in index: {:?}{}",
preview,
if unknown.len() > 5 { ", ..." } else { "" },
)));
}
Some(slice)
}
None => None,
};
let (scores, ids) = self.inner.search_with_allowlist(q_slice, k, allow_slice);
// For empty queries (nq=0), match TurboQuantIndex's shape
// contract: effective_k is `min(k, n_vectors, n_allowed)`. The
// kernel dedups the allowlist via a packed bool mask for nq>0,
// so we have to dedup here too — otherwise `allowlist=[1, 1, 1]`
// returns shape `(0, 3)` for empty queries but `(N, 1)` for
// non-empty queries, a silent shape divergence.
let effective_k = if nq == 0 {
let n_allowed = match allow_slice {
Some(s) => {
let mut seen: std::collections::HashSet<u64> =
std::collections::HashSet::with_capacity(s.len());
s.iter().filter(|id| seen.insert(**id)).count()
}
None => self.inner.len(),
};
k.min(self.inner.len()).min(n_allowed)
} else {
scores.len() / nq
};
let scores_arr = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), scores)
.map_err(shape_err)?
.into_pyarray(py);
let ids_arr = numpy::ndarray::Array2::from_shape_vec((nq, effective_k), ids)
.map_err(shape_err)?
.into_pyarray(py);
Ok((scores_arr, ids_arr))
}
fn contains(&self, id: u64) -> bool {
self.inner.contains(id)
}
fn prepare(&self) {
self.inner.prepare();
}
/// Serialize the index and id-map side-tables to a `.tvim` file.
fn write(&self, path: &str) -> PyResult<()> {
self.inner.write(path).map_err(|e| {
pyo3::exceptions::PyIOError::new_err(format!("{}", e))
})
}
/// Load an `IdMapIndex` from a `.tvim` file previously written by
/// [`IdMapIndex.write`].
#[classmethod]
fn load(_cls: &Bound<PyType>, path: &str) -> PyResult<Self> {
let inner = turbovec_core::IdMapIndex::load(path).map_err(|e| {
pyo3::exceptions::PyIOError::new_err(format!("{}", e))
})?;
Ok(Self { inner })
}
fn __len__(&self) -> usize {
self.inner.len()
}
fn __repr__(&self) -> String {
let dim = self
.inner
.dim_opt()
.map_or_else(|| "None".to_string(), |d| d.to_string());
format!(
"turbovec.IdMapIndex(dim={}, bit_width={}, n_vectors={})",
dim,
self.inner.bit_width(),
self.inner.len()
)
}
fn __contains__(&self, id: u64) -> bool {
self.inner.contains(id)
}
/// Vector dimensionality. Returns ``None`` when the index was
/// constructed lazily and hasn't seen an add yet; otherwise ``int``.
#[getter]
fn dim(&self) -> Option<usize> {
self.inner.dim_opt()
}
#[getter]
fn bit_width(&self) -> usize {
self.inner.bit_width()
}
}
#[pymodule]
fn _turbovec(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<TurboQuantIndex>()?;
m.add_class::<IdMapIndex>()?;
Ok(())
}
turbovec-python/tests/test_agno.py
+1203
View File
File diff suppressed because it is too large Load Diff
turbovec-python/tests/test_dedup.py
+39
View File
@@ -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]
turbovec-python/tests/test_filtering.py
+205
View File
@@ -0,0 +1,205 @@
"""Tests for the `mask=` and `allowlist=` filtering kwargs."""
from __future__ import annotations
import numpy as np
import pytest
from turbovec import IdMapIndex, TurboQuantIndex
DIM = 128
def unit_vectors(n: int, dim: int = DIM, seed: int = 0) -> np.ndarray:
rng = np.random.default_rng(seed)
v = rng.standard_normal((n, dim)).astype(np.float32)
v /= np.linalg.norm(v, axis=1, keepdims=True) + 1e-9
return v
# ------------------- TurboQuantIndex.search(mask=...) -------------------
def test_mask_none_matches_unmasked():
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(100, seed=1))
queries = unit_vectors(3, seed=2)
s1, i1 = idx.search(queries, 10)
s2, i2 = idx.search(queries, 10, mask=None)
np.testing.assert_array_equal(s1, s2)
np.testing.assert_array_equal(i1, i2)
def test_mask_all_true_matches_unmasked():
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(80, seed=3))
queries = unit_vectors(2, seed=4)
s1, i1 = idx.search(queries, 5)
mask = np.ones(len(idx), dtype=bool)
s2, i2 = idx.search(queries, 5, mask=mask)
np.testing.assert_array_equal(s1, s2)
np.testing.assert_array_equal(i1, i2)
def test_mask_restricts_returned_indices():
n = 200
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(n, seed=5))
queries = unit_vectors(4, seed=6)
mask = np.zeros(n, dtype=bool)
allowed = [3, 7, 19, 42, 88, 121, 150, 175, 198]
mask[allowed] = True
scores, indices = idx.search(queries, 5, mask=mask)
assert scores.shape == (4, 5)
assert indices.shape == (4, 5)
for row in indices:
for slot in row:
assert slot in allowed, f"kernel returned disallowed slot {slot}"
# Scores descending per row.
for row in scores:
assert np.all(np.diff(row) <= 1e-6)
def test_mask_shrinks_effective_k():
n = 100
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(n, seed=7))
queries = unit_vectors(2, seed=8)
mask = np.zeros(n, dtype=bool)
mask[[5, 10, 15]] = True
scores, indices = idx.search(queries, 10, mask=mask)
assert scores.shape == (2, 3), "effective_k should be popcount(mask) = 3"
assert indices.shape == (2, 3)
def test_mask_all_false_returns_empty_columns():
n = 50
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(n, seed=9))
queries = unit_vectors(2, seed=10)
mask = np.zeros(n, dtype=bool)
scores, indices = idx.search(queries, 5, mask=mask)
assert scores.shape == (2, 0)
assert indices.shape == (2, 0)
def test_mask_wrong_length_raises():
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(50, seed=11))
queries = unit_vectors(1, seed=12)
with pytest.raises(ValueError, match="mask length"):
idx.search(queries, 5, mask=np.ones(10, dtype=bool))
def test_mask_must_be_bool_dtype():
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(50, seed=13))
queries = unit_vectors(1, seed=14)
with pytest.raises((TypeError, ValueError)):
idx.search(queries, 5, mask=np.ones(50, dtype=np.uint8))
def test_mask_matches_post_hoc_filter():
n = 256
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(n, seed=15))
queries = unit_vectors(5, seed=16)
mask = np.zeros(n, dtype=bool)
mask[::3] = True
k = 7
unfiltered_scores, unfiltered_idx = idx.search(queries, n)
expected_indices = []
expected_scores = []
for qi in range(queries.shape[0]):
row_idx = unfiltered_idx[qi]
row_s = unfiltered_scores[qi]
keep = mask[row_idx]
expected_indices.append(row_idx[keep][:k])
expected_scores.append(row_s[keep][:k])
masked_scores, masked_idx = idx.search(queries, k, mask=mask)
assert masked_idx.shape == (queries.shape[0], k)
for qi in range(queries.shape[0]):
np.testing.assert_array_equal(masked_idx[qi], expected_indices[qi])
# Score parity is exact when both calls go through the same code
# path (single batched call on both sides).
np.testing.assert_allclose(masked_scores[qi], expected_scores[qi], rtol=1e-4, atol=1e-5)
# ------------------- IdMapIndex.search(allowlist=...) -------------------
def test_allowlist_none_matches_unfiltered():
idx = IdMapIndex(dim=DIM, bit_width=4)
ids = np.arange(7000, 7100, dtype=np.uint64)
idx.add_with_ids(unit_vectors(100, seed=20), ids)
queries = unit_vectors(2, seed=21)
s1, i1 = idx.search(queries, 10)
s2, i2 = idx.search(queries, 10, allowlist=None)
np.testing.assert_array_equal(s1, s2)
np.testing.assert_array_equal(i1, i2)
def test_allowlist_restricts_returned_ids():
idx = IdMapIndex(dim=DIM, bit_width=4)
ids = np.arange(1000, 1100, dtype=np.uint64)
idx.add_with_ids(unit_vectors(100, seed=22), ids)
queries = unit_vectors(3, seed=23)
allowed = np.array([1003, 1010, 1042, 1077, 1099], dtype=np.uint64)
scores, returned = idx.search(queries, 10, allowlist=allowed)
assert scores.shape == (3, len(allowed))
assert returned.shape == (3, len(allowed))
for row in returned:
for rid in row:
assert rid in allowed
def test_allowlist_empty_raises_value_error():
idx = IdMapIndex(dim=DIM, bit_width=4)
idx.add_with_ids(
unit_vectors(5, seed=24),
np.array([1, 2, 3, 4, 5], dtype=np.uint64),
)
queries = unit_vectors(1, seed=25)
with pytest.raises(ValueError, match="allowlist is empty"):
idx.search(queries, 3, allowlist=np.array([], dtype=np.uint64))
def test_allowlist_unknown_id_raises_key_error():
idx = IdMapIndex(dim=DIM, bit_width=4)
idx.add_with_ids(
unit_vectors(5, seed=26),
np.array([1, 2, 3, 4, 5], dtype=np.uint64),
)
queries = unit_vectors(1, seed=27)
with pytest.raises(KeyError, match="not present"):
idx.search(queries, 3, allowlist=np.array([2, 999], dtype=np.uint64))
def test_allowlist_kwarg_is_keyword_only():
idx = IdMapIndex(dim=DIM, bit_width=4)
ids = np.arange(100, 110, dtype=np.uint64)
idx.add_with_ids(unit_vectors(10, seed=28), ids)
queries = unit_vectors(1, seed=29)
# Positional third arg should be rejected by PyO3 signature.
with pytest.raises(TypeError):
idx.search(queries, 3, ids)
def test_mask_kwarg_is_keyword_only():
idx = TurboQuantIndex(dim=DIM, bit_width=4)
idx.add(unit_vectors(10, seed=30))
queries = unit_vectors(1, seed=31)
mask = np.ones(10, dtype=bool)
with pytest.raises(TypeError):
idx.search(queries, 3, mask)
turbovec-python/tests/test_haystack.py
+1200
View File
File diff suppressed because it is too large Load Diff
turbovec-python/tests/test_id_map.py
+211
View File
@@ -0,0 +1,211 @@
"""Tests for IdMapIndex — the stable-id wrapper around TurboQuantIndex."""
from __future__ import annotations
import numpy as np
import pytest
from turbovec import IdMapIndex
def unit_vectors(n: int, dim: int, seed: int = 0) -> np.ndarray:
rng = np.random.default_rng(seed)
v = rng.standard_normal((n, dim)).astype(np.float32)
v /= np.linalg.norm(v, axis=1, keepdims=True) + 1e-9
return v
def test_add_with_ids_updates_len_and_contains():
idx = IdMapIndex(dim=128, bit_width=4)
idx.add_with_ids(unit_vectors(5, 128), np.array([10, 20, 30, 40, 50], dtype=np.uint64))
assert len(idx) == 5
assert idx.contains(30)
assert not idx.contains(99)
# __contains__ sugar
assert 30 in idx
assert 99 not in idx
def test_search_returns_external_ids():
idx = IdMapIndex(dim=256, bit_width=4)
vectors = unit_vectors(10, 256, seed=0)
ids = np.arange(1_000_000, 1_000_010, dtype=np.uint64)
idx.add_with_ids(vectors, ids)
_, got = idx.search(vectors, k=1)
# Each self-query should return its own external id as top-1.
np.testing.assert_array_equal(got[:, 0], ids)
def test_remove_returns_true_false_correctly():
idx = IdMapIndex(dim=128, bit_width=4)
idx.add_with_ids(unit_vectors(3, 128), np.array([1, 2, 3], dtype=np.uint64))
assert idx.remove(2) is True
assert len(idx) == 2
assert idx.remove(2) is False # already gone
assert idx.remove(999) is False # never existed
def test_remove_then_re_add_same_id():
idx = IdMapIndex(dim=128, bit_width=4)
idx.add_with_ids(unit_vectors(5, 128), np.array([1, 2, 3, 4, 5], dtype=np.uint64))
assert idx.remove(3)
new_vec = unit_vectors(1, 128, seed=42)
idx.add_with_ids(new_vec, np.array([3], dtype=np.uint64))
assert 3 in idx
assert len(idx) == 5
def test_remaining_ids_self_query_after_removes():
dim = 256
idx = IdMapIndex(dim=dim, bit_width=4)
vectors = unit_vectors(15, dim, seed=0)
ids = np.array([i * 7 + 11 for i in range(15)], dtype=np.uint64)
idx.add_with_ids(vectors, ids)
# Remove a few positions.
removed_positions = [5, 14, 0]
for p in removed_positions:
assert idx.remove(int(ids[p]))
for i, id_val in enumerate(ids):
if i in removed_positions:
continue
_, got = idx.search(vectors[i:i + 1], k=1)
assert got[0, 0] == id_val, (
f"id {id_val} (row {i}) didn't self-query after removes"
)
def test_add_with_ids_rejects_duplicate_id():
idx = IdMapIndex(dim=128, bit_width=4)
idx.add_with_ids(unit_vectors(2, 128), np.array([1, 2], dtype=np.uint64))
# Second call includes id=2 which is already present.
with pytest.raises(BaseException): # pyo3 surfaces Rust panic as PanicException/RuntimeError
idx.add_with_ids(unit_vectors(1, 128, seed=1), np.array([2], dtype=np.uint64))
def test_write_and_load_round_trip(tmp_path):
idx = IdMapIndex(dim=256, bit_width=4)
vectors = unit_vectors(10, 256, seed=0)
ids = np.arange(5000, 5010, dtype=np.uint64)
idx.add_with_ids(vectors, ids)
idx.remove(5004)
idx.remove(5007)
path = tmp_path / "idx.tvim"
idx.write(str(path))
restored = IdMapIndex.load(str(path))
assert len(restored) == 8
assert 5000 in restored
assert 5004 not in restored
assert 5007 not in restored
for i, id_val in enumerate(ids):
if id_val in (5004, 5007):
continue
_, got = restored.search(vectors[i:i + 1], k=1)
assert got[0, 0] == id_val
def test_load_rejects_nonexistent_file():
with pytest.raises(IOError):
IdMapIndex.load("/nonexistent/path/does-not-exist.tvim")
@pytest.mark.parametrize("bad_bit_width", [0, 1, 5, 8])
def test_constructor_rejects_bad_bit_width(bad_bit_width):
with pytest.raises(ValueError, match="bit_width"):
IdMapIndex(dim=128, bit_width=bad_bit_width)
@pytest.mark.parametrize("bad_dim", [0, 1, 4, 7, 9])
def test_constructor_rejects_bad_dim(bad_dim):
with pytest.raises(ValueError, match="dim"):
IdMapIndex(dim=bad_dim, bit_width=4)
def test_search_on_empty_eager_index_returns_zero_effective_k():
idx = IdMapIndex(dim=128, bit_width=4)
q = unit_vectors(1, 128)
_, ids = idx.search(q, k=3)
assert ids.shape == (1, 0)
# ---- Wave 5: typed-exception hygiene + cross-class consistency ----
def test_search_empty_queries_returns_consistent_shape_across_index_types():
# Cross-class invariant: passing an empty queries array with k>n
# must produce the same shape on both index types. Previously,
# IdMapIndex returned (0, k) (raw k) while TurboQuantIndex returned
# (0, min(k, n_vectors)) — a silent divergence in result shape.
from turbovec import TurboQuantIndex
tq = TurboQuantIndex(dim=64, bit_width=4)
im = IdMapIndex(dim=64, bit_width=4)
tq.add(unit_vectors(3, 64))
im.add_with_ids(unit_vectors(3, 64), np.array([1, 2, 3], dtype=np.uint64))
empty_queries = np.empty((0, 64), dtype=np.float32)
tq_scores, tq_indices = tq.search(empty_queries, k=5)
im_scores, im_ids = im.search(empty_queries, k=5)
assert tq_scores.shape == im_scores.shape
assert tq_indices.shape == im_ids.shape
# effective_k should be min(k=5, n_vectors=3) = 3.
assert tq_scores.shape == (0, 3)
def test_search_query_dim_mismatch_raises_value_error():
idx = IdMapIndex(dim=128, bit_width=4)
idx.add_with_ids(unit_vectors(3, 128), np.array([1, 2, 3], dtype=np.uint64))
wrong = unit_vectors(1, 64)
with pytest.raises(ValueError, match="query dim"):
idx.search(wrong, k=1)
def test_add_with_ids_noncontiguous_vectors_raises_value_error():
idx = IdMapIndex(dim=128, bit_width=4)
full = unit_vectors(2, 256)
sliced = full[:, ::2]
assert not sliced.flags["C_CONTIGUOUS"]
with pytest.raises(ValueError, match="contiguous"):
idx.add_with_ids(sliced, np.array([1, 2], dtype=np.uint64))
def test_add_with_ids_rejects_nan_with_value_error():
idx = IdMapIndex(dim=64, bit_width=4)
data = unit_vectors(1, 64).copy()
data[0, 5] = np.nan
with pytest.raises(ValueError, match="invalid input value"):
idx.add_with_ids(data, np.array([1], dtype=np.uint64))
def test_search_empty_queries_dedups_allowlist_for_effective_k():
# Wave-6 fix for a bug introduced in wave-5: the `effective_k`
# computation for `nq == 0` counted the raw allowlist length, but
# the kernel dedups the allowlist via a packed-bool mask for
# `nq > 0`. So `allowlist=[1, 1, 1]` returned shape `(0, 3)` for
# empty queries but `(N, 1)` for non-empty — silent divergence.
idx = IdMapIndex(dim=64, bit_width=4)
idx.add_with_ids(
unit_vectors(3, 64),
np.array([10, 20, 30], dtype=np.uint64),
)
# Allowlist with three copies of the same id. Effective n_allowed
# is 1 (after dedup), not 3.
allowlist_with_dupes = np.array([10, 10, 10], dtype=np.uint64)
empty_queries = np.empty((0, 64), dtype=np.float32)
real_query = unit_vectors(1, 64)
_, empty_ids = idx.search(empty_queries, k=5, allowlist=allowlist_with_dupes)
_, real_ids = idx.search(real_query, k=5, allowlist=allowlist_with_dupes)
# Both should have effective_k = 1 (only one unique id in the
# allowlist), differing only in the leading dimension.
assert empty_ids.shape[1] == real_ids.shape[1]
assert empty_ids.shape == (0, 1)
assert real_ids.shape == (1, 1)
turbovec-python/tests/test_index.py
+272
View File
@@ -0,0 +1,272 @@
"""Tests for TurboQuantIndex — the pyo3 binding surface.
These exercise the public Python API exposed by the Rust extension
(add, search, save/load, prepare, len, dim, bit_width), independent of
the LangChain / LlamaIndex wrappers.
"""
from __future__ import annotations
import numpy as np
import pytest
from turbovec import TurboQuantIndex
def unit_vectors(n: int, dim: int, seed: int = 0) -> np.ndarray:
rng = np.random.default_rng(seed)
v = rng.standard_normal((n, dim)).astype(np.float32)
v /= np.linalg.norm(v, axis=1, keepdims=True) + 1e-9
return v
def test_new_reports_dim_and_bit_width():
idx = TurboQuantIndex(dim=128, bit_width=4)
assert idx.dim == 128
assert idx.bit_width == 4
assert len(idx) == 0
@pytest.mark.parametrize("bit_width", [2, 3, 4])
def test_bit_width_options(bit_width):
idx = TurboQuantIndex(dim=128, bit_width=bit_width)
assert idx.bit_width == bit_width
idx.add(unit_vectors(20, 128))
assert len(idx) == 20
def test_add_updates_length():
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(50, 128))
assert len(idx) == 50
def test_add_is_incremental():
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(20, 128, seed=1))
idx.add(unit_vectors(30, 128, seed=2))
assert len(idx) == 50
def test_search_shape():
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(100, 128))
scores, indices = idx.search(unit_vectors(5, 128, seed=99), k=10)
assert scores.shape == (5, 10)
assert indices.shape == (5, 10)
def test_search_single_query():
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(100, 128))
scores, indices = idx.search(unit_vectors(1, 128, seed=99), k=5)
assert scores.shape == (1, 5)
assert indices.shape == (1, 5)
def test_self_query_recall_at_1():
vectors = unit_vectors(100, 256, seed=42)
idx = TurboQuantIndex(dim=256, bit_width=4)
idx.add(vectors)
hits = 0
for i in range(20):
_, indices = idx.search(vectors[i:i + 1], k=1)
if indices[0, 0] == i:
hits += 1
assert hits == 20, f"recall@1 failed: {hits}/20"
def test_save_load_roundtrip(tmp_path):
vectors = unit_vectors(80, 128, seed=7)
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(vectors)
idx.prepare()
path = str(tmp_path / "idx.tv")
idx.write(path)
loaded = TurboQuantIndex.load(path)
assert len(loaded) == 80
assert loaded.dim == 128
assert loaded.bit_width == 4
q = unit_vectors(3, 128, seed=8)
s_orig, i_orig = idx.search(q, k=10)
s_load, i_load = loaded.search(q, k=10)
np.testing.assert_array_equal(i_orig, i_load)
np.testing.assert_allclose(s_orig, s_load, rtol=1e-5)
def test_prepare_is_idempotent():
idx = TurboQuantIndex(dim=64, bit_width=4)
idx.add(unit_vectors(20, 64))
idx.prepare()
idx.prepare()
assert len(idx) == 20
def test_batch_query_matches_individual():
idx = TurboQuantIndex(dim=256, bit_width=4)
vectors = unit_vectors(50, 256, seed=0)
idx.add(vectors)
queries = unit_vectors(5, 256, seed=99)
_, batch_indices = idx.search(queries, k=3)
for i in range(5):
_, single_indices = idx.search(queries[i:i + 1], k=3)
np.testing.assert_array_equal(
batch_indices[i:i + 1], single_indices
)
def test_noncontiguous_input_is_handled():
# A strided slice of a larger array should still work.
big = unit_vectors(100, 128)
strided = big[::2]
assert not strided.flags["C_CONTIGUOUS"]
idx = TurboQuantIndex(dim=128, bit_width=4)
# PyO3 layer asserts contiguity; caller is expected to convert.
# This test documents that behaviour: a contiguous copy works.
idx.add(np.ascontiguousarray(strided))
assert len(idx) == 50
def test_swap_remove_shrinks_length():
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(10, 128))
moved_from = idx.swap_remove(3)
assert moved_from == 9
assert len(idx) == 9
def test_swap_remove_last_is_no_swap():
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(5, 128))
assert idx.swap_remove(4) == 4
assert len(idx) == 4
def test_search_after_swap_remove_reflects_new_layout():
# Cache-invalidation regression: the vector that moves into the
# deleted slot must be findable immediately after the delete.
idx = TurboQuantIndex(dim=256, bit_width=4)
vectors = unit_vectors(20, 256, seed=0)
idx.add(vectors)
# Prime the cache with a self-query.
_, pre = idx.search(vectors[5:6], k=1)
assert pre[0, 0] == 5
# Delete slot 5 — the last vector (index 19) moves into slot 5.
idx.swap_remove(5)
assert len(idx) == 19
_, post = idx.search(vectors[19:20], k=1)
assert post[0, 0] == 5, "vector that moved into slot 5 not found there"
def test_add_with_mismatched_dim_raises_value_error():
idx = TurboQuantIndex(dim=128, bit_width=4)
with pytest.raises(ValueError, match="dim mismatch"):
idx.add(unit_vectors(3, 256))
@pytest.mark.parametrize("bad_bit_width", [0, 1, 5, 8])
def test_constructor_rejects_bad_bit_width(bad_bit_width):
with pytest.raises(ValueError, match="bit_width"):
TurboQuantIndex(dim=128, bit_width=bad_bit_width)
@pytest.mark.parametrize("bad_dim", [0, 1, 4, 7, 9])
def test_constructor_rejects_bad_dim(bad_dim):
with pytest.raises(ValueError, match="dim"):
TurboQuantIndex(dim=bad_dim, bit_width=4)
def test_search_on_empty_eager_index_returns_zero_effective_k():
idx = TurboQuantIndex(dim=128, bit_width=4)
q = unit_vectors(1, 128)
scores, indices = idx.search(q, k=3)
assert scores.shape == (1, 0)
assert indices.shape == (1, 0)
# ---- Wave 5: typed-exception hygiene at the binding layer ----
def test_add_noncontiguous_vectors_raises_value_error():
# A non-C-contiguous numpy view used to surface as a Rust panic via
# `as_slice().expect("vectors must be contiguous")`. The binding now
# converts that to a clean ValueError pointing at the caller-side
# fix (`np.ascontiguousarray`).
idx = TurboQuantIndex(dim=128, bit_width=4)
full = unit_vectors(4, 256) # 256 cols
sliced = full[:, ::2] # take every other column → non-contiguous, 128 cols
assert not sliced.flags["C_CONTIGUOUS"]
with pytest.raises(ValueError, match="contiguous"):
idx.add(sliced)
def test_search_query_dim_mismatch_raises_value_error():
# A query with the wrong number of columns previously panicked
# inside the core's `assert_eq!(queries.len(), nq * dim)` — pyo3
# surfaces that as a PanicException, not the ValueError users
# naturally try to catch. The binding now validates ncols against
# the index dim before forwarding.
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(5, 128))
wrong_dim_query = unit_vectors(1, 64)
with pytest.raises(ValueError, match="query dim"):
idx.search(wrong_dim_query, k=1)
def test_search_noncontiguous_query_raises_value_error():
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(3, 128))
full = unit_vectors(2, 256)
sliced = full[:, ::2]
assert not sliced.flags["C_CONTIGUOUS"]
with pytest.raises(ValueError, match="contiguous"):
idx.search(sliced, k=1)
def test_swap_remove_out_of_bounds_raises_index_error():
# Previously panicked in core's `assert!(idx < n_vectors)`. The
# binding now raises a clean IndexError before the inner call.
idx = TurboQuantIndex(dim=128, bit_width=4)
idx.add(unit_vectors(3, 128))
with pytest.raises(IndexError, match="out of range"):
idx.swap_remove(99)
# ---- Wave 5: input validation surfaces from the Rust core ----
def test_add_rejects_nan_with_value_error():
# The Rust core's `AddError::InvalidInputValue` is mapped through
# the binding's `add_2d` Result → PyValueError path.
idx = TurboQuantIndex(dim=64, bit_width=4)
data = unit_vectors(1, 64).copy()
data[0, 5] = np.nan
with pytest.raises(ValueError, match="invalid input value"):
idx.add(data)
def test_add_rejects_huge_magnitude_with_value_error():
idx = TurboQuantIndex(dim=64, bit_width=4)
data = unit_vectors(1, 64).copy()
data[0, 5] = 1e20
with pytest.raises(ValueError, match="invalid input value"):
idx.add(data)
def test_search_with_nan_query_raises():
# NaN in the query reaches the core via search_with_mask, which
# now panics with a clear message; pyo3 maps Rust panics to
# PanicException → BaseException, so the test catches broadly but
# asserts on the message substring.
idx = TurboQuantIndex(dim=64, bit_width=4)
idx.add(unit_vectors(3, 64))
q = unit_vectors(1, 64).copy()
q[0, 0] = np.nan
with pytest.raises(BaseException, match="invalid query value"):
idx.search(q, k=1)
turbovec-python/tests/test_langchain.py
+783
View File
@@ -0,0 +1,783 @@
from __future__ import annotations
import numpy as np
import pytest
pytest.importorskip("langchain_core")
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from turbovec import IdMapIndex
from turbovec.langchain import TurboQuantVectorStore
class StubEmbeddings(Embeddings):
"""Deterministic text->vector function for tests.
Hashes the input string to seed an RNG, producing a reproducible
unit-norm vector. Similar strings do not map to similar vectors —
that's fine for structural tests, and callers shouldn't rely on
semantic ordering here.
"""
def __init__(self, dim: int = 64) -> None:
self.dim = dim
def _embed(self, text: str) -> list[float]:
rng = np.random.default_rng(abs(hash(text)) % (2**32))
v = rng.standard_normal(self.dim).astype(np.float32)
v /= np.linalg.norm(v) + 1e-9
return v.tolist()
def embed_documents(self, texts):
return [self._embed(t) for t in texts]
def embed_query(self, text):
return self._embed(text)
def test_from_texts_infers_dim_and_indexes():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["apple", "banana", "cherry", "date"], emb, bit_width=4
)
assert len(store._str_to_u64) == 4
assert store._index.dim == 64
assert store._index.bit_width == 4
def test_similarity_search_returns_documents():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["a", "b", "c"], emb, bit_width=4)
results = store.similarity_search("a", k=2)
assert len(results) == 2
assert all(isinstance(r, Document) for r in results)
def test_similarity_search_results_carry_document_id():
# Returned Documents must have `.id` populated — both the explicit ids
# callers pass to `add_texts` and the UUIDs the store generates by
# default. Matches the InMemoryVectorStore reference behaviour and what
# downstream LangChain callers (retrievers, chains) expect.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c"], emb, bit_width=4, ids=["id-a", "id-b", "id-c"]
)
results = store.similarity_search("a", k=3)
assert {r.id for r in results} == {"id-a", "id-b", "id-c"}
scored = store.similarity_search_with_score("a", k=3)
assert {doc.id for doc, _ in scored} == {"id-a", "id-b", "id-c"}
by_vec = store.similarity_search_by_vector(emb._embed("a"), k=3)
assert {r.id for r in by_vec} == {"id-a", "id-b", "id-c"}
def test_similarity_search_callable_filter_receives_document_id():
# Predicate Document must carry `.id` so callers can filter on it,
# not just on page_content / metadata.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c"], emb, bit_width=4, ids=["keep-1", "drop", "keep-2"]
)
results = store.similarity_search(
"a", k=10, filter=lambda doc: doc.id.startswith("keep")
)
assert {r.id for r in results} == {"keep-1", "keep-2"}
# ---- Reference-parity tests against InMemoryVectorStore. Each pins a
# behaviour the langchain-core in-tree suite tests; the bug class is
# "drop-in regression that only shows up when users compare against
# InMemoryVectorStore". ----
def test_async_methods_await_aembed_functions():
# If our async paths silently fall back to sync embedding, callers
# with a real async embedder (e.g. OpenAI async client) end up
# blocking the event loop. AsyncMock makes the side-effect visible:
# if aembed_documents / aembed_query weren't awaited, await_count
# stays zero.
import asyncio
from unittest.mock import AsyncMock
class AsyncStub(StubEmbeddings):
def __init__(self, dim: int = 64) -> None:
super().__init__(dim)
self.aembed_documents = AsyncMock(
side_effect=lambda texts: [self._embed(t) for t in texts]
)
self.aembed_query = AsyncMock(
side_effect=lambda text: self._embed(text)
)
emb = AsyncStub(dim=64)
store = TurboQuantVectorStore(embedding=emb)
async def run() -> None:
await store.aadd_texts(["a", "b"])
await store.asimilarity_search("a", k=1)
asyncio.run(run())
assert emb.aembed_documents.await_count >= 1
assert emb.aembed_query.await_count >= 1
def test_add_documents_upsert_replaces_metadata():
# Re-adding a Document with the same id and new metadata must let
# the new metadata win — not silently retain the old one. Reference
# `InMemoryVectorStore.add_documents` overwrites on duplicate id.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts([], emb, bit_width=4)
store.add_documents([Document(id="x", page_content="v1", metadata={"tag": "old"})])
store.add_documents([Document(id="x", page_content="v2", metadata={"tag": "new"})])
[doc] = store.get_by_ids(["x"])
assert doc.metadata == {"tag": "new"}
assert doc.page_content == "v2"
def test_add_documents_does_not_mutate_inputs():
# Caller-passed Documents must not be mutated by add_documents — no
# in-place .id assignment, no metadata mutation. Reference behaviour;
# easy regression target when refactoring the add path.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts([], emb, bit_width=4)
docs = [
Document(page_content="a", metadata={"k": 1}),
Document(id="explicit", page_content="b", metadata={"k": 2}),
]
original_meta_ids = [id(d.metadata) for d in docs]
store.add_documents(docs)
assert docs[0].id is None
assert docs[1].id == "explicit"
assert docs[0].metadata == {"k": 1}
assert docs[1].metadata == {"k": 2}
# Same dict objects — we didn't replace caller-provided metadata dicts.
assert [id(d.metadata) for d in docs] == original_meta_ids
def test_add_documents_with_ids_is_idempotent():
# Re-running ingestion on an unchanged corpus must not accrete
# duplicates. Reference upserts on same id; size stays constant.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts([], emb, bit_width=4)
docs = [
Document(id="a", page_content="hello"),
Document(id="b", page_content="world"),
]
store.add_documents(docs)
store.add_documents(docs)
assert len(store._docs) == 2
assert set(store._docs.keys()) == {"a", "b"}
def test_add_texts_intra_batch_duplicate_ids_keep_last():
# Two rows sharing an id in a single call must not orphan a vector.
# Reference InMemoryVectorStore overwrites on a repeated id (last wins);
# we dedup the batch the same way so the index holds one vector per id.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts([], emb, bit_width=4)
ret = store.add_texts(["alpha", "beta"], ids=["dup", "dup"])
# Return value still mirrors the input (one entry per input text).
assert ret == ["dup", "dup"]
# No orphaned vector: index and id maps agree at one entry.
assert len(store._index) == 1
assert len(store._u64_to_str) == 1
assert set(store._str_to_u64) == {"dup"}
# Last occurrence wins.
assert store._docs["dup"][0] == "beta"
def test_upsert_dim_mismatch_preserves_existing_data():
# An upsert whose new embeddings fail validation must not destroy the
# existing entry: the delete is deferred until the add succeeds.
class VarDimEmbeddings(Embeddings):
def __init__(self):
self.call = 0
def embed_documents(self, texts):
self.call += 1
dim = 64 if self.call == 1 else 32
return [[float(i + 1)] * dim for i in range(len(texts))]
def embed_query(self, t):
return [1.0] * 64
store = TurboQuantVectorStore(VarDimEmbeddings())
store.add_texts(["hello"], ids=["my-id"]) # dim=64, OK
with pytest.raises(ValueError):
store.add_texts(["world"], ids=["my-id"]) # dim=32, must reject
# Original data survives the failed upsert.
assert "my-id" in store._docs
assert store._docs["my-id"][0] == "hello"
assert len(store._index) == 1
assert len(store._u64_to_str) == 1
def test_get_by_ids_empty_input_and_order_preserved():
# Two contract points the reference makes that our existing tests
# don't pin: (1) empty input returns [] without erroring; (2) output
# is in the order of input ids so callers can zip with parallel
# arrays (e.g. scores from a separate retriever).
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c"], emb, bit_width=4, ids=["id-a", "id-b", "id-c"]
)
assert store.get_by_ids([]) == []
docs = store.get_by_ids(["id-c", "id-a", "id-b"])
assert [d.id for d in docs] == ["id-c", "id-a", "id-b"]
def test_similarity_search_with_dict_filter():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["alpha", "beta", "gamma", "delta", "epsilon"],
emb,
metadatas=[
{"tier": "free"},
{"tier": "pro"},
{"tier": "free"},
{"tier": "pro"},
{"tier": "pro"},
],
bit_width=4,
)
results = store.similarity_search("alpha", k=10, filter={"tier": "pro"})
assert len(results) == 3
assert all(r.metadata["tier"] == "pro" for r in results)
def test_similarity_search_with_callable_filter():
# Predicate receives a langchain_core Document (matching the in-tree
# InMemoryVectorStore convention), not a bare metadata dict.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c", "d"],
emb,
metadatas=[{"n": 1}, {"n": 2}, {"n": 3}, {"n": 4}],
bit_width=4,
)
results = store.similarity_search(
"a", k=10, filter=lambda doc: doc.metadata.get("n", 0) > 2
)
assert {r.metadata["n"] for r in results} == {3, 4}
def test_similarity_search_callable_filter_can_use_page_content():
# Document is passed to the predicate, so page_content is reachable.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["alpha", "beta", "alphabet"], emb, bit_width=4,
)
results = store.similarity_search(
"alpha", k=10, filter=lambda doc: doc.page_content.startswith("alpha")
)
contents = {r.page_content for r in results}
assert contents == {"alpha", "alphabet"}
def test_similarity_search_filter_with_scores():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c"],
emb,
metadatas=[{"k": 1}, {"k": 2}, {"k": 1}],
bit_width=4,
)
results = store.similarity_search_with_score("a", k=10, filter={"k": 1})
assert len(results) == 2
for doc, score in results:
assert doc.metadata["k"] == 1
assert isinstance(score, float)
def test_similarity_search_filter_no_matches_returns_empty():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b"], emb, metadatas=[{"k": 1}, {"k": 2}], bit_width=4
)
assert store.similarity_search("a", k=5, filter={"k": 999}) == []
def test_similarity_search_filter_invalid_type_raises():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["a"], emb, bit_width=4)
with pytest.raises(TypeError):
store.similarity_search("a", k=1, filter=42)
def test_metadata_roundtrip():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["hello", "world"],
emb,
metadatas=[{"source": "a"}, {"source": "b"}],
bit_width=4,
)
scored = store.similarity_search_with_score("hello", k=2)
assert len(scored) == 2
sources = {doc.metadata["source"] for doc, _ in scored}
assert sources == {"a", "b"}
def test_add_texts_uses_provided_ids():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts([], emb, bit_width=4)
returned = store.add_texts(["x", "y"], ids=["id-x", "id-y"])
assert returned == ["id-x", "id-y"]
assert set(store._docs.keys()) == {"id-x", "id-y"}
def test_k_larger_than_ntotal_is_clamped():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["one", "two"], emb, bit_width=4)
results = store.similarity_search("one", k=100)
assert len(results) == 2
def test_empty_store_search_returns_empty():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts([], emb, bit_width=4)
assert store.similarity_search("anything", k=5) == []
def test_dump_and_load_roundtrip(tmp_path):
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["one", "two", "three"],
emb,
metadatas=[{"n": 1}, {"n": 2}, {"n": 3}],
bit_width=4,
)
store.dump(tmp_path)
loaded = TurboQuantVectorStore.load(tmp_path, emb)
assert len(loaded._docs) == 3
results = loaded.similarity_search("one", k=3)
assert {doc.page_content for doc in results} == {"one", "two", "three"}
def test_dump_writes_json_sidecar(tmp_path):
# Side-car is plain JSON. A reviewer auditing a turbovec-saved store
# should be able to read it with a text editor.
import json
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["x"], emb, bit_width=4)
store.dump(tmp_path)
assert (tmp_path / "docstore.json").exists()
assert not (tmp_path / "docstore.pkl").exists()
with open(tmp_path / "docstore.json") as f:
data = json.load(f)
assert data["schema_version"] >= 1
def test_load_rejects_unknown_schema_version(tmp_path):
import json
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["x"], emb, bit_width=4)
store.dump(tmp_path)
with open(tmp_path / "docstore.json") as f:
data = json.load(f)
data["schema_version"] = 99
with open(tmp_path / "docstore.json", "w") as f:
json.dump(data, f)
with pytest.raises(ValueError, match="schema version"):
TurboQuantVectorStore.load(tmp_path, emb)
def test_delete_removes_documents_returns_none():
# Match InMemoryVectorStore convention: delete returns None.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["apple", "banana", "cherry"],
emb,
ids=["a", "b", "c"],
bit_width=4,
)
result = store.delete(["b"])
assert result is None
assert set(store._docs.keys()) == {"a", "c"}
assert len(store._index) == 2
def test_delete_missing_ids_silently_skips():
# Match InMemoryVectorStore convention: missing ids are silently
# skipped, no error.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b"], emb, ids=["id-a", "id-b"], bit_width=4
)
assert store.delete(["id-a", "ghost"]) is None
assert "id-a" not in store._docs
assert "id-b" in store._docs
def test_delete_none_ids_is_noop():
# InMemoryVectorStore treats `delete(None)` as a no-op rather than
# raising. Match that.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["x"], emb, bit_width=4)
assert store.delete(None) is None
assert "x" not in store._docs # uuid-based ids, but the store has 1 doc
assert len(store._docs) == 1
def test_add_texts_upsert_replaces_existing_id():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["v1"], emb, ids=["same-id"], bit_width=4
)
# Re-add with the same id but different text.
store.add_texts(["v2"], ids=["same-id"])
assert len(store._docs) == 1
assert store._docs["same-id"][0] == "v2"
def test_mismatched_dim_raises():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore(emb, index=IdMapIndex(32, 4))
with pytest.raises(ValueError, match="embedding dimension"):
store.add_texts(["hi"])
# ---- Lazy index construction (Tier 4) -------------------------------------
def test_constructor_no_index_is_lazy():
# Without an `index`, the underlying IdMapIndex is constructed in its
# lazy-uncommitted state — `dim` is None until the first add.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore(emb)
assert store._index.dim is None
# Search before any add returns empty rather than raising.
assert store.similarity_search("anything", k=3) == []
def test_lazy_index_dim_locked_on_first_add():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore(emb, bit_width=2)
store.add_texts(["hello"])
assert store._index.dim == 64
assert store._index.bit_width == 2
def test_from_texts_no_dim_arg_required():
# Tier 4: dim is inferred from the embedding model, no explicit param.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["one", "two"], emb, bit_width=4
)
assert store._index.dim == 64
# ---- get_by_ids -----------------------------------------------------------
def test_get_by_ids_returns_documents():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c"], emb,
metadatas=[{"n": 1}, {"n": 2}, {"n": 3}],
ids=["id-a", "id-b", "id-c"],
bit_width=4,
)
docs = store.get_by_ids(["id-a", "id-c"])
assert {d.id for d in docs} == {"id-a", "id-c"}
assert {d.metadata["n"] for d in docs} == {1, 3}
def test_get_by_ids_silently_skips_missing():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a"], emb, ids=["id-a"], bit_width=4
)
docs = store.get_by_ids(["id-a", "id-missing"])
assert len(docs) == 1
assert docs[0].id == "id-a"
# ---- Relevance score normalization ----------------------------------------
def test_select_relevance_score_fn_maps_to_unit_interval():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["hello"], emb, bit_width=4)
fn = store._select_relevance_score_fn()
# Cosine similarity in [-1, 1] → relevance in [0, 1].
assert fn(-1.0) == 0.0
assert fn(0.0) == 0.5
assert fn(1.0) == 1.0
def test_similarity_search_with_relevance_scores_in_zero_one():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["one", "two", "three"], emb, bit_width=4
)
results = store.similarity_search_with_relevance_scores("one", k=3)
assert len(results) == 3
for _doc, score in results:
assert 0.0 <= score <= 1.0
# ---- MMR raises with explanation ------------------------------------------
def test_max_marginal_relevance_search_raises_with_message():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["a", "b"], emb, bit_width=4)
with pytest.raises(NotImplementedError, match="full-precision"):
store.max_marginal_relevance_search("a", k=2)
def test_max_marginal_relevance_search_by_vector_raises():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(["a", "b"], emb, bit_width=4)
with pytest.raises(NotImplementedError, match="full-precision"):
store.max_marginal_relevance_search_by_vector(emb._embed("a"), k=2)
# ---- add_documents partial-id support ------------------------------------
def test_add_documents_honors_partial_ids():
# Tier 3: per-Document fallback — Documents with .id set keep their
# id, others get a UUID. Base-class default (without override) would
# drop all ids if any Document had .id=None.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore(emb)
docs = [
Document(id="explicit-1", page_content="a"),
Document(page_content="b"), # no id → UUID
Document(id="explicit-2", page_content="c"),
]
returned_ids = store.add_documents(docs)
assert "explicit-1" in returned_ids
assert "explicit-2" in returned_ids
# The UUID-generated id is some non-explicit string.
uuid_id = [i for i in returned_ids if i not in ("explicit-1", "explicit-2")]
assert len(uuid_id) == 1
# ---- Async surfaces -------------------------------------------------------
def test_async_add_search_delete():
import asyncio
async def runner():
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore(emb)
ids = await store.aadd_texts(["alpha", "beta", "gamma"])
assert len(ids) == 3
results = await store.asimilarity_search("alpha", k=2)
assert len(results) == 2
scored = await store.asimilarity_search_with_score("alpha", k=2)
assert len(scored) == 2 and isinstance(scored[0][1], float)
by_vec = await store.asimilarity_search_by_vector(emb._embed("alpha"), k=1)
assert len(by_vec) == 1
got = await store.aget_by_ids(ids)
assert len(got) == 3
await store.adelete([ids[0]])
assert ids[0] not in store._docs
asyncio.run(runner())
def test_async_add_documents_and_afrom_texts():
import asyncio
async def runner():
emb = StubEmbeddings(dim=64)
store = await TurboQuantVectorStore.afrom_texts(
["x", "y"], emb, bit_width=4
)
assert len(store._docs) == 2
await store.aadd_documents([Document(page_content="z")])
assert len(store._docs) == 3
asyncio.run(runner())
def test_async_mmr_raises():
import asyncio
async def runner():
emb = StubEmbeddings(dim=64)
store = await TurboQuantVectorStore.afrom_texts(["x"], emb, bit_width=4)
with pytest.raises(NotImplementedError, match="full-precision"):
await store.amax_marginal_relevance_search("x", k=1)
asyncio.run(runner())
# ---- Empty-store persistence round-trip (lazy index) ---------------------
# ---- End-to-end smoke tests: framework wiring ---------------------------
def test_as_retriever_invoke_returns_documents():
# Smoke test: wire the store into LangChain's VectorStoreRetriever via
# `as_retriever()` and run a query through the .invoke() interface.
# This is the canonical way users plug a VectorStore into a Chain,
# so it exercises the base-class wiring that calls similarity_search
# on our store from the framework side.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["alpha", "beta", "gamma", "delta"],
emb,
metadatas=[{"tag": "a"}, {"tag": "b"}, {"tag": "a"}, {"tag": "b"}],
bit_width=4,
)
retriever = store.as_retriever(search_kwargs={"k": 2})
docs = retriever.invoke("alpha")
assert len(docs) == 2
assert all(isinstance(d, Document) for d in docs)
def test_as_retriever_with_filter_kwarg():
# The retriever passes search_kwargs (including `filter`) through to
# similarity_search. This verifies the keyword reaches our store
# without being dropped by the base class.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["alpha", "beta", "gamma"],
emb,
metadatas=[{"tag": "keep"}, {"tag": "drop"}, {"tag": "keep"}],
bit_width=4,
)
retriever = store.as_retriever(
search_kwargs={"k": 5, "filter": {"tag": "keep"}}
)
docs = retriever.invoke("alpha")
assert len(docs) == 2
assert all(d.metadata["tag"] == "keep" for d in docs)
def test_as_retriever_similarity_score_threshold():
# `similarity_score_threshold` is the search_type that uses
# similarity_search_with_relevance_scores under the hood, which
# depends on our _select_relevance_score_fn override. If that's
# missing or broken, this test fails with NotImplementedError.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["alpha", "beta", "gamma"], emb, bit_width=4
)
retriever = store.as_retriever(
search_type="similarity_score_threshold",
search_kwargs={"k": 3, "score_threshold": 0.0},
)
docs = retriever.invoke("alpha")
# All scores should be >= threshold (relevance in [0, 1] >= 0).
assert len(docs) >= 1
def test_dump_and_load_empty_store(tmp_path):
# When no documents have been added the underlying IdMapIndex is in
# its lazy-uncommitted state (dim=None). dump/load must round-trip
# that without losing the bit_width or accidentally committing a dim.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore(emb, bit_width=2)
store.dump(tmp_path)
loaded = TurboQuantVectorStore.load(tmp_path, emb)
assert loaded._index.dim is None
assert loaded._index.bit_width == 2
# Subsequent search returns empty; subsequent add commits the dim.
assert loaded.similarity_search("anything", k=1) == []
loaded.add_texts(["new"])
assert loaded._index.dim == 64
# ---- Tier-2 field-completeness tests. Each pins a value that a future
# refactor could silently drop (the #81 family: populated but
# unasserted). ----
def test_similarity_search_with_score_returns_descending_scores_and_self_match():
# Pins the actual semantics of the float in (Document, float) — that
# it's a real similarity score, that the self-match wins, and that
# results are monotonically non-increasing. Without this, the tuple
# position 1 could silently regress to a constant or get swapped
# with `handle` and only `isinstance(score, float)` would still pass.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["alpha", "beta", "gamma"], emb, bit_width=4
)
scored = store.similarity_search_with_score("alpha", k=3)
assert scored[0][0].page_content == "alpha"
scores = [s for _, s in scored]
assert all(a >= b for a, b in zip(scores, scores[1:]))
def test_load_then_add_assigns_fresh_handles_without_collision(tmp_path):
# `_next_u64` is persisted across dump/load; if it were silently
# dropped (back to 0 on load), a fresh add would reuse a handle
# that's still mapped to an old doc, corrupting search results.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c"], emb, bit_width=4, ids=["id-a", "id-b", "id-c"]
)
store.dump(tmp_path)
loaded = TurboQuantVectorStore.load(tmp_path, emb)
loaded.add_texts(["d"], ids=["id-d"])
# All four ids reachable; all four handles distinct.
docs = loaded.get_by_ids(["id-a", "id-b", "id-c", "id-d"])
assert [d.id for d in docs] == ["id-a", "id-b", "id-c", "id-d"]
handles = list(loaded._str_to_u64.values())
assert len(set(handles)) == len(handles)
def test_embeddings_property_returns_supplied_embedder():
# Pins the `embeddings` property override. A refactor dropping it
# would silently make the base class return None, breaking
# `similarity_search_with_relevance_scores` discovery and some
# retriever wiring.
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore(emb)
assert store.embeddings is emb
def test_aget_by_ids_preserves_order_and_returns_documents_with_id():
# Async mirror of `test_get_by_ids_empty_input_and_order_preserved`.
import asyncio
emb = StubEmbeddings(dim=64)
store = TurboQuantVectorStore.from_texts(
["a", "b", "c"], emb, bit_width=4, ids=["id-a", "id-b", "id-c"]
)
async def run() -> list[Document]:
empty = await store.aget_by_ids([])
assert empty == []
return await store.aget_by_ids(["id-c", "id-a", "id-b"])
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)

Some files were not shown because too many files have changed in this diff Show More