From 992b77155ea319dd9ff1b58478f4c6ecfa41de43 Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Mon, 23 Feb 2026 15:36:42 -0600 Subject: [PATCH 1/7] fix(rust): change Index::search to take &self instead of self The search methods on CAGRA, IVF-PQ, IVF-Flat, and Brute Force indexes were taking `self` by value, consuming the index and preventing reuse. This required users to rebuild the index after every search operation, making real-time/high-throughput search impractical. This change allows indexes to be reused for multiple searches: Before: ```rust let index = Index::build(...)?; index.search(...)?; // First search - OK index.search(...)?; // COMPILE ERROR: use of moved value ``` After: ```rust let index = Index::build(...)?; index.search(...)?; // First search - OK index.search(...)?; // Second search - OK index.search(...)?; // Third search - OK ``` The underlying C FFI functions do not mutate the index during search, so taking a shared reference (&self) is safe and more ergonomic. This makes the Rust API consistent with other ANN libraries like faiss-rs, hnsw_rs, and annoy-rs. Fixes #1838 --- rust/cuvs/src/brute_force.rs | 2 +- rust/cuvs/src/cagra/index.rs | 2 +- rust/cuvs/src/ivf_flat/index.rs | 2 +- rust/cuvs/src/ivf_pq/index.rs | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index a0759ee49b..aed2aa44e1 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -62,7 +62,7 @@ impl Index { /// * `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors /// * `distances` - Matrix in device memory that receives the distances of the nearest neighbors pub fn search( - self, + &self, res: &Resources, queries: &ManagedTensor, neighbors: &ManagedTensor, diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs index 42f55659bd..046157672f 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/cagra/index.rs @@ -59,7 +59,7 @@ impl Index { /// * `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors /// * `distances` - Matrix in device memory that receives the distances of the nearest neighbors pub fn search( - self, + &self, res: &Resources, params: &SearchParams, queries: &ManagedTensor, diff --git a/rust/cuvs/src/ivf_flat/index.rs b/rust/cuvs/src/ivf_flat/index.rs index fa630a917c..fb7def9e89 100644 --- a/rust/cuvs/src/ivf_flat/index.rs +++ b/rust/cuvs/src/ivf_flat/index.rs @@ -59,7 +59,7 @@ impl Index { /// * `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors /// * `distances` - Matrix in device memory that receives the distances of the nearest neighbors pub fn search( - self, + &self, res: &Resources, params: &SearchParams, queries: &ManagedTensor, diff --git a/rust/cuvs/src/ivf_pq/index.rs b/rust/cuvs/src/ivf_pq/index.rs index 3a66b3d457..59dae7c0be 100644 --- a/rust/cuvs/src/ivf_pq/index.rs +++ b/rust/cuvs/src/ivf_pq/index.rs @@ -59,7 +59,7 @@ impl Index { /// * `neighbors` - Matrix in device memory that receives the indices of the nearest neighbors /// * `distances` - Matrix in device memory that receives the distances of the nearest neighbors pub fn search( - self, + &self, res: &Resources, params: &SearchParams, queries: &ManagedTensor, From c1d50ab29b650a4a1166dde868e089a1840b9996 Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Mon, 23 Feb 2026 21:19:08 -0600 Subject: [PATCH 2/7] test(rust): add multiple-search tests to validate index reuse Add tests for brute_force, cagra, ivf_flat, and ivf_pq that build one index and search it three times. This validates the &self change on search() allows proper index reuse. For brute_force, pass host data to build so the C++ index owns the device copy. The C++ brute_force::index stores a non-owning view when built with device data, so passing a device ManagedTensor (which is dropped after build) would cause use-after-free during search. --- rust/cuvs/Cargo.toml | 1 - rust/cuvs/src/brute_force.rs | 58 +++++++++++++++++++++++++++++++++ rust/cuvs/src/cagra/index.rs | 54 ++++++++++++++++++++++++++++++ rust/cuvs/src/ivf_flat/index.rs | 56 +++++++++++++++++++++++++++++++ rust/cuvs/src/ivf_pq/index.rs | 56 +++++++++++++++++++++++++++++++ 5 files changed, 224 insertions(+), 1 deletion(-) diff --git a/rust/cuvs/Cargo.toml b/rust/cuvs/Cargo.toml index 15bf8cf920..611d4e431d 100644 --- a/rust/cuvs/Cargo.toml +++ b/rust/cuvs/Cargo.toml @@ -14,4 +14,3 @@ ndarray = "0.15" [dev-dependencies] ndarray-rand = "0.14" -mark-flaky-tests = "1" diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index aed2aa44e1..394fff52f9 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -172,4 +172,62 @@ mod tests { fn test_l2() { test_bfknn(DistanceType::L2Expanded); } + + /// Test that an index can be searched multiple times without rebuilding. + /// This validates that search() takes &self instead of self. + #[test] + fn test_brute_force_multiple_searches() { + let res = Resources::new().unwrap(); + + // Create a random dataset + let n_datapoints = 64; + let n_features = 8; + let dataset_host = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + // Pass host data directly so the C++ index copies it to device and owns + // the copy. Passing a device ManagedTensor would create a non-owning view + // that becomes a dangling pointer once the ManagedTensor is dropped after + // build, causing use-after-free during search. + let index = Index::build(&res, DistanceType::L2Expanded, None, &dataset_host) + .expect("failed to create brute force index"); + + res.sync_stream().unwrap(); + + let k = 4; + + // Perform multiple searches on the same index + for search_iter in 0..3 { + let n_queries = 4; + let queries = dataset_host.slice(s![0..n_queries, ..]); + let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); + + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host) + .to_device(&res) + .unwrap(); + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host) + .to_device(&res) + .unwrap(); + + // This should work on every iteration because search() takes &self + index + .search(&res, &queries, &neighbors, &distances) + .expect(&format!("search iteration {} failed", search_iter)); + + // Copy back to host memory + distances.to_host(&res, &mut distances_host).unwrap(); + neighbors.to_host(&res, &mut neighbors_host).unwrap(); + res.sync_stream().unwrap(); + + // Verify results are consistent + assert_eq!( + neighbors_host[[0, 0]], 0, + "iteration {}: first query should find itself", + search_iter + ); + } + } } diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs index 046157672f..c14dc7e447 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/cagra/index.rs @@ -167,4 +167,58 @@ mod tests { .set_compression(CompressionParams::new().unwrap()); test_cagra(build_params); } + + /// Test that an index can be searched multiple times without rebuilding. + /// This validates that search() takes &self instead of self. + #[test] + fn test_cagra_multiple_searches() { + let res = Resources::new().unwrap(); + let build_params = IndexParams::new().unwrap(); + + // Create a random dataset + let n_datapoints = 256; + let n_features = 16; + let dataset = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + // Build the index once + let index = + Index::build(&res, &build_params, &dataset).expect("failed to create cagra index"); + + let search_params = SearchParams::new().unwrap(); + let k = 5; + + // Perform multiple searches on the same index + for search_iter in 0..3 { + let n_queries = 4; + let queries = dataset.slice(s![0..n_queries, ..]); + let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); + + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host) + .to_device(&res) + .unwrap(); + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host) + .to_device(&res) + .unwrap(); + + // This should work on every iteration because search() takes &self + index + .search(&res, &search_params, &queries, &neighbors, &distances) + .expect(&format!("search iteration {} failed", search_iter)); + + // Copy back to host memory + distances.to_host(&res, &mut distances_host).unwrap(); + neighbors.to_host(&res, &mut neighbors_host).unwrap(); + + // Verify results are consistent across searches + assert_eq!( + neighbors_host[[0, 0]], 0, + "iteration {}: first query should find itself", + search_iter + ); + } + } } diff --git a/rust/cuvs/src/ivf_flat/index.rs b/rust/cuvs/src/ivf_flat/index.rs index fb7def9e89..a0c653a478 100644 --- a/rust/cuvs/src/ivf_flat/index.rs +++ b/rust/cuvs/src/ivf_flat/index.rs @@ -157,4 +157,60 @@ mod tests { assert_eq!(neighbors_host[[2, 0]], 2); assert_eq!(neighbors_host[[3, 0]], 3); } + + /// Test that an index can be searched multiple times without rebuilding. + /// This validates that search() takes &self instead of self. + #[test] + fn test_ivf_flat_multiple_searches() { + let build_params = IndexParams::new().unwrap().set_n_lists(64); + let res = Resources::new().unwrap(); + + // Create a random dataset + let n_datapoints = 1024; + let n_features = 16; + let dataset = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + let dataset_device = ManagedTensor::from(&dataset).to_device(&res).unwrap(); + + // Build the index once + let index = Index::build(&res, &build_params, dataset_device) + .expect("failed to create ivf-flat index"); + + let search_params = SearchParams::new().unwrap(); + let k = 5; + + // Perform multiple searches on the same index + for search_iter in 0..3 { + let n_queries = 4; + let queries = dataset.slice(s![0..n_queries, ..]); + let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); + + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host) + .to_device(&res) + .unwrap(); + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host) + .to_device(&res) + .unwrap(); + + // This should work on every iteration because search() takes &self + index + .search(&res, &search_params, &queries, &neighbors, &distances) + .expect(&format!("search iteration {} failed", search_iter)); + + // Copy back to host memory + distances.to_host(&res, &mut distances_host).unwrap(); + neighbors.to_host(&res, &mut neighbors_host).unwrap(); + + // Verify results are consistent + assert_eq!( + neighbors_host[[0, 0]], 0, + "iteration {}: first query should find itself", + search_iter + ); + } + } } diff --git a/rust/cuvs/src/ivf_pq/index.rs b/rust/cuvs/src/ivf_pq/index.rs index 59dae7c0be..df3a8f3653 100644 --- a/rust/cuvs/src/ivf_pq/index.rs +++ b/rust/cuvs/src/ivf_pq/index.rs @@ -151,4 +151,60 @@ mod tests { assert_eq!(neighbors_host[[2, 0]], 2); assert_eq!(neighbors_host[[3, 0]], 3); } + + /// Test that an index can be searched multiple times without rebuilding. + /// This validates that search() takes &self instead of self. + #[test] + fn test_ivf_pq_multiple_searches() { + let build_params = IndexParams::new().unwrap().set_n_lists(64); + let res = Resources::new().unwrap(); + + // Create a random dataset + let n_datapoints = 1024; + let n_features = 16; + let dataset = + ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); + + let dataset_device = ManagedTensor::from(&dataset).to_device(&res).unwrap(); + + // Build the index once + let index = Index::build(&res, &build_params, dataset_device) + .expect("failed to create ivf-pq index"); + + let search_params = SearchParams::new().unwrap(); + let k = 5; + + // Perform multiple searches on the same index + for search_iter in 0..3 { + let n_queries = 4; + let queries = dataset.slice(s![0..n_queries, ..]); + let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); + + let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); + let neighbors = ManagedTensor::from(&neighbors_host) + .to_device(&res) + .unwrap(); + + let mut distances_host = ndarray::Array::::zeros((n_queries, k)); + let distances = ManagedTensor::from(&distances_host) + .to_device(&res) + .unwrap(); + + // This should work on every iteration because search() takes &self + index + .search(&res, &search_params, &queries, &neighbors, &distances) + .expect(&format!("search iteration {} failed", search_iter)); + + // Copy back to host memory + distances.to_host(&res, &mut distances_host).unwrap(); + neighbors.to_host(&res, &mut neighbors_host).unwrap(); + + // Verify results are consistent + assert_eq!( + neighbors_host[[0, 0]], 0, + "iteration {}: first query should find itself", + search_iter + ); + } + } } From 8511b1bc4a4d406f3b4a4d4abd9f3aad84030966 Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Tue, 24 Feb 2026 07:15:49 -0600 Subject: [PATCH 3/7] style: fix copyright dates and cargo-fmt formatting Update copyright years to 2024-2026 and fix assert_eq! macro formatting to match stable rustfmt style. --- rust/cuvs/src/brute_force.rs | 5 +++-- rust/cuvs/src/cagra/index.rs | 5 +++-- rust/cuvs/src/ivf_flat/index.rs | 5 +++-- rust/cuvs/src/ivf_pq/index.rs | 5 +++-- 4 files changed, 12 insertions(+), 8 deletions(-) diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index 394fff52f9..9100775880 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ //! Brute Force KNN @@ -224,7 +224,8 @@ mod tests { // Verify results are consistent assert_eq!( - neighbors_host[[0, 0]], 0, + neighbors_host[[0, 0]], + 0, "iteration {}: first query should find itself", search_iter ); diff --git a/rust/cuvs/src/cagra/index.rs b/rust/cuvs/src/cagra/index.rs index c14dc7e447..789f72b603 100644 --- a/rust/cuvs/src/cagra/index.rs +++ b/rust/cuvs/src/cagra/index.rs @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -215,7 +215,8 @@ mod tests { // Verify results are consistent across searches assert_eq!( - neighbors_host[[0, 0]], 0, + neighbors_host[[0, 0]], + 0, "iteration {}: first query should find itself", search_iter ); diff --git a/rust/cuvs/src/ivf_flat/index.rs b/rust/cuvs/src/ivf_flat/index.rs index a0c653a478..c38be828de 100644 --- a/rust/cuvs/src/ivf_flat/index.rs +++ b/rust/cuvs/src/ivf_flat/index.rs @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -207,7 +207,8 @@ mod tests { // Verify results are consistent assert_eq!( - neighbors_host[[0, 0]], 0, + neighbors_host[[0, 0]], + 0, "iteration {}: first query should find itself", search_iter ); diff --git a/rust/cuvs/src/ivf_pq/index.rs b/rust/cuvs/src/ivf_pq/index.rs index df3a8f3653..f61e3b771a 100644 --- a/rust/cuvs/src/ivf_pq/index.rs +++ b/rust/cuvs/src/ivf_pq/index.rs @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -201,7 +201,8 @@ mod tests { // Verify results are consistent assert_eq!( - neighbors_host[[0, 0]], 0, + neighbors_host[[0, 0]], + 0, "iteration {}: first query should find itself", search_iter ); From cc39b315cec8145bd54df8cb9367e4be2d68c0da Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Wed, 25 Feb 2026 07:59:33 -0600 Subject: [PATCH 4/7] fix: restore mark-flaky-tests dev-dependency removed during merge The merge with main brought in the `#[flaky]` attribute usage in brute_force.rs tests, but the Cargo.toml diff inadvertently removed the mark-flaky-tests dependency, causing rust-build CI to fail with "unresolved import mark_flaky_tests". --- rust/cuvs/Cargo.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/rust/cuvs/Cargo.toml b/rust/cuvs/Cargo.toml index 611d4e431d..15bf8cf920 100644 --- a/rust/cuvs/Cargo.toml +++ b/rust/cuvs/Cargo.toml @@ -14,3 +14,4 @@ ndarray = "0.15" [dev-dependencies] ndarray-rand = "0.14" +mark-flaky-tests = "1" From 43450cf87154ece52a9a5641151aedc7da3b6748 Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Wed, 25 Feb 2026 13:41:38 -0600 Subject: [PATCH 5/7] fix(rust): use device memory for brute force multiple-search test The brute_force C++ build function requires device memory (uses raft::device_matrix_view exclusively), unlike CAGRA which handles both host and device inputs. Convert the dataset to device memory before building, matching the pattern used by IVF-Flat and IVF-PQ tests. This fixes the "device_type mismatch between return mdspan and DLTensor" error in rust-build CI. --- rust/cuvs/src/brute_force.rs | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index 9100775880..6cff393a75 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -182,14 +182,12 @@ mod tests { // Create a random dataset let n_datapoints = 64; let n_features = 8; - let dataset_host = + let dataset = ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - // Pass host data directly so the C++ index copies it to device and owns - // the copy. Passing a device ManagedTensor would create a non-owning view - // that becomes a dangling pointer once the ManagedTensor is dropped after - // build, causing use-after-free during search. - let index = Index::build(&res, DistanceType::L2Expanded, None, &dataset_host) + // Build the brute force index from device memory + let dataset_device = ManagedTensor::from(&dataset).to_device(&res).unwrap(); + let index = Index::build(&res, DistanceType::L2Expanded, None, dataset_device) .expect("failed to create brute force index"); res.sync_stream().unwrap(); @@ -199,7 +197,7 @@ mod tests { // Perform multiple searches on the same index for search_iter in 0..3 { let n_queries = 4; - let queries = dataset_host.slice(s![0..n_queries, ..]); + let queries = dataset.slice(s![0..n_queries, ..]); let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); From 1b5f88ad68960445e028cd4685f74d15b977dcca Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Wed, 25 Feb 2026 20:21:58 -0600 Subject: [PATCH 6/7] fix(rust): keep dataset alive in brute force index to prevent use-after-free The brute force C library stores a non-owning view into the dataset (unlike CAGRA/IVF which build self-contained structures). Previously, Index::build() consumed the ManagedTensor dataset, which was dropped at the end of build(), freeing device memory while the index still held a pointer to it. This caused test_brute_force_multiple_searches to fail on CUDA 12.9.1 (returning garbage neighbor indices) while passing on 13.1.1 by luck (freed memory not yet overwritten). Fix: Store the ManagedTensor in the Index struct so device memory lives as long as the index. Drop order is correct: Index::drop() calls cuvsBruteForceIndexDestroy first, then _dataset is freed. Also fixes a copy-paste typo in Drop ("cagraIndexDestroy" -> "bruteForceIndexDestroy"). --- rust/cuvs/src/brute_force.rs | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index 6cff393a75..d8b8abba22 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -12,8 +12,18 @@ use crate::error::{check_cuvs, Result}; use crate::resources::Resources; /// Brute Force KNN Index +/// +/// Unlike CAGRA or IVF indexes which build self-contained data structures, +/// the brute force index holds a non-owning view into the original dataset. +/// The `_dataset` field keeps the device memory alive for the lifetime of +/// the index. #[derive(Debug)] -pub struct Index(ffi::cuvsBruteForceIndex_t); +pub struct Index { + inner: ffi::cuvsBruteForceIndex_t, + // The C library stores a view into the dataset (not a copy). + // We must keep the ManagedTensor alive so the device memory is not freed. + _dataset: Option, +} impl Index { /// Builds a new Brute Force KNN Index from the dataset for efficient search. @@ -31,16 +41,18 @@ impl Index { dataset: T, ) -> Result { let dataset: ManagedTensor = dataset.into(); - let index = Index::new()?; + let mut index = Index::new()?; unsafe { check_cuvs(ffi::cuvsBruteForceBuild( res.0, dataset.as_ptr(), metric, metric_arg.unwrap_or(2.0), - index.0, + index.inner, ))?; } + // Store the dataset so the device memory remains valid for searches. + index._dataset = Some(dataset); Ok(index) } @@ -49,7 +61,10 @@ impl Index { unsafe { let mut index = std::mem::MaybeUninit::::uninit(); check_cuvs(ffi::cuvsBruteForceIndexCreate(index.as_mut_ptr()))?; - Ok(Index(index.assume_init())) + Ok(Index { + inner: index.assume_init(), + _dataset: None, + }) } } @@ -76,7 +91,7 @@ impl Index { check_cuvs(ffi::cuvsBruteForceSearch( res.0, - self.0, + self.inner, queries.as_ptr(), neighbors.as_ptr(), distances.as_ptr(), @@ -88,8 +103,8 @@ impl Index { impl Drop for Index { fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.0) }) { - write!(stderr(), "failed to call cagraIndexDestroy {:?}", e) + if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.inner) }) { + write!(stderr(), "failed to call bruteForceIndexDestroy {:?}", e) .expect("failed to write to stderr"); } } From f9b434a15b0d7c3d26067c0c1ed8515d5491c867 Mon Sep 17 00:00:00 2001 From: Zachary Bennett Date: Tue, 3 Mar 2026 15:40:39 -0600 Subject: [PATCH 7/7] refactor(rust): split out brute force dataset lifetime to follow-up PR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Revert the brute force dataset ownership changes (storing ManagedTensor in the Index struct) as requested by reviewer. The dataset lifetime safety will be addressed in a follow-up PR with compile-time enforcement via DatasetOwnership<'a> across all 4 index types. This keeps PR #1839 focused on the self → &self search fix. Also removes the brute force multiple-search test since it cannot be safe without dataset lifetime enforcement (the C++ brute_force::index stores a non-owning view into the dataset). Fixes the "cagraIndexDestroy" → "bruteForceIndexDestroy" typo in Drop. --- rust/cuvs/src/brute_force.rs | 89 +++++------------------------------- 1 file changed, 12 insertions(+), 77 deletions(-) diff --git a/rust/cuvs/src/brute_force.rs b/rust/cuvs/src/brute_force.rs index d8b8abba22..1440bb3205 100644 --- a/rust/cuvs/src/brute_force.rs +++ b/rust/cuvs/src/brute_force.rs @@ -12,18 +12,8 @@ use crate::error::{check_cuvs, Result}; use crate::resources::Resources; /// Brute Force KNN Index -/// -/// Unlike CAGRA or IVF indexes which build self-contained data structures, -/// the brute force index holds a non-owning view into the original dataset. -/// The `_dataset` field keeps the device memory alive for the lifetime of -/// the index. #[derive(Debug)] -pub struct Index { - inner: ffi::cuvsBruteForceIndex_t, - // The C library stores a view into the dataset (not a copy). - // We must keep the ManagedTensor alive so the device memory is not freed. - _dataset: Option, -} +pub struct Index(ffi::cuvsBruteForceIndex_t); impl Index { /// Builds a new Brute Force KNN Index from the dataset for efficient search. @@ -41,18 +31,16 @@ impl Index { dataset: T, ) -> Result { let dataset: ManagedTensor = dataset.into(); - let mut index = Index::new()?; + let index = Index::new()?; unsafe { check_cuvs(ffi::cuvsBruteForceBuild( res.0, dataset.as_ptr(), metric, metric_arg.unwrap_or(2.0), - index.inner, + index.0, ))?; } - // Store the dataset so the device memory remains valid for searches. - index._dataset = Some(dataset); Ok(index) } @@ -61,10 +49,7 @@ impl Index { unsafe { let mut index = std::mem::MaybeUninit::::uninit(); check_cuvs(ffi::cuvsBruteForceIndexCreate(index.as_mut_ptr()))?; - Ok(Index { - inner: index.assume_init(), - _dataset: None, - }) + Ok(Index(index.assume_init())) } } @@ -91,7 +76,7 @@ impl Index { check_cuvs(ffi::cuvsBruteForceSearch( res.0, - self.inner, + self.0, queries.as_ptr(), neighbors.as_ptr(), distances.as_ptr(), @@ -103,7 +88,7 @@ impl Index { impl Drop for Index { fn drop(&mut self) { - if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.inner) }) { + if let Err(e) = check_cuvs(unsafe { ffi::cuvsBruteForceIndexDestroy(self.0) }) { write!(stderr(), "failed to call bruteForceIndexDestroy {:?}", e) .expect("failed to write to stderr"); } @@ -188,60 +173,10 @@ mod tests { test_bfknn(DistanceType::L2Expanded); } - /// Test that an index can be searched multiple times without rebuilding. - /// This validates that search() takes &self instead of self. - #[test] - fn test_brute_force_multiple_searches() { - let res = Resources::new().unwrap(); - - // Create a random dataset - let n_datapoints = 64; - let n_features = 8; - let dataset = - ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - - // Build the brute force index from device memory - let dataset_device = ManagedTensor::from(&dataset).to_device(&res).unwrap(); - let index = Index::build(&res, DistanceType::L2Expanded, None, dataset_device) - .expect("failed to create brute force index"); - - res.sync_stream().unwrap(); - - let k = 4; - - // Perform multiple searches on the same index - for search_iter in 0..3 { - let n_queries = 4; - let queries = dataset.slice(s![0..n_queries, ..]); - let queries = ManagedTensor::from(&queries).to_device(&res).unwrap(); - - let mut neighbors_host = ndarray::Array::::zeros((n_queries, k)); - let neighbors = ManagedTensor::from(&neighbors_host) - .to_device(&res) - .unwrap(); - - let mut distances_host = ndarray::Array::::zeros((n_queries, k)); - let distances = ManagedTensor::from(&distances_host) - .to_device(&res) - .unwrap(); - - // This should work on every iteration because search() takes &self - index - .search(&res, &queries, &neighbors, &distances) - .expect(&format!("search iteration {} failed", search_iter)); - - // Copy back to host memory - distances.to_host(&res, &mut distances_host).unwrap(); - neighbors.to_host(&res, &mut neighbors_host).unwrap(); - res.sync_stream().unwrap(); - - // Verify results are consistent - assert_eq!( - neighbors_host[[0, 0]], - 0, - "iteration {}: first query should find itself", - search_iter - ); - } - } + // NOTE: brute_force multiple-search test is omitted here because the C++ + // brute_force::index stores a non-owning view into the dataset. Building + // from device data via `build()` drops the ManagedTensor after the call, + // leaving a dangling pointer. A follow-up PR will add dataset lifetime + // enforcement (DatasetOwnership<'a>) to make this safe. + // See: https://github.com/rapidsai/cuvs/issues/1838 }