Skip to content

feat: add RaBitQ navigation to DiskANN - #1839

Draft
CLiqing wants to merge 11 commits into
zilliztech:mainfrom
CLiqing:feat/diskann-rabitq-navigation
Draft

CLiqing wants to merge 11 commits into
zilliztech:mainfrom
CLiqing:feat/diskann-rabitq-navigation

Conversation

@CLiqing

@CLiqing CLiqing commented Sep 22, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Add RaBitQ as an optional in-memory navigation codec for DiskANN, reusing the Faiss RaBitQ implementation already integrated by #1814. PQ remains the default. The disk graph, traversal, candidate queues, and SSD I/O paths remain shared.

issue: #1747.

Functionality

  • Select navigation_codec="RABITQ" on DISKANN, or use the DISKANN_RABITQ alias. RaBitQ supports FP32 input with L2, IP, and COSINE metrics.
  • Reuse Faiss encoding, random rotation, serialization, and distance kernels. Database bits are configurable from 1 to 9; search-time rbq_bits_query supports 0 to 8 (default 4).
  • Use probabilistic one-bit screening before multi-bit evaluation; this is approximate pruning, not a guaranteed distance bound. Coarse and full scoring support batch-4 evaluation. No separate full-evaluation mode or public cache-seed parameter is exposed.
  • Introduce a navigation store and query-local distance-computer interface without putting codec-specific parameters into the graph traversal interface. Preserve the existing PQ fast path.
  • Keep navigation encoding independent of SSD payload encoding: RaBitQ navigation works with either uncompressed floating-point SSD vectors or existing SSD PQ. RaBitQ does not require an unused navigation PQ budget, codebook, codes, or scratch buffers.
  • Support filtered TopK, warmup, BFS node caching, sidecar validation, and failure-safe loading. Include Python factory/roundtrip support.
  • Discover persisted PQ/RaBitQ navigation when loading through DISKANN, without repeating the build-time codec or bits. Explicit codec/alias constraints must match; ambiguous, missing, or corrupt selected models fail rather than silently falling back. SSD PQ files are independent of navigation discovery.
  • Share a prepared-build context, scoped temporary-file ownership, and a navigation builder/codec registry. PQ and RaBitQ use the same build orchestration and file declarations; graph/SSD construction does not select the navigation codec. Local rollback cleans this build’s outputs without overwriting preexisting files. Remote rollback remains best-effort and depends on FileManager support; safe reuse after a failed upload is not guaranteed.

Shared correctness fixes

  • Use consistent SSD PQ scoring in ordinary search, filtered brute-force fallback, iterators, and distance-by-ID evaluation. Apply IP score restoration to both ordinary and exhausted iterator batches.
  • Report SSD PQ raw-vector capabilities consistently through static and runtime APIs; reject raw-vector retrieval when only SSD PQ is stored.
  • Account for the actual aligned coordinate-cache allocation and avoid treating the default EmbList strategy as an explicit EmbList input.
  • Handle empty Faiss array I/O safely in the serialization helper.
  • Do not infer persisted PQ or one-bit RaBitQ from omitted configuration when estimating load resources. Incomplete model information uses the full serialized file size plus the existing engine allowance; complete model information retains codec-specific estimation. Cache budgets are included for both PQ and RaBitQ.

Scope

RaBitQ iterator/range-search and embedding-list support are not included. AiSAQ continues to use PQ navigation. This does not change HNSW RaBitQ kernels, add TQ/IVF support, or relax the existing PQ chunk build limit.

Validation

Latest resource-estimation fix at 4155d20e:

  • Release: 65 related test cases / 3,183,039 assertions passed; Python roundtrips: 9/9 passed.
  • Targeted resource checks: 3 cases / 1,591 assertions, including automatic loading of RBQ8/9 with SSD PQ and the unchanged one-bit build default. The missing-codec regression reproduces against the archived previous library.
  • Sanitizers and performance were not rerun for this static-estimation change; results below are from the previous checkpoint.

Earlier local x86-64 validation at 48768d18:

  • Release: 63 test cases / 3,179,121 assertions passed, including DiskANN, AiSAQ, and HNSW RaBitQ regressions.
  • ASan/UBSan: 18 targeted cases / 118,567 assertions passed, with leak checking enabled.
  • Python: 9 roundtrip cases passed for L2/IP/COSINE and both RaBitQ entry points plus default PQ.
  • Added coverage for public static configuration/capability APIs, zero navigation-PQ budgets, SSD FP32/PQ, cached/uncached search, filtered fallback, and iterator exhaustion with independent score references.
  • Added automatic navigation discovery, explicit conflicts, ambiguity/corruption, failed-load retry, and build-registration rollback coverage. Python roundtrips load fresh generic DISKANN nodes without repeating build-time codec parameters.
  • Existing-index PQ/RaBitQ checks on SIFT, GIST, Cohere, and OpenAI: 96 matrix measurements plus 8 longer spot checks against 36f1cd09, with identical recall, IDs, and distance hashes. High-dimensional RaBitQ QPS changes ranged from -0.39% to +1.18%; a longer GIST PQ check reproduced 2.35% lower QPS and 3.74% higher CPU/query. This is not a blanket no-regression guarantee.
  • Actual mid-write navigation/graph/SSD failure injection remains incomplete; simulated temporary-artifact cleanup is not a substitute for all storage-failure scenarios.

ARM/Cardinal builds and upstream CI have not been validated locally. All commits include DCO sign-offs. Experiment reports and datasets are not part of this PR.

Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
…urces

Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
@sre-ci-robot

Copy link
Copy Markdown
Collaborator

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: CLiqing

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Details Needs approval from an approver in each of these files:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

@mergify

mergify Bot commented Sep 22, 2026

Copy link
Copy Markdown

@CLiqing 🔍 Important: PR Classification Needed!

For efficient project management and a seamless review process, it's essential to classify your PR correctly. Here's how:

  1. If you're fixing a bug, label it as kind/bug.
  2. For small tweaks (less than 20 lines without altering any functionality), please use kind/improvement.
  3. Significant changes that don't modify existing functionalities should be tagged as kind/enhancement.
  4. Adjusting APIs or changing functionality? Go with kind/feature.

For any PR outside the kind/improvement category, ensure you link to the associated issue using the format: “issue: #”.

Thanks for your efforts and contribution to the community!.

…itches

Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
…decs

Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
…fore writes

Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
@alexanderguzhva

Copy link
Copy Markdown
Collaborator

@CLiqing is the PR ready to be reviewed?

…ults

Signed-off-by: ChenLiqing <23721160+CLiqing@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants