Skip to content

feat: add RaBitQ navigation to DiskANN - #1839

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

CLiqing wants to merge 12 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. Generation failures clean this build’s unfinished outputs without overwriting preexisting files. Once local generation completes, retain formal outputs before the first upload. Publication failures release temporary files but preserve formal outputs, block automatic same-prefix rebuilds, and do not call unsupported FileManager rollback operations. Remote partial uploads remain subject to the consumer’s existing lifecycle; this does not make the FileManager safe for arbitrary reuse.

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 publication-failure fix at 1a122663:

  • Rebuilt Release library and tests: all 33 DiskANN/AiSAQ test cases / 2,311,805 assertions passed.
  • Added PQ/RaBitQ + IP + SSD PQ failure coverage: first upload and later optional-file upload, false returns and exceptions, unsupported RemoveFile, complete local output retention, temporary-file cleanup, and same-prefix retry rejection with both reused and fresh managers. Rejected retries do not add uploads or duplicate registered byte counts.
  • The new regression fails against the archived 4155d20e library, which still attempts unsupported rollback.
  • No Milvus changes or navigation capability extensions. Python, sanitizers, performance, and remote upload integration were not rerun for this lifecycle fix.

Previous resource-estimation validation 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-file ownership 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.

Review ranges

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>
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