diff --git a/rust/cuvs/examples/cagra.rs b/rust/cuvs/examples/cagra.rs index b118c3abc5..8566e9c85b 100644 --- a/rust/cuvs/examples/cagra.rs +++ b/rust/cuvs/examples/cagra.rs @@ -10,7 +10,11 @@ use ndarray::s; use ndarray_rand::rand_distr::Uniform; use ndarray_rand::RandomExt; -/// Example showing how to index and search data with CAGRA +/// Example showing how to index and search data with CAGRA using the validated builder API. +/// +/// `IndexParams::builder()` validates parameters before any GPU allocation, surfacing +/// misconfiguration immediately with a clear error message instead of an opaque CUDA +/// assertion 1-2 seconds into `Index::build()`. fn cagra_example() -> Result<()> { let res = Resources::new()?; @@ -20,8 +24,14 @@ fn cagra_example() -> Result<()> { let dataset = ndarray::Array::::random((n_datapoints, n_features), Uniform::new(0., 1.0)); - // build the cagra index - let build_params = IndexParams::new()?; + // Build the CAGRA index using the validated builder. + // Parameters are checked in Rust before any FFI call — invalid values (e.g. + // graph_degree=0) produce an error here, not inside Index::build(). + let build_params = IndexParams::builder() + .graph_degree(32) + .intermediate_graph_degree(64) + .nn_descent_niter(20) + .build()?; let index = Index::build(&res, &build_params, &dataset)?; println!( "Indexed {}x{} datapoints into cagra index", diff --git a/rust/cuvs/src/cagra/index_params.rs b/rust/cuvs/src/cagra/index_params.rs index ea34959147..6655b24741 100644 --- a/rust/cuvs/src/cagra/index_params.rs +++ b/rust/cuvs/src/cagra/index_params.rs @@ -137,6 +137,31 @@ impl IndexParams { } } +impl IndexParams { + /// Returns a builder for constructing [`IndexParams`] with validated parameters. + /// + /// Unlike the `IndexParams::new()?.set_*()` setter chain, [`IndexParamsBuilder::build`] + /// validates all parameters in Rust before any FFI allocation. Invalid values produce a + /// clear error message naming the offending field and its valid range, before any GPU + /// work begins. + /// + /// # Example + /// + /// ```no_run + /// use cuvs::cagra::IndexParams; + /// + /// let params = IndexParams::builder() + /// .graph_degree(32) + /// .intermediate_graph_degree(64) + /// .nn_descent_niter(20) + /// .build() + /// .unwrap(); + /// ``` + pub fn builder() -> IndexParamsBuilder { + IndexParamsBuilder::default() + } +} + impl fmt::Debug for IndexParams { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { // custom debug trait here, default value will show the pointer address @@ -177,6 +202,115 @@ impl Drop for CompressionParams { } } +/// Builder for [`IndexParams`] with pre-validated parameters. +/// +/// Construct via [`IndexParams::builder()`]. Call [`IndexParamsBuilder::build`] to +/// validate all parameters and allocate the FFI struct in one step. +/// +/// Defaults match the cuVS C API defaults: `graph_degree=64`, +/// `intermediate_graph_degree=128`, `nn_descent_niter=20`. +pub struct IndexParamsBuilder { + graph_degree: usize, + intermediate_graph_degree: usize, + nn_descent_niter: usize, + build_algo: Option, + compression: Option, +} + +impl Default for IndexParamsBuilder { + fn default() -> Self { + Self { + graph_degree: 64, + intermediate_graph_degree: 128, + nn_descent_niter: 20, + build_algo: None, + compression: None, + } + } +} + +impl IndexParamsBuilder { + /// Degree of output graph. + /// + /// Must be > 0. Values that are multiples of 32 are preferred for warp alignment. + pub fn graph_degree(mut self, v: usize) -> Self { + self.graph_degree = v; + self + } + + /// Degree of input graph for pruning. + /// + /// Must be >= `graph_degree`. + pub fn intermediate_graph_degree(mut self, v: usize) -> Self { + self.intermediate_graph_degree = v; + self + } + + /// Number of iterations to run if building with NN_DESCENT. + /// + /// Must be > 0. + pub fn nn_descent_niter(mut self, v: usize) -> Self { + self.nn_descent_niter = v; + self + } + + /// ANN algorithm to build knn graph. + pub fn build_algo(mut self, v: BuildAlgo) -> Self { + self.build_algo = Some(v); + self + } + + /// Vector compression parameters. + pub fn compression(mut self, v: CompressionParams) -> Self { + self.compression = Some(v); + self + } + + /// Validate all parameters without allocating any GPU resources. + /// + /// Returns `Ok(())` if all parameters are valid, or `Err` with a message naming + /// the offending field and its valid range. + pub fn validate(&self) -> crate::error::Result<()> { + if self.graph_degree == 0 { + return Err(format!("graph_degree must be > 0; got {}", self.graph_degree).into()); + } + if self.intermediate_graph_degree < self.graph_degree { + return Err(format!( + "intermediate_graph_degree ({}) must be >= graph_degree ({})", + self.intermediate_graph_degree, self.graph_degree + ) + .into()); + } + if self.nn_descent_niter == 0 { + return Err(format!( + "nn_descent_niter must be > 0; got {}", + self.nn_descent_niter + ) + .into()); + } + Ok(()) + } + + /// Validate all parameters and allocate the FFI struct. + /// + /// Returns `Err` with a message naming the offending field and its valid range + /// before any GPU work begins. + pub fn build(self) -> crate::error::Result { + self.validate()?; + let mut params = IndexParams::new()? + .set_graph_degree(self.graph_degree) + .set_intermediate_graph_degree(self.intermediate_graph_degree) + .set_nn_descent_niter(self.nn_descent_niter); + if let Some(algo) = self.build_algo { + params = params.set_build_algo(algo); + } + if let Some(compression) = self.compression { + params = params.set_compression(compression); + } + Ok(params) + } +} + #[cfg(test)] mod tests { use super::*; @@ -206,4 +340,95 @@ mod tests { assert_eq!((*(*params.0).compression).pq_bits, 4); } } + + // --- IndexParamsBuilder tests --- + + #[test] + fn builder_rejects_zero_graph_degree() { + let err = IndexParams::builder() + .graph_degree(0) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("graph_degree"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_rejects_invalid_intermediate_degree() { + let err = IndexParams::builder() + .graph_degree(32) + .intermediate_graph_degree(16) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("intermediate_graph_degree"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_rejects_zero_niter() { + let err = IndexParams::builder() + .nn_descent_niter(0) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("nn_descent_niter"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_accepts_valid_params() { + assert!(IndexParams::builder() + .graph_degree(32) + .intermediate_graph_degree(64) + .nn_descent_niter(20) + .validate() + .is_ok()); + } + + #[test] + fn builder_round_trips_to_ffi() { + // Built params must produce the same FFI struct values as the manual setter chain. + let via_builder = IndexParams::builder() + .graph_degree(32) + .intermediate_graph_degree(64) + .nn_descent_niter(20) + .build() + .unwrap(); + let via_setters = IndexParams::new() + .unwrap() + .set_graph_degree(32) + .set_intermediate_graph_degree(64) + .set_nn_descent_niter(20); + unsafe { + assert_eq!((*via_builder.0).graph_degree, (*via_setters.0).graph_degree); + assert_eq!( + (*via_builder.0).intermediate_graph_degree, + (*via_setters.0).intermediate_graph_degree + ); + assert_eq!( + (*via_builder.0).nn_descent_niter, + (*via_setters.0).nn_descent_niter + ); + } + } + + #[test] + fn existing_setter_api_unchanged() { + // Ensure the original API still compiles and sets values correctly. + let params = IndexParams::new() + .unwrap() + .set_graph_degree(32) + .set_intermediate_graph_degree(64) + .set_nn_descent_niter(20); + unsafe { + assert_eq!((*params.0).graph_degree, 32); + assert_eq!((*params.0).intermediate_graph_degree, 64); + assert_eq!((*params.0).nn_descent_niter, 20); + } + } } diff --git a/rust/cuvs/src/cagra/mod.rs b/rust/cuvs/src/cagra/mod.rs index 9043b17386..6b5b9fb74d 100644 --- a/rust/cuvs/src/cagra/mod.rs +++ b/rust/cuvs/src/cagra/mod.rs @@ -94,5 +94,5 @@ mod index_params; mod search_params; pub use index::Index; -pub use index_params::{BuildAlgo, CompressionParams, IndexParams}; -pub use search_params::{HashMode, SearchAlgo, SearchParams}; +pub use index_params::{BuildAlgo, CompressionParams, IndexParams, IndexParamsBuilder}; +pub use search_params::{HashMode, SearchAlgo, SearchParams, SearchParamsBuilder}; diff --git a/rust/cuvs/src/cagra/search_params.rs b/rust/cuvs/src/cagra/search_params.rs index 59537d7718..7914b1dbc2 100644 --- a/rust/cuvs/src/cagra/search_params.rs +++ b/rust/cuvs/src/cagra/search_params.rs @@ -122,6 +122,15 @@ impl SearchParams { } } +impl SearchParams { + /// Returns a builder for constructing [`SearchParams`] with validated parameters. + /// + /// See [`SearchParamsBuilder`] for details. + pub fn builder() -> SearchParamsBuilder { + SearchParamsBuilder::default() + } +} + impl fmt::Debug for SearchParams { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { // custom debug trait here, default value will show the pointer address @@ -143,6 +152,180 @@ impl Drop for SearchParams { } } +/// Builder for [`SearchParams`] with pre-validated parameters. +/// +/// Construct via [`SearchParams::builder()`]. Call [`SearchParamsBuilder::build`] to +/// validate all parameters and allocate the FFI struct in one step. +pub struct SearchParamsBuilder { + itopk_size: usize, + max_queries: usize, + max_iterations: usize, + min_iterations: usize, + team_size: usize, + thread_block_size: usize, + hashmap_max_fill_rate: f32, + hashmap_min_bitlen: usize, + num_random_samplings: u32, + rand_xor_mask: u64, + algo: Option, + hashmap_mode: Option, +} + +impl Default for SearchParamsBuilder { + fn default() -> Self { + Self { + itopk_size: 64, + max_queries: 0, + max_iterations: 0, + min_iterations: 0, + team_size: 0, + thread_block_size: 0, + hashmap_max_fill_rate: 0.5, + hashmap_min_bitlen: 0, + num_random_samplings: 1, + rand_xor_mask: 0x128394, + algo: None, + hashmap_mode: None, + } + } +} + +impl SearchParamsBuilder { + /// Number of intermediate search results retained during the search. + /// + /// Must be a power of 2 (or 0 to use the cuVS default). + pub fn itopk_size(mut self, v: usize) -> Self { + self.itopk_size = v; + self + } + + /// Maximum number of queries to search at the same time. 0 = auto. + pub fn max_queries(mut self, v: usize) -> Self { + self.max_queries = v; + self + } + + /// Upper limit of search iterations. 0 = auto. + pub fn max_iterations(mut self, v: usize) -> Self { + self.max_iterations = v; + self + } + + /// Lower limit of search iterations. + pub fn min_iterations(mut self, v: usize) -> Self { + self.min_iterations = v; + self + } + + /// Number of threads used to calculate a single distance. + /// + /// Must be 0 (auto), 4, 8, 16, or 32. + pub fn team_size(mut self, v: usize) -> Self { + self.team_size = v; + self + } + + /// Thread block size. 0 (auto), 64, 128, 256, 512, or 1024. + pub fn thread_block_size(mut self, v: usize) -> Self { + self.thread_block_size = v; + self + } + + /// Upper limit of hashmap fill rate. + /// + /// Must be in the exclusive range (0.1, 0.9). + pub fn hashmap_max_fill_rate(mut self, v: f32) -> Self { + self.hashmap_max_fill_rate = v; + self + } + + /// Lower limit of hashmap bit length. + pub fn hashmap_min_bitlen(mut self, v: usize) -> Self { + self.hashmap_min_bitlen = v; + self + } + + /// Number of iterations of initial random seed node selection. + pub fn num_random_samplings(mut self, v: u32) -> Self { + self.num_random_samplings = v; + self + } + + /// Bit mask used for initial random seed node selection. + pub fn rand_xor_mask(mut self, v: u64) -> Self { + self.rand_xor_mask = v; + self + } + + /// Which search implementation to use. + pub fn algo(mut self, v: SearchAlgo) -> Self { + self.algo = Some(v); + self + } + + /// Hashmap type. + pub fn hashmap_mode(mut self, v: HashMode) -> Self { + self.hashmap_mode = Some(v); + self + } + + /// Validate all parameters without allocating any GPU resources. + /// + /// Returns `Ok(())` if all parameters are valid, or `Err` with a message naming + /// the offending field and its valid range. + pub fn validate(&self) -> crate::error::Result<()> { + if self.itopk_size != 0 && !self.itopk_size.is_power_of_two() { + return Err(format!( + "itopk_size must be a power of 2 or 0 (auto); got {}", + self.itopk_size + ) + .into()); + } + const VALID_TEAM_SIZES: &[usize] = &[0, 4, 8, 16, 32]; + if !VALID_TEAM_SIZES.contains(&self.team_size) { + return Err(format!( + "team_size must be one of {{0, 4, 8, 16, 32}}; got {}", + self.team_size + ) + .into()); + } + if self.hashmap_max_fill_rate <= 0.1 || self.hashmap_max_fill_rate >= 0.9 { + return Err(format!( + "hashmap_max_fill_rate must be in (0.1, 0.9); got {}", + self.hashmap_max_fill_rate + ) + .into()); + } + Ok(()) + } + + /// Validate all parameters and allocate the FFI struct. + /// + /// Returns `Err` with a message naming the offending field and its valid range + /// before any GPU work begins. + pub fn build(self) -> crate::error::Result { + self.validate()?; + let mut params = SearchParams::new()? + .set_itopk_size(self.itopk_size) + .set_max_queries(self.max_queries) + .set_max_iterations(self.max_iterations) + .set_min_iterations(self.min_iterations) + .set_team_size(self.team_size) + .set_thread_block_size(self.thread_block_size) + .set_hashmap_max_fill_rate(self.hashmap_max_fill_rate) + .set_hashmap_min_bitlen(self.hashmap_min_bitlen) + .set_num_random_samplings(self.num_random_samplings) + .set_rand_xor_mask(self.rand_xor_mask); + if let Some(algo) = self.algo { + params = params.set_algo(algo); + } + if let Some(mode) = self.hashmap_mode { + params = params.set_hashmap_mode(mode); + } + Ok(params) + } +} + #[cfg(test)] mod tests { use super::*; @@ -155,4 +338,67 @@ mod tests { assert_eq!((*params.0).itopk_size, 128); } } + + // --- SearchParamsBuilder tests --- + + #[test] + fn builder_rejects_non_power_of_two_itopk() { + let err = SearchParams::builder() + .itopk_size(100) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("itopk_size"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_rejects_invalid_team_size() { + let err = SearchParams::builder().team_size(7).validate().unwrap_err(); + assert!( + err.to_string().contains("team_size"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_rejects_fill_rate_too_high() { + let err = SearchParams::builder() + .hashmap_max_fill_rate(0.95) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("hashmap_max_fill_rate"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_rejects_fill_rate_too_low() { + let err = SearchParams::builder() + .hashmap_max_fill_rate(0.05) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("hashmap_max_fill_rate"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_accepts_valid_params() { + assert!(SearchParams::builder() + .itopk_size(64) + .team_size(8) + .hashmap_max_fill_rate(0.5) + .validate() + .is_ok()); + } + + #[test] + fn builder_accepts_zero_itopk_as_auto() { + // itopk_size=0 means "auto select" in cuVS — should be valid + assert!(SearchParams::builder().itopk_size(0).validate().is_ok()); + } } diff --git a/rust/cuvs/src/error.rs b/rust/cuvs/src/error.rs index f7b78ec74d..0dbeef5a4e 100644 --- a/rust/cuvs/src/error.rs +++ b/rust/cuvs/src/error.rs @@ -64,3 +64,12 @@ pub fn check_cuda(err: ffi::cudaError_t) -> Result<()> { _ => Err(Error::CudaError(err)), } } + +impl From for Error { + fn from(text: String) -> Self { + Error::CuvsError(CuvsError { + code: ffi::cuvsError_t::CUVS_ERROR, + text, + }) + } +} diff --git a/rust/cuvs/src/ivf_flat/index_params.rs b/rust/cuvs/src/ivf_flat/index_params.rs index 523bc7619e..7eac14e413 100644 --- a/rust/cuvs/src/ivf_flat/index_params.rs +++ b/rust/cuvs/src/ivf_flat/index_params.rs @@ -73,6 +73,15 @@ impl IndexParams { } } +impl IndexParams { + /// Returns a builder for constructing [`IndexParams`] with validated parameters. + /// + /// See [`IndexParamsBuilder`] for details. + pub fn builder() -> IndexParamsBuilder { + IndexParamsBuilder::default() + } +} + impl fmt::Debug for IndexParams { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { // custom debug trait here, default value will show the pointer address @@ -94,6 +103,101 @@ impl Drop for IndexParams { } } +/// Builder for IVF-Flat [`IndexParams`] with pre-validated parameters. +/// +/// Construct via [`IndexParams::builder()`]. Defaults match the cuVS C API defaults. +pub struct IndexParamsBuilder { + n_lists: u32, + metric: Option, + metric_arg: f32, + kmeans_n_iters: u32, + kmeans_trainset_fraction: f64, + add_data_on_build: bool, +} + +impl Default for IndexParamsBuilder { + fn default() -> Self { + Self { + n_lists: 1024, + metric: None, + metric_arg: 2.0, + kmeans_n_iters: 20, + kmeans_trainset_fraction: 0.5, + add_data_on_build: true, + } + } +} + +impl IndexParamsBuilder { + /// The number of clusters used in the coarse quantizer. + /// + /// Must be > 0. + pub fn n_lists(mut self, v: u32) -> Self { + self.n_lists = v; + self + } + + /// DistanceType to use for building the index. + pub fn metric(mut self, v: DistanceType) -> Self { + self.metric = Some(v); + self + } + + /// Metric argument (e.g. p for Minkowski distance). + pub fn metric_arg(mut self, v: f32) -> Self { + self.metric_arg = v; + self + } + + /// Number of iterations searching for kmeans centers during index building. + pub fn kmeans_n_iters(mut self, v: u32) -> Self { + self.kmeans_n_iters = v; + self + } + + /// Fraction of dataset used for kmeans training. Must be in (0, 1]. + pub fn kmeans_trainset_fraction(mut self, v: f64) -> Self { + self.kmeans_trainset_fraction = v; + self + } + + /// Populate the index with the dataset during build. When false, use `extend`. + pub fn add_data_on_build(mut self, v: bool) -> Self { + self.add_data_on_build = v; + self + } + + /// Validate all parameters without allocating any GPU resources. + pub fn validate(&self) -> crate::error::Result<()> { + if self.n_lists == 0 { + return Err(format!("n_lists must be > 0; got {}", self.n_lists).into()); + } + if self.kmeans_trainset_fraction <= 0.0 || self.kmeans_trainset_fraction > 1.0 { + return Err(format!( + "kmeans_trainset_fraction must be in (0, 1]; got {}", + self.kmeans_trainset_fraction + ) + .into()); + } + Ok(()) + } + + /// Validate all parameters and allocate the FFI struct. + pub fn build(self) -> crate::error::Result { + self.validate()?; + let mut params = IndexParams::new()? + .set_n_lists(self.n_lists) + .set_metric_arg(self.metric_arg) + .set_kmeans_n_iters(self.kmeans_n_iters) + .set_kmeans_trainset_fraction(self.kmeans_trainset_fraction) + .set_add_data_on_build(self.add_data_on_build); + if let Some(metric) = self.metric { + params = params.set_metric(metric); + } + Ok(params) + } +} + #[cfg(test)] mod tests { use super::*; @@ -110,4 +214,22 @@ mod tests { assert_eq!((*params.0).add_data_on_build, false); } } + + #[test] + fn builder_rejects_zero_n_lists() { + let err = IndexParams::builder().n_lists(0).validate().unwrap_err(); + assert!( + err.to_string().contains("n_lists"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_accepts_valid_params() { + assert!(IndexParams::builder() + .n_lists(256) + .kmeans_trainset_fraction(0.5) + .validate() + .is_ok()); + } } diff --git a/rust/cuvs/src/ivf_flat/mod.rs b/rust/cuvs/src/ivf_flat/mod.rs index 7417116965..492939c69e 100644 --- a/rust/cuvs/src/ivf_flat/mod.rs +++ b/rust/cuvs/src/ivf_flat/mod.rs @@ -72,5 +72,5 @@ mod index_params; mod search_params; pub use index::Index; -pub use index_params::IndexParams; +pub use index_params::{IndexParams, IndexParamsBuilder}; pub use search_params::SearchParams; diff --git a/rust/cuvs/src/ivf_pq/index_params.rs b/rust/cuvs/src/ivf_pq/index_params.rs index e1f2d53656..7e276d3d4a 100644 --- a/rust/cuvs/src/ivf_pq/index_params.rs +++ b/rust/cuvs/src/ivf_pq/index_params.rs @@ -149,6 +149,15 @@ impl IndexParams { } } +impl IndexParams { + /// Returns a builder for constructing [`IndexParams`] with validated parameters. + /// + /// See [`IndexParamsBuilder`] for details. + pub fn builder() -> IndexParamsBuilder { + IndexParamsBuilder::default() + } +} + impl fmt::Debug for IndexParams { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { // custom debug trait here, default value will show the pointer address @@ -170,6 +179,159 @@ impl Drop for IndexParams { } } +/// Builder for IVF-PQ [`IndexParams`] with pre-validated parameters. +/// +/// Construct via [`IndexParams::builder()`]. Defaults match the cuVS C API defaults. +pub struct IndexParamsBuilder { + n_lists: u32, + metric: Option, + metric_arg: f32, + kmeans_n_iters: u32, + kmeans_trainset_fraction: f64, + pq_bits: u32, + pq_dim: u32, + codebook_kind: Option, + codes_layout: Option, + force_random_rotation: bool, + max_train_points_per_pq_code: u32, + add_data_on_build: bool, +} + +impl Default for IndexParamsBuilder { + fn default() -> Self { + Self { + n_lists: 1024, + metric: None, + metric_arg: 2.0, + kmeans_n_iters: 20, + kmeans_trainset_fraction: 0.5, + pq_bits: 8, + pq_dim: 0, + codebook_kind: None, + codes_layout: None, + force_random_rotation: false, + max_train_points_per_pq_code: 256, + add_data_on_build: true, + } + } +} + +impl IndexParamsBuilder { + /// The number of clusters used in the coarse quantizer. + /// + /// Must be > 0. + pub fn n_lists(mut self, v: u32) -> Self { + self.n_lists = v; + self + } + + /// DistanceType to use for building the index. + pub fn metric(mut self, v: crate::distance_type::DistanceType) -> Self { + self.metric = Some(v); + self + } + + /// Metric argument (e.g. p for Minkowski distance). + pub fn metric_arg(mut self, v: f32) -> Self { + self.metric_arg = v; + self + } + + /// Number of iterations searching for kmeans centers during index building. + pub fn kmeans_n_iters(mut self, v: u32) -> Self { + self.kmeans_n_iters = v; + self + } + + /// Fraction of dataset used for kmeans training. Must be in (0, 1]. + pub fn kmeans_trainset_fraction(mut self, v: f64) -> Self { + self.kmeans_trainset_fraction = v; + self + } + + /// Bit length of the vector element after quantization. Typically 4 or 8. + pub fn pq_bits(mut self, v: u32) -> Self { + self.pq_bits = v; + self + } + + /// Dimensionality of the vector after product quantization. 0 = auto. + pub fn pq_dim(mut self, v: u32) -> Self { + self.pq_dim = v; + self + } + + /// Codebook generation method. + pub fn codebook_kind(mut self, v: cuvsIvfPqCodebookGen) -> Self { + self.codebook_kind = Some(v); + self + } + + /// Memory layout of IVF-PQ list data. + pub fn codes_layout(mut self, v: cuvsIvfPqListLayout) -> Self { + self.codes_layout = Some(v); + self + } + + /// Apply a random rotation matrix on input data and queries. + pub fn force_random_rotation(mut self, v: bool) -> Self { + self.force_random_rotation = v; + self + } + + /// Max number of data points per PQ code during codebook training. + pub fn max_train_points_per_pq_code(mut self, v: u32) -> Self { + self.max_train_points_per_pq_code = v; + self + } + + /// Populate the index with the dataset during build. When false, use `extend`. + pub fn add_data_on_build(mut self, v: bool) -> Self { + self.add_data_on_build = v; + self + } + + /// Validate all parameters without allocating any GPU resources. + pub fn validate(&self) -> crate::error::Result<()> { + if self.n_lists == 0 { + return Err(format!("n_lists must be > 0; got {}", self.n_lists).into()); + } + if self.kmeans_trainset_fraction <= 0.0 || self.kmeans_trainset_fraction > 1.0 { + return Err(format!( + "kmeans_trainset_fraction must be in (0, 1]; got {}", + self.kmeans_trainset_fraction + ) + .into()); + } + Ok(()) + } + + /// Validate all parameters and allocate the FFI struct. + pub fn build(self) -> crate::error::Result { + self.validate()?; + let mut params = IndexParams::new()? + .set_n_lists(self.n_lists) + .set_metric_arg(self.metric_arg) + .set_kmeans_n_iters(self.kmeans_n_iters) + .set_kmeans_trainset_fraction(self.kmeans_trainset_fraction) + .set_pq_bits(self.pq_bits) + .set_pq_dim(self.pq_dim) + .set_force_random_rotation(self.force_random_rotation) + .set_max_train_points_per_pq_code(self.max_train_points_per_pq_code) + .set_add_data_on_build(self.add_data_on_build); + if let Some(metric) = self.metric { + params = params.set_metric(metric); + } + if let Some(kind) = self.codebook_kind { + params = params.set_codebook_kind(kind); + } + if let Some(layout) = self.codes_layout { + params = params.set_codes_layout(layout); + } + Ok(params) + } +} + #[cfg(test)] mod tests { use super::*; @@ -186,4 +348,34 @@ mod tests { assert_eq!((*params.0).add_data_on_build, false); } } + + #[test] + fn builder_rejects_zero_n_lists() { + let err = IndexParams::builder().n_lists(0).validate().unwrap_err(); + assert!( + err.to_string().contains("n_lists"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_accepts_valid_params() { + assert!(IndexParams::builder() + .n_lists(256) + .kmeans_trainset_fraction(0.5) + .validate() + .is_ok()); + } + + #[test] + fn existing_setter_api_unchanged() { + let params = IndexParams::new() + .unwrap() + .set_n_lists(128) + .set_add_data_on_build(false); + unsafe { + assert_eq!((*params.0).n_lists, 128); + assert_eq!((*params.0).add_data_on_build, false); + } + } } diff --git a/rust/cuvs/src/ivf_pq/mod.rs b/rust/cuvs/src/ivf_pq/mod.rs index c4676cd1aa..aaa3626984 100644 --- a/rust/cuvs/src/ivf_pq/mod.rs +++ b/rust/cuvs/src/ivf_pq/mod.rs @@ -69,5 +69,5 @@ mod index_params; mod search_params; pub use index::Index; -pub use index_params::IndexParams; +pub use index_params::{IndexParams, IndexParamsBuilder}; pub use search_params::SearchParams; diff --git a/rust/cuvs/src/vamana/index_params.rs b/rust/cuvs/src/vamana/index_params.rs index c52c287238..3138390a96 100644 --- a/rust/cuvs/src/vamana/index_params.rs +++ b/rust/cuvs/src/vamana/index_params.rs @@ -96,6 +96,15 @@ impl IndexParams { } } +impl IndexParams { + /// Returns a builder for constructing [`IndexParams`] with validated parameters. + /// + /// See [`IndexParamsBuilder`] for details. + pub fn builder() -> IndexParamsBuilder { + IndexParamsBuilder::default() + } +} + impl fmt::Debug for IndexParams { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { // custom debug trait here, default value will show the pointer address @@ -117,6 +126,135 @@ impl Drop for IndexParams { } } +/// Builder for Vamana [`IndexParams`] with pre-validated parameters. +/// +/// Construct via [`IndexParams::builder()`]. Defaults match the cuVS C API defaults. +pub struct IndexParamsBuilder { + graph_degree: u32, + visited_size: u32, + vamana_iters: f32, + alpha: f32, + max_fraction: f32, + batch_base: f32, + queue_size: u32, + reverse_batchsize: u32, + metric: Option, +} + +impl Default for IndexParamsBuilder { + fn default() -> Self { + Self { + graph_degree: 64, + visited_size: 75, + vamana_iters: 2.0, + alpha: 1.2, + max_fraction: 0.06, + batch_base: 2.0, + queue_size: 255, + reverse_batchsize: 1_000_000, + metric: None, + } + } +} + +impl IndexParamsBuilder { + /// Maximum degree of the output graph (R parameter in Vamana literature). + /// + /// Must be > 0. + pub fn graph_degree(mut self, v: u32) -> Self { + self.graph_degree = v; + self + } + + /// Maximum number of visited nodes per search (L parameter in Vamana literature). + /// + /// Must be >= `graph_degree`. + pub fn visited_size(mut self, v: u32) -> Self { + self.visited_size = v; + self + } + + /// Number of Vamana vector insertion iterations. + pub fn vamana_iters(mut self, v: f32) -> Self { + self.vamana_iters = v; + self + } + + /// Alpha for pruning parameter. + /// + /// Must be > 0. + pub fn alpha(mut self, v: f32) -> Self { + self.alpha = v; + self + } + + /// Maximum fraction of dataset inserted per batch. + pub fn max_fraction(mut self, v: f32) -> Self { + self.max_fraction = v; + self + } + + /// Base of growth rate of batch sizes. + pub fn batch_base(mut self, v: f32) -> Self { + self.batch_base = v; + self + } + + /// Size of candidate queue structure. + pub fn queue_size(mut self, v: u32) -> Self { + self.queue_size = v; + self + } + + /// Max batchsize of reverse edge processing. + pub fn reverse_batchsize(mut self, v: u32) -> Self { + self.reverse_batchsize = v; + self + } + + /// DistanceType to use for building the index. + pub fn metric(mut self, v: DistanceType) -> Self { + self.metric = Some(v); + self + } + + /// Validate all parameters without allocating any GPU resources. + pub fn validate(&self) -> crate::error::Result<()> { + if self.graph_degree == 0 { + return Err(format!("graph_degree must be > 0; got {}", self.graph_degree).into()); + } + if self.visited_size < self.graph_degree { + return Err(format!( + "visited_size ({}) must be >= graph_degree ({})", + self.visited_size, self.graph_degree + ) + .into()); + } + if self.alpha <= 0.0 { + return Err(format!("alpha must be > 0; got {}", self.alpha).into()); + } + Ok(()) + } + + /// Validate all parameters and allocate the FFI struct. + pub fn build(self) -> crate::error::Result { + self.validate()?; + let mut params = IndexParams::new()? + .set_graph_degree(self.graph_degree) + .set_visited_size(self.visited_size) + .set_vamana_iters(self.vamana_iters) + .set_alpha(self.alpha) + .set_max_fraction(self.max_fraction) + .set_batch_base(self.batch_base) + .set_queue_size(self.queue_size) + .set_reverse_batchsize(self.reverse_batchsize); + if let Some(metric) = self.metric { + params = params.set_metric(metric); + } + Ok(params) + } +} + #[cfg(test)] mod tests { use super::*; @@ -133,4 +271,39 @@ mod tests { assert_eq!((*params.0).visited_size, 128); } } + + #[test] + fn builder_rejects_zero_graph_degree() { + let err = IndexParams::builder() + .graph_degree(0) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("graph_degree"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_rejects_visited_size_less_than_graph_degree() { + let err = IndexParams::builder() + .graph_degree(64) + .visited_size(32) + .validate() + .unwrap_err(); + assert!( + err.to_string().contains("visited_size"), + "error message should name the field: {err}" + ); + } + + #[test] + fn builder_accepts_valid_params() { + assert!(IndexParams::builder() + .graph_degree(32) + .visited_size(75) + .alpha(1.2) + .validate() + .is_ok()); + } } diff --git a/rust/cuvs/src/vamana/mod.rs b/rust/cuvs/src/vamana/mod.rs index a3ae4ee9ff..631dd37d0e 100644 --- a/rust/cuvs/src/vamana/mod.rs +++ b/rust/cuvs/src/vamana/mod.rs @@ -8,4 +8,4 @@ mod index; mod index_params; pub use index::Index; -pub use index_params::IndexParams; +pub use index_params::{IndexParams, IndexParamsBuilder};