cuvs/
ivf_flat.rs

1use std::{marker::PhantomData, ptr, ptr::NonNull};
2
3use crate::{
4    brute_force::{DistanceMetric, KnnSearchOutput, matrix_layout, validate_matrix_binding},
5    device::{DeviceBuffer, MatrixLayout, ValidatedDeviceTensor},
6    dlpack::DlpackTensor,
7    elements::IvfElement,
8    error::{Result, check_cuvs, from_cudf, invalid_input, invalid_state, require_nonzero},
9    params_handle::ParamsHandle,
10    resources::CuvsResources,
11};
12
13#[derive(Debug, Clone, Copy)]
14pub struct IvfFlatIndexParams {
15    pub metric: DistanceMetric,
16    pub n_lists: u32,
17    pub kmeans_n_iters: u32,
18    pub kmeans_trainset_fraction: f64,
19}
20
21impl IvfFlatIndexParams {
22    pub fn new(metric: DistanceMetric, n_lists: u32) -> Self {
23        Self { metric, n_lists, kmeans_n_iters: 20, kmeans_trainset_fraction: 0.5 }
24    }
25}
26
27#[derive(Debug, Clone, Copy)]
28pub struct IvfFlatSearchParams {
29    pub n_probes: u32,
30}
31
32impl Default for IvfFlatSearchParams {
33    fn default() -> Self {
34        Self { n_probes: 20 }
35    }
36}
37
38/// Query-local owned IVF-Flat index. Native build copies the indexed vectors.
39pub struct IvfFlatIndex<T: IvfElement> {
40    raw: Option<NonNull<cuvs_sys::cuvsIvfFlatIndex>>,
41    resources: CuvsResources,
42    rows: usize,
43    columns: usize,
44    n_lists: u32,
45    _element: PhantomData<T>,
46}
47
48impl<T: IvfElement> IvfFlatIndex<T> {
49    pub fn build(
50        resources: &CuvsResources,
51        dataset: &ValidatedDeviceTensor<'_, T>,
52        params: IvfFlatIndexParams,
53    ) -> Result<Self> {
54        let operation = "IvfFlatIndex::build";
55        validate_matrix_binding(resources, dataset, operation)?;
56        if matrix_layout(dataset, operation)? != MatrixLayout::RowMajor {
57            return Err(invalid_input(operation, "IVF-Flat dataset must be contiguous row-major"));
58        }
59        let rows = dataset.rows()?;
60        let columns = dataset.columns()?;
61        require_nonzero(rows, operation, "dataset rows")?;
62        if params.n_lists == 0 || params.n_lists as usize > rows {
63            return Err(invalid_input(
64                operation,
65                format!("n_lists ({}) must be in 1..=dataset rows ({rows})", params.n_lists),
66            ));
67        }
68        if params.kmeans_n_iters == 0
69            || !params.kmeans_trainset_fraction.is_finite()
70            || !(0.0..=1.0).contains(&params.kmeans_trainset_fraction)
71            || params.kmeans_trainset_fraction == 0.0
72        {
73            return Err(invalid_input(
74                operation,
75                "kmeans_n_iters must be positive and kmeans_trainset_fraction must be in (0, 1]",
76            ));
77        }
78        if params.metric == DistanceMetric::CosineExpanded && columns < 2 {
79            return Err(invalid_input(operation, "cosine IVF-Flat requires at least two dimensions"));
80        }
81        resources.with_native(operation, |native| {
82            let mut config = ParamsHandle::new(
83                cuvs_sys::cuvsIvfFlatIndexParamsCreate,
84                cuvs_sys::cuvsIvfFlatIndexParamsDestroy,
85                operation,
86                "IVF-Flat params",
87            )?;
88            unsafe {
89                let config_ptr = config.as_ptr();
90                (*config_ptr).metric = params.metric.raw();
91                (*config_ptr).n_lists = params.n_lists;
92                (*config_ptr).kmeans_n_iters = params.kmeans_n_iters;
93                (*config_ptr).kmeans_trainset_fraction = params.kmeans_trainset_fraction;
94            }
95            let mut raw = ptr::null_mut();
96            unsafe { check_cuvs(cuvs_sys::cuvsIvfFlatIndexCreate(&mut raw), operation)? };
97            let raw =
98                NonNull::new(raw).ok_or_else(|| invalid_state(operation, "cuVS returned a null IVF-Flat index"))?;
99            let mut tensor = DlpackTensor::from_const(dataset.as_tensor());
100            let built = unsafe {
101                check_cuvs(
102                    cuvs_sys::cuvsIvfFlatBuild(native, config.as_ptr(), tensor.as_mut_ptr(), raw.as_ptr()),
103                    operation,
104                )
105            };
106            let closed = config.close(operation);
107            if let Err(error) = built.and(closed) {
108                unsafe {
109                    let _ = cuvs_sys::cuvsIvfFlatIndexDestroy(raw.as_ptr());
110                }
111                return Err(error);
112            }
113            Ok(Self {
114                raw: Some(raw),
115                resources: resources.retain(),
116                rows,
117                columns,
118                n_lists: params.n_lists,
119                _element: PhantomData,
120            })
121        })
122    }
123
124    /// Search outputs are `[queries.rows(), k]` row-major. Missing neighbors have ordinal `i64::MAX`.
125    pub fn search_into(
126        &self,
127        queries: &ValidatedDeviceTensor<'_, T>,
128        k: usize,
129        params: IvfFlatSearchParams,
130        neighbors: &mut DeviceBuffer<i64>,
131        distances: &mut DeviceBuffer<f32>,
132    ) -> Result<KnnSearchOutput> {
133        let operation = "IvfFlatIndex::search_into";
134        require_nonzero(k, operation, "k")?;
135        if k > self.rows || k > u32::MAX as usize {
136            return Err(invalid_input(
137                operation,
138                format!("k ({k}) must be <= indexed rows ({}) and fit u32", self.rows),
139            ));
140        }
141        if params.n_probes == 0 || params.n_probes > self.n_lists {
142            return Err(invalid_input(
143                operation,
144                format!("n_probes ({}) must be in 1..=n_lists ({})", params.n_probes, self.n_lists),
145            ));
146        }
147        validate_matrix_binding(&self.resources, queries, operation)?;
148        if matrix_layout(queries, operation)? != MatrixLayout::RowMajor || queries.columns()? != self.columns {
149            return Err(invalid_input(
150                operation,
151                "IVF-Flat queries must be contiguous row-major with the indexed dimension",
152            ));
153        }
154        let n_queries = queries.rows()?;
155        let output_len = n_queries
156            .checked_mul(k)
157            .ok_or_else(|| invalid_input(operation, "search output element count overflowed usize"))?;
158        if neighbors.len() != output_len || distances.len() != output_len {
159            return Err(invalid_input(
160                operation,
161                format!("neighbors and distances must each contain n_queries * k ({output_len}) elements"),
162            ));
163        }
164        let mut neighbors_view = neighbors.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
165        let mut distances_view = distances.as_tensor_mut(&self.resources, &[n_queries, k], None, 0)?;
166        let raw = self.raw.ok_or_else(|| invalid_state(operation, "the IVF-Flat index has been closed"))?;
167        self.resources.with_native(operation, |native| {
168            let mut config = ParamsHandle::new(
169                cuvs_sys::cuvsIvfFlatSearchParamsCreate,
170                cuvs_sys::cuvsIvfFlatSearchParamsDestroy,
171                operation,
172                "IVF-Flat search params",
173            )?;
174            unsafe {
175                (*config.as_ptr()).n_probes = params.n_probes;
176            }
177            let mut query = DlpackTensor::from_const(queries.as_tensor());
178            let mut neighbor = DlpackTensor::from_mut(&mut neighbors_view);
179            let mut distance = DlpackTensor::from_mut(&mut distances_view);
180            let filter = cuvs_sys::cuvsFilter { addr: 0, type_: cuvs_sys::cuvsFilterType::NO_FILTER };
181            let result = unsafe {
182                check_cuvs(
183                    cuvs_sys::cuvsIvfFlatSearch(
184                        native,
185                        config.as_ptr(),
186                        raw.as_ptr(),
187                        query.as_mut_ptr(),
188                        neighbor.as_mut_ptr(),
189                        distance.as_mut_ptr(),
190                        filter,
191                    ),
192                    operation,
193                )
194            };
195            result.and(config.close(operation))
196        })?;
197        let indexed_rows =
198            i64::try_from(self.rows).map_err(|_| invalid_input(operation, "indexed row count exceeds Int64"))?;
199        // cuVS local top-k can return an in-range placeholder ID with an infinite distance.
200        neighbors
201            .inner_mut()
202            .fill_int64_where_float32_infinite(distances.inner(), i64::MAX)
203            .map_err(|error| from_cudf(operation, error))?;
204        neighbors.validate_native_index_range_or_sentinel(0, indexed_rows, i64::MAX, operation)?;
205        distances.validate_native_distances_for_valid_neighbors(neighbors, i64::MAX, operation)?;
206        Ok(KnnSearchOutput { n_queries, k })
207    }
208
209    pub fn rows(&self) -> usize {
210        self.rows
211    }
212    pub fn columns(&self) -> usize {
213        self.columns
214    }
215
216    /// Report native destruction errors; a failed destruction is retried on drop.
217    pub fn close(mut self) -> Result<()> {
218        self.close_inner("IvfFlatIndex::close")
219    }
220
221    fn close_inner(&mut self, operation: &'static str) -> Result<()> {
222        let Some(raw) = self.raw.take() else { return Ok(()) };
223        let result = self.resources.with_native(operation, |_| unsafe {
224            check_cuvs(cuvs_sys::cuvsIvfFlatIndexDestroy(raw.as_ptr()), operation)
225        });
226        if result.is_err() {
227            self.raw = Some(raw);
228        }
229        result
230    }
231}
232
233impl<T: IvfElement> Drop for IvfFlatIndex<T> {
234    fn drop(&mut self) {
235        let _ = self.close_inner("IvfFlatIndex::drop");
236    }
237}
238
239impl<T: IvfElement> std::fmt::Debug for IvfFlatIndex<T> {
240    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
241        formatter
242            .debug_struct("IvfFlatIndex")
243            .field("rows", &self.rows)
244            .field("columns", &self.columns)
245            .finish_non_exhaustive()
246    }
247}
cuvs/ivf_flat.rs

cuvs/
ivf_flat.rs
