llama-index-vector-stores-rabitqlib

A LlamaIndex vector store backed by RaBitQ-quantized HNSW (rabitqlib.HnswIndex). It's a local, in-process index, a drop-in for SimpleVectorStore or FaissVectorStore, and stores compressed vectors (default 7 bits per dimension).

Status: v0.1, not yet published to PyPI.

Install

pip install -e .   # from this folder; pulls llama-index-core and rabitqlib>=0.8.3

Quickstart

from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage
from llama_index.vector_stores.rabitqlib import RabitqHnswVectorStore

vector_store = RabitqHnswVectorStore(dim=1536)  # your embedding model's dimension
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

nodes = index.as_retriever(similarity_top_k=5).retrieve("my question")
index.delete_ref_doc("some-doc-id", delete_from_docstore=True)

# Persist and reload. The vector store has to be loaded explicitly, because
# StorageContext.from_defaults(persist_dir=...) alone would load SimpleVectorStore.
index.storage_context.persist(persist_dir="./storage")
vector_store = RabitqHnswVectorStore.from_persist_dir("./storage")
storage_context = StorageContext.from_defaults(
    vector_store=vector_store, persist_dir="./storage"
)
index = load_index_from_storage(storage_context)

examples/quickstart.py is a runnable offline version that uses a toy embedding and no LLM. It passes min_train_size=1 so its three documents go into the HNSW index. With the default of 1024, a store that small is answered by an exact scan of the buffer and never reaches RaBitQ (see "Training is deferred" below).

Options

How it works

• Text lives in LlamaIndex's docstore (stores_text = False), as with FAISS. The store keeps only embeddings plus a node-ID ↔ index-ID map and a ref-doc ↔ node map.
• Training is deferred. RaBitQ trains its centroids once, on the first vectors it gets, and can't retrain afterwards. Training on a handful of vectors would hurt every later query. So until min_train_size embeddings have arrived they're kept in a float32 buffer and searched exactly. When the threshold is reached, the whole buffer is added to a new HnswIndex in one call, which trains on it. Pass train_vectors= or centroids= to skip the buffer.
• Adds are batched. Each add(nodes) makes one HnswIndex.add call. Re-adding an existing node ID replaces its embedding.
• Scores are higher-is-better, computed from RaBitQ distance estimates. They're close to exact at nbits=7, but not exact:
  • cosine: cosine similarity, clipped to [-1, 1];
  • ip: inner product (unbounded);
  • l2: 1 / (1 + squared_l2), in (0, 1].
• Persistence. persist(path) writes JSON metadata to path (for example default__vector_store.json). Alongside it go path.hnsw (the index), path.buffer.npy (vectors still buffered) or path.train.npy (training data not yet used). The JSON has a format_version. Only the local filesystem is supported.
• Thread safety. Queries and persist run concurrently with each other, though two persist calls run one at a time. Add, delete and clear are exclusive. Waiting writers don't block new queries, so a continuous stream of overlapping queries can hold an add back indefinitely. That's unlikely with RAG-style traffic.

Limitations (v0.1)

• No metadata filters. A query with filters, or delete_nodes(filters=...), raises ValueError, and so does a query restricted to doc_ids. Only the default query mode is supported, so no MMR, hybrid or sparse search.
• node_ids on a query is ignored, as in FaissVectorStore, so index.as_retriever(node_ids=subset) searches every node. as_retriever() always passes the node IDs it knew about when it was created. Honouring that list would break long-lived retrievers after an insert, and would cost O(n) per query.
• Deletes don't reclaim memory. Deleted nodes are excluded from results but stay in storage. To reclaim it, rebuild from your documents.
• Accuracy is bounded by quantization. Recall can plateau below 1.0 however large ef is (see below). Raise nbits for higher accuracy.
• No compaction or retraining. Centroids are fixed after the first training. If your data drifts a lot, rebuild.

Measured recall

Run with examples/recall.py. Setup: 10,000 base and 200 query vectors, dim 128, drawn from 50 Gaussian clusters (seed 0); default M=16, ef_construction=200, num_clusters=16; graph built with num_threads=1, so runs are reproducible; top-10 compared against an exact NumPy scan. Machine: 96-core Linux, rabitqlib 0.8.3, llama-index-core 0.14.25. With multi-threaded builds, expect differences of about ±0.002.

On this dataset, recall is limited by the quantized distance estimates, not by graph search: calling rabitqlib.HnswIndex directly gives the same numbers, and they level off from about ef=64. HnswIndex has no exact re-ranking step, so raising ef can't push recall past this ceiling; only a higher nbits can. This is one synthetic dataset, so measure on your own embeddings.

Development

python -m pytest --cov=llama_index.vector_stores.rabitqlib
ruff check . && ruff format --check .
mypy llama_index tests

License

Apache-2.0, the same as RaBitQ-Library.