Skip to content

feat(rust): add IVF-SQ bindings#2229

Open
jamie8johnson wants to merge 2 commits into
rapidsai:mainfrom
jamie8johnson:rust-ivf-sq
Open

feat(rust): add IVF-SQ bindings#2229
jamie8johnson wants to merge 2 commits into
rapidsai:mainfrom
jamie8johnson:rust-ivf-sq

Conversation

@jamie8johnson

Copy link
Copy Markdown
Contributor

Adds Rust bindings for the IVF-SQ index, wrapping the C API introduced in v26.06 ("IVF-SQ C API"). Follow-up to the CAGRA serialize/deserialize and search_with_filter wrappers (#2019 lineage).

What

  • New rust/cuvs/src/ivf_sq/ module mirroring the cagra module structure: IndexParams / SearchParams builders, Index with build, search, search_with_filter (cuvsFilter bitset), extend, serialize / deserialize, and the n_lists / dim / size getters.
  • rust/cuvs-sys/src/bindings.rs regenerated via rust/scripts/generate-bindings.sh — purely additive (cuvsIvfSq* block).

Notes for reviewers

  • cuvsIvfSqIndexGetCenters is deliberately not wrapped, matching the cagra and ivf_pq modules — it fills a caller-allocated DLManagedTensor, which has no ergonomic safe-Rust shape yet (commented in index.rs).
  • IVF-SQ exposes getters + extend where the ivf_pq Rust module currently doesn't; the C API supports them for both, so this module simply wraps the full practical surface. Happy to trim for symmetry if preferred.

Testing

9 GPU tests (single-process): params builders, build + self-neighbor search, repeated search, getters, extend-from-empty (NULL indices path), bitset-filtered search asserting exclusion, serialize/deserialize round-trip re-verifying search on the loaded index, interior-NUL path rejection. cargo fmt / clippy clean. Validated against both the v26.06.00 tag (conda libcuvs 26.06) and current main — the ivf_sq.h header is byte-identical between the two.

Wraps the IVF-SQ C API (added in v26.06) in the Rust crate, mirroring
the cagra module structure: IndexParams and SearchParams builders,
Index with build / search / search_with_filter (cuvsFilter bitset) /
extend / serialize / deserialize, and the NLists/Dim/Size getters.
cuvsIvfSqIndexGetCenters is deliberately not wrapped, matching cagra
and ivf_pq. rust/cuvs-sys/src/bindings.rs regenerated (additive).

9 GPU tests: params builders, build + self-neighbor search, repeated
search, getters, extend from an empty index, bitset-filtered search
(verifies exclusion), serialize/deserialize round-trip (re-verifies
search on the loaded index), interior-NUL path rejection.
@copy-pr-bot

copy-pr-bot Bot commented Jun 9, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jun 9, 2026

Copy link
Copy Markdown

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c887fe6d-2200-460e-bd93-0e8503ae6d6d

📥 Commits

Reviewing files that changed from the base of the PR and between 2beefce and 0fa1808.

📒 Files selected for processing (1)
  • rust/cuvs/src/ivf_sq/index.rs
🚧 Files skipped from review as they are similar to previous changes (1)
  • rust/cuvs/src/ivf_sq/index.rs

📝 Walkthrough

Summary by CodeRabbit

  • New Features

    • Added IVF-SQ (IVF-Scalar Quantization) approximate nearest-neighbor indexing: build, search (with optional filtering), extend, and save/load.
    • Exposed configurable index and search parameters and index property queries.
    • Public API surface updated to include the new IVF-SQ module and types.
  • Tests

    • Added unit tests covering build/search/extend/serialize-deserialize and parameter behavior.

Walkthrough

Adds Rust FFI bindings and safe Rust wrappers for an IVF‑SQ approximate nearest‑neighbor index: parameter types, search parameters, Index wrapper (build/search/extend/serialize/deserialize), module exports, and unit tests.

Changes

IVF-SQ Approximate Neighbor Index

Layer / File(s) Summary
FFI Bindings and C Contract
rust/cuvs-sys/src/bindings.rs
C FFI definitions for cuvsIvfSqIndexParams, cuvsIvfSqSearchParams, cuvsIvfSqIndex handle, and core operations: cuvsIvfSqBuild, cuvsIvfSqSearch, cuvsIvfSqExtend, cuvsIvfSqSerialize/cuvsIvfSqDeserialize, plus metadata getters.
Parameter Configuration Wrappers
rust/cuvs/src/ivf_sq/index_params.rs, rust/cuvs/src/ivf_sq/search_params.rs
Rust IndexParams and SearchParams wrappers with fluent setters, Debug impls, Drop cleanup, and unit tests verifying underlying fields.
Index Type and Core Operations
rust/cuvs/src/ivf_sq/index.rs
Index wrapper providing build(), new(), search(), search_with_filter(), extend(), n_lists()/dim()/size(), serialize()/deserialize(), Drop, and GPU-based unit tests covering search, extend, filtering, and persistence.
Module Declaration and Public Exports
rust/cuvs/src/ivf_sq/mod.rs, rust/cuvs/src/lib.rs
Module-level documentation with example, submodule wiring, and re-exports of Index, IndexParams, and SearchParams; registers ivf_sq in the crate.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

  • rapidsai/cuvs#1913: Introduces the underlying C IVF‑SQ API that these Rust bindings consume.
  • rapidsai/cuvs#1910: Implements the C-level IVF‑SQ index and core operations referenced by the new bindings.

Suggested labels

improvement, non-breaking, Rust

Suggested reviewers

  • robertmaynard
  • AyodeAwe
  • dantegd
🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title 'feat(rust): add IVF-SQ bindings' directly and concisely summarizes the main change: adding Rust bindings for the IVF-SQ index. It is clear, specific, and accurately represents the primary objective of the changeset.
Description check ✅ Passed The description is directly related to the changeset, providing context about the IVF-SQ C API bindings, explaining the module structure, implementation details, design decisions (e.g., why cuvsIvfSqIndexGetCenters is not wrapped), and comprehensive testing information.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands and usage tips.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant