feat(rust): add tiered index bindings#2235
Open
jamie8johnson wants to merge 3 commits into
Open
Conversation
Wrap the tiered index C API (cuvsTieredIndex*) in a safe Rust module under rust/cuvs/src/tiered_index/. The tiered index couples a brute-force tier that absorbs incremental inserts with an ANN tier (CAGRA by default; IVF-Flat / IVF-PQ selectable via AnnAlgo), so vectors added with extend() are immediately searchable. - IndexParams: metric, algo, min_ann_rows, create_ann_index_on_extend, plus embedded per-backend ANN params (set_cagra_params / set_ivf_flat_params / set_ivf_pq_params) whose wrappers are retained so the C struct outlives the raw pointer it stores. - Index: build, extend, search, and search_with_filter (the C cuvsTieredIndexSearch always takes a cuvsFilter prefilter; search() passes a NO_FILTER sentinel). - Balanced Create/Destroy via Drop; assume_init only after check_cuvs. The C API offers no serialize/compact/merge-less persistence, so no filesystem paths are wrapped. The cuvsTieredIndex* symbols are already emitted by the checked-in bindings.rs (via cuvs/core/all.h), so no bindings regeneration diff is needed. Tests cover build+search, the key insert-after-build visibility contract, repeated extends, and filtered search. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR introduces a new ChangesTiered Index Module
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/cuvs/src/tiered_index/index.rs`:
- Line 8: Index::search and Index::search_with_filter currently always accept
crate::cagra::SearchParams and cast it to *mut c_void, which breaks
cuvsTieredIndexSearch when index.algo is IVF-* because the C API expects a
different cuvs*C*SearchParams_t; update the code so the tiered index dispatches
search params by algo: introduce a SearchParams enum/trait wrapper (e.g.,
TieredSearchParams with variants Cagra(crate::cagra::SearchParams) and
Ivf(crate::ivf::SearchParams) or a trait object) and in Index::search /
Index::search_with_filter match on index.algo (from IndexParams::set_algo) to
cast the correct variant to *mut c_void (or call the correct FFI binding),
ensuring you use the correct cuvs*C*SearchParams_t type for each backend
(reference cuvsTieredIndexSearch, index.algo, IndexParams::set_algo,
crate::cagra::SearchParams and the IVF search params type); alternatively, if
you prefer a smaller change, restrict IndexParams::set_algo to only allow CAGRA
for tiered indexes until you implement the multi-backend search params dispatch.
- Around line 116-136: The public safe method Index::search_with_filter
currently accepts ffi::cuvsFilter (a #[repr(C)] pub struct with an address
field) and passes it directly to cuvsTieredIndexSearch, allowing callers to
supply arbitrary/dangling addresses; fix by making search_with_filter unsafe
(i.e., require callers to uphold the invariant that cuvsFilter.addr is valid) or
replace the parameter with a new safe wrapper type (e.g., a validated
FilterHandle/newtype around cuvsFilter that enforces/validates addr and type_)
and adjust all call sites; ensure the change is applied to the
search_with_filter signature and its unsafe ffi::cuvsTieredIndexSearch
invocation so only validated filters reach the FFI boundary.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 49ddfffe-085d-4c11-ab84-1eb82e4d8e18
📒 Files selected for processing (4)
rust/cuvs/src/lib.rsrust/cuvs/src/tiered_index/index.rsrust/cuvs/src/tiered_index/index_params.rsrust/cuvs/src/tiered_index/mod.rs
…lter API Two CodeRabbit findings on the tiered-index Rust bindings: 1. (Critical, soundness) `Index::search_with_filter` took a raw `ffi::cuvsFilter` in a safe function, letting a caller pass an arbitrary address into the FFI search. It now takes a `&ManagedTensor` bitset and builds the `cuvsFilter` (BITSET) internally, mirroring the cagra binding. The bitset is documented as a 1-D u32 device tensor with ceil(n_rows/32) elements, bit 1 = include. Both `search` (NO_FILTER sentinel) and `search_with_filter` now route through a private `search_impl` helper that performs the single FFI call. 2. (Major, UB-reachable) Search params were hardcoded to `cagra::SearchParams` while `IndexParams::set_algo` permits IVF_FLAT/IVF_PQ. `cuvsTieredIndexSearch` reinterprets the opaque `void*` per the index's build-time algo, so wrong-backend params were undefined behavior. Added a backend-aware `tiered_index::SearchParams` enum (Cagra/IvfFlat/IvfPq) exposing `algo()` and `as_void_ptr()`. `Index` now captures the build-time `algo`; `search`/`search_with_filter` validate the params variant against it and return `Error::InvalidArgument` on mismatch. The empty-handle constructor is now private (`create_handle`) so `build` is the only public constructor and always records the backend. Updated the module doc example and tests to the new enum, switched the filtered-search test to pass the bitset tensor directly, and added `test_search_params_backend_mismatch` (CAGRA index searched with IVF-Flat params asserts `InvalidArgument`). All 7 lib tests + the doctest pass; clippy clean for the module. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
Both findings addressed in e2811ed (follow-up commit, no history rewrite):
7/7 tests + doctest pass on GPU against libcuvs 26.06. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
rust/cuvs/src/tiered_index/index.rs (1)
399-440: ⚡ Quick winAdd one successful non-CAGRA search smoke test.
The new mismatch test proves rejection, but the suite still never builds a tiered index with a non-CAGRA backend and searches it successfully. A bad
algo()/as_void_ptr()mapping for IVF-Flat would still pass this PR.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/cuvs/src/tiered_index/index.rs` around lines 399 - 440, Add a positive smoke test that builds a tiered index with a non-CAGRA backend and verifies a successful search to ensure IVF-Flat mapping is correct: create a new test (e.g., test_non_cagra_search_smoke) that uses SearchParams::IvfFlat to build the index via Index::build (instead of default_params()), run Index::search with queries/neighbors/distances like in test_search_params_backend_mismatch, assert the search returns Ok and that neighbors/distances contain expected (non-empty) values; this will exercise the algo()/as_void_ptr() path for IVF-Flat and catch any incorrect mapping.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rust/cuvs/src/tiered_index/index.rs`:
- Around line 47-57: Index::create_handle() allocates a native cuvsTieredIndex_t
and currently leaks it if ffi::cuvsTieredIndexBuild(...) returns Err because
Index is never constructed; modify the build path so that if
cuvsTieredIndexBuild fails you immediately call the destructor
(ffi::cuvsTieredIndexDestroy(handle)) before returning the Err. Locate the code
around Index::create_handle(), the unsafe block calling cuvsTieredIndexBuild,
and the Index { handle, algo: ... } return, and ensure on any early error you
destroy the raw handle (the same handle variable) so Drop is not the only
cleanup path.
---
Nitpick comments:
In `@rust/cuvs/src/tiered_index/index.rs`:
- Around line 399-440: Add a positive smoke test that builds a tiered index with
a non-CAGRA backend and verifies a successful search to ensure IVF-Flat mapping
is correct: create a new test (e.g., test_non_cagra_search_smoke) that uses
SearchParams::IvfFlat to build the index via Index::build (instead of
default_params()), run Index::search with queries/neighbors/distances like in
test_search_params_backend_mismatch, assert the search returns Ok and that
neighbors/distances contain expected (non-empty) values; this will exercise the
algo()/as_void_ptr() path for IVF-Flat and catch any incorrect mapping.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 45b3a31f-893e-4f77-ad89-9e87ef1968ed
📒 Files selected for processing (4)
rust/cuvs/src/tiered_index/index.rsrust/cuvs/src/tiered_index/index_params.rsrust/cuvs/src/tiered_index/mod.rsrust/cuvs/src/tiered_index/search_params.rs
✅ Files skipped from review due to trivial changes (1)
- rust/cuvs/src/tiered_index/search_params.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- rust/cuvs/src/tiered_index/mod.rs
Comment on lines
+47
to
+57
| let handle = Index::create_handle()?; | ||
| unsafe { | ||
| check_cuvs(ffi::cuvsTieredIndexBuild( | ||
| res.0, | ||
| params.as_ptr(), | ||
| dataset.as_ptr(), | ||
| handle, | ||
| ))?; | ||
| } | ||
| // Capture the backend so search() can reject mismatched params. | ||
| Ok(Index { handle, algo: params.algo() }) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '41,70p' rust/cuvs/src/tiered_index/index.rs
echo
echo "== Tiered index C API declarations/usages =="
rg -n -C2 "cuvsTieredIndex(Create|Build|Destroy)" rust/cuvs-sys rust/c/include c/src rust/cuvs/src
echo
echo "== Other Rust wrappers with create/build/destroy patterns =="
rg -n -C3 "Create\\(|Build\\(|Destroy\\(" rust/cuvs/src --glob '*.rs'Repository: rapidsai/cuvs
Length of output: 8305
Destroy native cuvsTieredIndex_t handle when cuvsTieredIndexBuild fails (rust/cuvs/src/tiered_index/index.rs:47-57)
Index::create_handle() allocates the native handle via cuvsTieredIndexCreate, and Drop is the only place that calls cuvsTieredIndexDestroy. If cuvsTieredIndexBuild returns Err, Index is never constructed, so Drop won’t run—leaving the allocated native handle to leak. Ensure cuvsTieredIndexDestroy(handle) is called on the build error path.
🛠️ Minimal fix
- let handle = Index::create_handle()?;
- unsafe {
- check_cuvs(ffi::cuvsTieredIndexBuild(
- res.0,
- params.as_ptr(),
- dataset.as_ptr(),
- handle,
- ))?;
- }
+ let handle = Index::create_handle()?;
+ if let Err(err) = unsafe {
+ check_cuvs(ffi::cuvsTieredIndexBuild(
+ res.0,
+ params.as_ptr(),
+ dataset.as_ptr(),
+ handle,
+ ))
+ } {
+ let _ = check_cuvs(unsafe { ffi::cuvsTieredIndexDestroy(handle) });
+ return Err(err);
+ }
// Capture the backend so search() can reject mismatched params.
Ok(Index { handle, algo: params.algo() })🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rust/cuvs/src/tiered_index/index.rs` around lines 47 - 57,
Index::create_handle() allocates a native cuvsTieredIndex_t and currently leaks
it if ffi::cuvsTieredIndexBuild(...) returns Err because Index is never
constructed; modify the build path so that if cuvsTieredIndexBuild fails you
immediately call the destructor (ffi::cuvsTieredIndexDestroy(handle)) before
returning the Err. Locate the code around Index::create_handle(), the unsafe
block calling cuvsTieredIndexBuild, and the Index { handle, algo: ... } return,
and ensure on any early error you destroy the raw handle (the same handle
variable) so Drop is not the only cleanup path.
Proves the SearchParams::algo()/as_void_ptr() mapping end-to-end for a non-CAGRA backend; the mismatch test alone only proved rejection.
Contributor
Author
|
Good catch — added test_tiered_ivf_flat_backend_search: builds an IVF-Flat-backed tiered index and verifies self-neighbor search end-to-end, so the algo/params mapping for a non-CAGRA backend is now positively exercised, not just the mismatch rejection. 8/8 tests green on GPU. |
Contributor
Author
|
Downstream data point for prioritization: we are now consuming this branch in production (pinned via a fork patch in cqs, a code-search engine with a file-watcher daemon). The tiered index is what lets us eliminate periodic full CAGRA/HNSW rebuilds — incremental adds land in the brute-force tier and compact into CAGRA, so the watcher's rebuild thread (and its crash-recovery machinery) goes away entirely. Practical notes from integration so far: the 26.06 conda |
jamie8johnson
added a commit
to jamie8johnson/cqs
that referenced
this pull request
Jun 11, 2026
…s replace periodic rebuilds (opt-in) (#1794) - New `tiered-index` cargo feature builds the cuVS tiered index backend (`src/tiered.rs`): a brute-force tier absorbs incremental `extend`s and stays searchable immediately, while the CAGRA ANN tier compacts inside cuVS — the watch loop's periodic full HNSW rebuild becomes a no-op when active. - `TieredBackend` registers at priority 150 (above CAGRA's 100) via the `register_index_backends!` table; gated runtime-on by `CQS_TIERED_INDEX=1`, so default and CAGRA selection paths are unchanged even with the feature compiled in. - Distance/score conventions mirror cagra.rs: Cosine keeps cuVS L2Expanded (`1 - d/2`), DotProduct sets InnerProduct (raw IP). ndarray_015 across the cuVS boundary; `unsafe impl Send/Sync` mirroring CagraIndex. - Watch integration: `tiered_index_active()` gates off the threshold-driven full HNSW rebuild when tiered serves (cheap incremental HNSW insert + base router stay live); HNSW/CAGRA rebuild paths fully intact. - Persistence: the cuVS C API has no tiered serialize/deserialize, so the tiered index is never persisted — the daemon rebuilds it from the store on restart. This still kills the *periodic* rebuild; only the cold-start build remains (same cost as the CAGRA build it replaces). - Fork pin: `cuvs::tiered_index` ships via `[patch.crates-io]` → jamie8johnson/cuvs branch cqs-tiered-26.6 (commit 5081e4eb), which is the upstream PR branch (rapidsai/cuvs#2235) + a version-pin 26.8.0→26.6.0. The patch block is COMMENTED OUT by default: plain `cuda-index` and crates.io builds resolve official cuvs 26.6 with no tiered symbols required (HARD RAIL, verified). Uncomment + `--features tiered-index` to build the backend. - Docs: CONTRIBUTING "Tiered-index fork pin" section (what/why/retirement via the #1679 playbook), README GPU section corrected (deleted a stale fork claim that lied about a #2019 patch + wrong 26.4 version; now states the current truth), env knobs documented, CHANGELOG entry. - Tests: backend ordering (3 backends, tiered leads), env override, build-param construction; an `#[ignore]` GPU smoke pinning build→search→extend→search (extended vector immediately findable as its own NN, no rebuild). Residual (issue-worthy): the watch loop does not yet hold a persistent tiered handle to call `extend` per-cycle — the served path rebuilds tiered from store on cache invalidation (still no periodic rebuild). Wiring a daemon-held tiered handle is the proper follow-up that turns each incremental add into an O(batch) `extend` instead of a from-store rebuild. Co-authored-by: jamie8johnson <jamie8johnson@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add Rust bindings for the tiered index
What
Adds a safe Rust wrapper for the tiered index C API
(
c/include/cuvs/neighbors/tiered_index.h) underrust/cuvs/src/tiered_index/.The tiered index couples a brute-force tier that absorbs incremental inserts
with an ANN tier (CAGRA by default; IVF-Flat and IVF-PQ also selectable). The
distinguishing property is that vectors added via
extendare immediatelysearchable, even before the ANN tier is (re)built.
New public API (
cuvs::tiered_index):IndexParams— wrapscuvsTieredIndexParams. Setters formetric,algo(
AnnAlgo=cuvsTieredIndexANNAlgo: CAGRA / IVF-Flat / IVF-PQ),min_ann_rows, andcreate_ann_index_on_extend. The embedded per-algorithmANN params (
cagra_params,ivf_flat_params,ivf_pq_params) are suppliedvia
set_cagra_params/set_ivf_flat_params/set_ivf_pq_params, whichmove ownership of the sub-params wrapper into
IndexParamsso the underlyingC struct outlives the raw pointer stored in the tiered params struct (same
pattern as
cagra::IndexParams::set_compression).Index— wrapscuvsTieredIndex.build,extend,search, andsearch_with_filter(the CcuvsTieredIndexSearchalways takes acuvsFilterprefilter;searchpasses aNO_FILTERsentinel, mirroring theexisting IVF-Flat wrapper).
All handles use balanced
Create/DestroywithDrop, andMaybeUninit::assume_initis only reached aftercheck_cuvssucceeds,consistent with the other neighbor modules.
The
cuvsTieredIndex*symbols are already produced by the checked-inrust/cuvs-sys/src/bindings.rs(the bindgen wrapper includescuvs/core/all.h, which transitively includestiered_index.h, and thecuvs.*allowlist captures the symbols), so no bindings change is required.rust/scripts/generate-bindings.shreproduces the same output.Reviewer notes
void*).cuvsTieredIndexSearchtakes the ANNbackend's search params as
void*. This wrapper passescagra::SearchParams(the default backend). If you'd prefer an enum thatdispatches over the three backends' search params, I'm happy to extend it —
the wrapping point is
Index::search_with_filter.entry points (confirmed via
nm -D libcuvs_c.so— onlyCreate/Destroy/ParamsCreate/ParamsDestroy/Build/Search/Extend/Mergeexist), so this PR intentionally omits persistence and takes no filesystem
paths.
Mergenot wrapped yet. The C API offerscuvsTieredIndexMerge; it isleft out of this first pass to keep the surface focused on
build/extend/search. Easy to add in a follow-up (or this PR) if desired.
26.8.0, but the newest libcuvs available in this contributor's environmentis the conda 26.06 build. The
cuvs-syscmake gate requires a libcuvs ofthe same major/minor and not older than the crate version, so against 26.06
the crate will not configure as-is. The tiered C symbols are present in
26.06 (verified with
nm -D), so to exercise the tests I temporarilypinned the workspace version to
26.6.0, ran the suite green on anRTX 4000(GPU 1), then reverted the version. CI on a matching 26.08 buildshould run the tests unmodified. Tested locally against libcuvs 26.06 with the workspace version temporarily pinned; CI against the current dev version should run unmodified.
Testing
Added unit tests in
tiered_index/index_params.rsandtiered_index/index.rs, plus a runnable doc example inmod.rs:test_index_params/test_embedded_cagra_params_retained— setters mutatethe C struct; embedded CAGRA params pointer stays live (not dangling).
test_tiered_build_and_search— build on an initial dataset; each queryfinds itself.
test_tiered_extend_visibility— the key test: vectors added viaextendafter build are immediately findable (each new vector is its ownrank-0 neighbor with ~0 self-distance).
test_tiered_repeated_extends— three successiveextendrounds; eachbatch is immediately findable (asserted as top-k membership + ~0
self-distance, robust to the ANN tier's approximate recall after an
on-extend rebuild).
test_tiered_filtered_search— a bitset prefilter excluding id 0 keeps id 0out of the result set.
Result (single-threaded,
CUDA_VISIBLE_DEVICES=1, conda libcuvs 26.06 withthe workspace temporarily pinned to 26.6.0 for linking):
Doc test for
mod.rsalso passes.cargo fmt --checkis clean andcargo clippyreports no findings intiered_index/(the pre-existingclippy::not_unsafe_ptr_arg_derefnote onresources.rsis unrelated to thischange and surfaces only under a newer local clippy).
Sibling-PR conflict note
This PR adds
pub mod tiered_index;torust/cuvs/src/lib.rs. Other queuedRust-binding PRs (brute-force serialize, scalar quantization, refine) also add
a
pub modline to the samelib.rsmodule list. These are adjacent one-lineadditions and will conflict trivially on the
modlist; resolve by keeping allthe new
pub modlines. No other files overlap.