feat(rust)!: compile-time dataset lifetime safety for all index types

[zbennett10]

This is follow-on from #1839 based on comments from @benfred

Introduces compile-time dataset lifetime enforcement across all 4 Rust index types (brute_force, cagra, ivf_flat, ivf_pq)... assuming we want to do it for all of them! It addresses the use-after-free risk when ManagedTensor::from(&ndarray) creates a non-owning view and the C library stores a pointer to the dataset.. and who knows what shenanigans may go on "down there".. 😄

Main Changes

• DatasetOwnership<'a> enum in dlpack.rs — tracks whether an index borrows or owns its dataset
• Index<'a> with dual constructors for all 4 index types:
  • build(&'a ManagedTensor) — borrowed path, compiler enforces dataset outlives index
  • build_owned(ManagedTensor) — owned path ('static), index is self-contained
• create_handle() helper — avoids double-free by separating FFI handle creation from Index struct construction

Design

pub struct Index<'a> { // Ties the lifetime of the index to it's underlying data
    inner: ffi::cuvsXxxIndex_t,
    _data: DatasetOwnership<'a>,
}

// Borrowed — compiler enforces arr outlives index
let tensor = ManagedTensor::from(&arr);
let index = Index::build(&res, &params, &tensor)?;

// Owned — 'static, self-contained
let device = ManagedTensor::from(&arr).to_device(&res)?;
let index = Index::build_owned(&res, &params, device)?;
drop(arr); // Fine — index owns the GPU copy

Breaking Change

Index::build() now takes &ManagedTensor instead of impl Into<ManagedTensor>. Users should:

• Use build(&tensor) with an explicit ManagedTensor reference (borrowed)
• Use build_owned(tensor) for the old move semantics (owned)

[copy-pr-bot]

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.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[zbennett10]

@benfred This does introduce a breaking change.. which I"m not super happy about it.. I mean we could do something like this where we have a backward-compatible approach.. basically keep build() as the owned path, add build_borrowed() as the new, safe path.. but idk maybe it's best to bite the bullet there.

// Keep the old build() signature — it becomes the "owned" path                                                                                                                               
impl Index<'static> {                                                                                                                                                                         
    pub fn build<T: Into<ManagedTensor>>(
        res: &Resources,
        params: &IndexParams,
        dataset: T,
    ) -> Result<Index<'static>> {
        let dataset: ManagedTensor = dataset.into();
        let inner = Self::create_handle()?;
        unsafe { check_cuvs(ffi::cuvsCagraBuild(res.0, params.0, dataset.as_ptr(), inner))?; }
        Ok(Index { inner, _data: DatasetOwnership::Owned(dataset) })
    }
}

// New safe borrowed path alongside it
impl<'a> Index<'a> {
    pub fn build_borrowed(
        res: &Resources,
        params: &IndexParams,
        dataset: &'a ManagedTensor,
    ) -> Result<Index<'a>> { ... }
}

Existing code like Index::build(&res, &params, &dataset) would still compile — the ndarray ref goes through Into<ManagedTensor> and gets owned.

But — the struct still has <'a>, so fn get_index() -> Index breaks. So for this to work we would need a type alias:
Something like pub type OwnedIndex = Index<'static>;

[aamijar]

/ok to test e006925

[benfred]

/ok to test 994f408

[cjnolet]

/ok to test 563762c

[yan-zaretskiy]

Hey Zachary, thank you for your contribution — the current state of the Rust bindings is indeed unsound for certain index types and this is a very important area of improvement. I've spent some time investigating the C++ internals to understand exactly which index types need lifetime enforcement, and here's what I found:

• BruteForce always stores a non-owning view of the dataset
• CAGRA behavior is runtime-dependent. make_aligned_dataset() in common.hpp stores a non_owning_dataset (just an mdspan view) when the input is device-accessible, row-major, and 16-byte aligned; and it performs a copy otherwise. Hence at the type level we have no choice but to also enforce a lifetime.
• IVF-Flat & IVF-PQ always copy the dataset, they don't need to reference the original data.

Our Rust bindings should retain the same ownership semantics as the underlying library. For this reason I feel that I'd need to push back on the proposed design.

Besides, the way it is implemented in this PR is too fragile. The ManagedTensor in itself does not indicate at the type level if it actually owns the underlying dataset, even if it has a non-null deleter. So the Owned variant of the DatasetOwnership<'a> does not express the true ownership.

Instead, I propose to simplify this PR as follows:

1. Revert IVF-Flat and IVF-PQ — these always copy
2. Remove the build_owned method altogether.
3. Use PhantomData<&'a ()> directly on the index structs and delete the DatasetOwnership. For the borrowed path, the caller keeps the tensor alive and the compiler enforces it via 'a.
4. deserialize should return Index<'static> — deserialized indices own their data.

[zbennett10]

Appreciate you @yan-zaretskiy — thanks a lot for the great feedback

The public API change is narrower now — just cagra::Index::build and brute_force::Index::build taking &ManagedTensor instead of impl Into<ManagedTensor>. IVF-Flat and IVF-PQ are completely untouched. Let me know what you think...

[yan-zaretskiy]

Looks good to me otherwise!

[yan-zaretskiy]

@zbennett10 Could you resolve the branch conflict so that we can merge this PR?

[zbennett10]

@zbennett10 Could you resolve the branch conflict so that we can merge this PR?

@yan-zaretskiy 👍