query_planning/plan/vector_algorithm/
definition.rs

1//! cuVS operation definitions and stable vector algorithm identities.
2
3use std::fmt;
4
5use serde_json::{Map, Value, json};
6
7use crate::plan::{
8    GpuOptionDefinition, GpuOptionKind, GpuOptionValidationError, GpuOutputFieldDefinition, VectorElementType,
9    gpu_function_options::{
10        canonicalize_options, dense_relation_binding_schema, option_schema as gpu_option_schema,
11        validate_dense_relation_binding_value,
12    },
13};
14
15/// Public SQL names for the query-local cuVS operations.
16pub const CUVS_BRUTE_FORCE_KNN_FUNCTION: &str = "cuvs_brute_force_knn";
17/// Canonical SQL function name for IVF-Flat nearest-neighbor search.
18pub const CUVS_IVF_FLAT_FUNCTION: &str = "cuvs_ivf_flat";
19/// Canonical SQL function name for IVF-PQ nearest-neighbor search.
20pub const CUVS_IVF_PQ_FUNCTION: &str = "cuvs_ivf_pq";
21/// Canonical SQL function name for CAGRA nearest-neighbor search.
22pub const CUVS_CAGRA_FUNCTION: &str = "cuvs_cagra";
23/// Canonical SQL function name for KMeans fit and assignment.
24pub const CUVS_KMEANS_FUNCTION: &str = "cuvs_kmeans";
25/// Canonical SQL function name for PCA fit-transform.
26pub const CUVS_PCA_FUNCTION: &str = "cuvs_pca";
27/// Canonical SQL function name for pairwise dense-vector distances.
28pub const CUVS_PAIRWISE_DISTANCE_FUNCTION: &str = "cuvs_pairwise_distance";
29/// Canonical SQL function name for selection of values within each vector.
30pub const CUVS_TOP_K_FUNCTION: &str = "cuvs_top_k";
31
32/// SQL relation argument roles in the canonical operation registry.
33#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
34pub enum CuvsRelationRole {
35    /// Dense rows indexed for nearest-neighbor search.
36    Dataset,
37    /// Dense query rows matched against the indexed dataset.
38    Queries,
39    /// Dense rows consumed by a single-relation operation.
40    Input,
41    /// Rows on the first axis of a pairwise distance result.
42    X,
43    /// Rows on the second axis of a pairwise distance result.
44    Y,
45}
46
47impl CuvsRelationRole {
48    /// Canonical lowercase relation argument name used by the SQL operation.
49    pub const fn as_str(self) -> &'static str {
50        match self {
51            Self::Dataset => "dataset",
52            Self::Queries => "queries",
53            Self::Input => "input",
54            Self::X => "x",
55            Self::Y => "y",
56        }
57    }
58
59    /// Canonical documentation description for this logical relation role.
60    pub const fn description(self) -> &'static str {
61        match self {
62            Self::Dataset => "Dense-vector rows indexed for nearest-neighbor search.",
63            Self::Queries => "Dense-vector query rows matched against the evaluated dataset.",
64            Self::Input => "Dense-vector rows consumed by this operation.",
65            Self::X => "Dense-vector rows on the first axis of the pairwise distance result.",
66            Self::Y => "Dense-vector rows on the second axis of the pairwise distance result.",
67        }
68    }
69}
70
71/// Stable relational output contract selected by an operation descriptor.
72#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
73pub enum CuvsOutputContract {
74    /// One row per neighbor with query and dataset identifiers, rank, and metric value.
75    KnnNeighbors,
76    /// Input identifiers and query-local cluster assignments.
77    KMeansAssignments,
78    /// Input identifiers followed by transformed principal components.
79    PcaComponents,
80    /// One row per pair of input rows with both identifiers and the metric value.
81    PairwiseDistances,
82    /// Selected values with input identifiers, rank, and vector dimension index.
83    TopKValues,
84}
85
86/// Verified metrics for one vector element type within an operation.
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub struct CuvsVectorElementSupport {
89    /// Dense vector element type supported by the operation.
90    pub element: VectorElementType,
91    /// Canonical metric option values supported for this element type; empty for operations without a metric.
92    pub metrics: &'static [&'static str],
93}
94
95/// One operation entry drives SQL names, relation shape, option metadata, and
96/// output-schema selection. Keep operation-specific facts in this registry.
97#[derive(Debug, Clone, Copy, PartialEq)]
98pub struct CuvsOperationDescriptor {
99    /// Identity selecting this registry entry.
100    pub operation: CuvsOperation,
101    /// Canonical table-function identifier.
102    pub sql_function_name: &'static str,
103    /// Human-readable function title; SQL names remain canonical identifiers.
104    pub display_name: &'static str,
105    /// Short catalog summary.
106    pub summary: &'static str,
107    /// Concrete selection guidance for the generated reference page.
108    pub use_this_when: &'static str,
109    /// Relation arguments, each passed as a parenthesized subquery; the SQL
110    /// surface admits no other relation syntax.
111    pub relation_roles: &'static [CuvsRelationRole],
112    /// Supported dense-vector element types for every relation binding.
113    pub vector_element_support: &'static [CuvsVectorElementSupport],
114    /// Result-shape family used to construct the Arrow schema.
115    pub output_contract: CuvsOutputContract,
116    /// Required options, accepted values, and defaults.
117    pub options: &'static [GpuOptionDefinition],
118    /// Ordered output field metadata, including repeated-field templates.
119    pub output_fields: &'static [GpuOutputFieldDefinition],
120    /// Generic SQL that demonstrates the real execution syntax beyond the
121    /// metadata-validation envelope.
122    pub quickstart_sql: &'static str,
123    /// Static and data-dependent lifecycle constraints that apply to this
124    /// operation in addition to the common relation/validation contract.
125    pub limitations: &'static [&'static str],
126}
127
128const FLOAT32_VECTOR_SUPPORT: &[CuvsVectorElementSupport] =
129    &[CuvsVectorElementSupport { element: VectorElementType::Float32, metrics: &[] }];
130const KNN_RELATIONS: &[CuvsRelationRole] = &[CuvsRelationRole::Dataset, CuvsRelationRole::Queries];
131const INPUT_RELATION: &[CuvsRelationRole] = &[CuvsRelationRole::Input];
132const PAIRWISE_RELATIONS: &[CuvsRelationRole] = &[CuvsRelationRole::X, CuvsRelationRole::Y];
133const TOP_K_SELECTIONS: &[&str] = &["min", "max"];
134const KNN_METRICS: &[&str] = &["l2_expanded", "l2_sqrt_expanded", "cosine", "inner_product"];
135const KMEANS_METRICS: &[&str] = &["l2_expanded", "l2_sqrt_expanded"];
136const IVF_PQ_INTEGER_METRICS: &[&str] = &["l2_expanded", "l2_sqrt_expanded", "inner_product"];
137const BRUTE_FORCE_SUPPORT: &[CuvsVectorElementSupport] =
138    &[CuvsVectorElementSupport { element: VectorElementType::Float32, metrics: KNN_METRICS }];
139const IVF_SUPPORT: &[CuvsVectorElementSupport] = &[
140    CuvsVectorElementSupport { element: VectorElementType::Float32, metrics: KNN_METRICS },
141    CuvsVectorElementSupport { element: VectorElementType::Int8, metrics: KNN_METRICS },
142    CuvsVectorElementSupport { element: VectorElementType::UInt8, metrics: KNN_METRICS },
143];
144const IVF_PQ_SUPPORT: &[CuvsVectorElementSupport] = &[
145    CuvsVectorElementSupport { element: VectorElementType::Float32, metrics: KNN_METRICS },
146    CuvsVectorElementSupport { element: VectorElementType::Int8, metrics: IVF_PQ_INTEGER_METRICS },
147    CuvsVectorElementSupport { element: VectorElementType::UInt8, metrics: IVF_PQ_INTEGER_METRICS },
148];
149const CAGRA_SUPPORT: &[CuvsVectorElementSupport] = &[
150    CuvsVectorElementSupport { element: VectorElementType::Float32, metrics: CAGRA_METRICS },
151    CuvsVectorElementSupport { element: VectorElementType::Int8, metrics: CAGRA_METRICS },
152    CuvsVectorElementSupport { element: VectorElementType::UInt8, metrics: CAGRA_METRICS },
153];
154const PAIRWISE_SUPPORT: &[CuvsVectorElementSupport] = &[
155    CuvsVectorElementSupport { element: VectorElementType::Float32, metrics: KNN_METRICS },
156    CuvsVectorElementSupport { element: VectorElementType::Float64, metrics: KNN_METRICS },
157];
158const KMEANS_SUPPORT: &[CuvsVectorElementSupport] = &[
159    CuvsVectorElementSupport { element: VectorElementType::Float32, metrics: KMEANS_METRICS },
160    CuvsVectorElementSupport { element: VectorElementType::Float64, metrics: KMEANS_METRICS },
161];
162const KMEANS_INIT: &[&str] = &["kmeans++", "random"];
163const PCA_SOLVERS: &[&str] = &["cov_eig_dq", "cov_eig_jacobi"];
164
165const KNN_OPTIONS: &[GpuOptionDefinition] = &[
166    GpuOptionDefinition {
167        name: "dataset",
168        required: true,
169        kind: GpuOptionKind::RelationBinding,
170        description: "Names the dataset ID column and one dense-vector binding shape.",
171    },
172    GpuOptionDefinition {
173        name: "queries",
174        required: true,
175        kind: GpuOptionKind::RelationBinding,
176        description: "Names the query ID column and one dense-vector binding shape.",
177    },
178    GpuOptionDefinition {
179        name: "k",
180        required: true,
181        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: None },
182        description: "Number of neighbors returned for each evaluated query row.",
183    },
184    GpuOptionDefinition {
185        name: "metric",
186        required: true,
187        kind: GpuOptionKind::Enum { values: KNN_METRICS, default: None },
188        description: "Distance or score metric. Distance metrics rank lower values first; inner product ranks higher values first.",
189    },
190];
191const CAGRA_METRICS: &[&str] = &["l2_expanded", "inner_product", "cosine"];
192const CAGRA_BUILD_ALGOS: &[&str] = &["ivf_pq", "nn_descent"];
193const CAGRA_SEARCH_ALGOS: &[&str] = &["auto", "single_cta", "multi_cta", "multi_kernel"];
194const CAGRA_OPTIONS: &[GpuOptionDefinition] = &[
195    KNN_OPTIONS[0],
196    KNN_OPTIONS[1],
197    KNN_OPTIONS[2],
198    GpuOptionDefinition {
199        name: "metric",
200        required: false,
201        kind: GpuOptionKind::Enum { values: CAGRA_METRICS, default: Some("l2_expanded") },
202        description: "Distance or score metric; native default l2_expanded. inner_product ranks higher scores first.",
203    },
204    GpuOptionDefinition {
205        name: "graph_degree",
206        required: false,
207        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: Some(64) },
208        description: "Output graph degree; native default 64. Must be less than dataset rows.",
209    },
210    GpuOptionDefinition {
211        name: "intermediate_graph_degree",
212        required: false,
213        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: Some(128) },
214        description: "Build graph degree; native default 128. Must exceed graph_degree and be less than dataset rows.",
215    },
216    GpuOptionDefinition {
217        name: "build_algo",
218        required: false,
219        kind: GpuOptionKind::Enum { values: CAGRA_BUILD_ALGOS, default: Some("ivf_pq") },
220        description: "Graph construction algorithm; native default ivf_pq. nn_descent does not support cosine.",
221    },
222    GpuOptionDefinition {
223        name: "itopk_size",
224        required: false,
225        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: Some(64) },
226        description: "Intermediate search results; native default 64. Must be at least k.",
227    },
228    GpuOptionDefinition {
229        name: "search_width",
230        required: false,
231        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: Some(1) },
232        description: "Starting graph nodes per iteration; native default 1.",
233    },
234    GpuOptionDefinition {
235        name: "max_iterations",
236        required: false,
237        kind: GpuOptionKind::Integer { minimum: 0, maximum: u32::MAX as u64, default: Some(0) },
238        description: "Maximum search iterations; native default 0 selects automatically.",
239    },
240    GpuOptionDefinition {
241        name: "search_algo",
242        required: false,
243        kind: GpuOptionKind::Enum { values: CAGRA_SEARCH_ALGOS, default: Some("auto") },
244        description: "Search kernel; auto is set explicitly (the native zero-initialized default is single_cta).",
245    },
246];
247const IVF_FLAT_OPTIONS: &[GpuOptionDefinition] = &[
248    KNN_OPTIONS[0],
249    KNN_OPTIONS[1],
250    KNN_OPTIONS[2],
251    GpuOptionDefinition {
252        name: "metric",
253        required: false,
254        kind: GpuOptionKind::Enum { values: KNN_METRICS, default: Some("l2_expanded") },
255        description: "Distance or score metric; native default l2_expanded. Distance metrics rank lower values first; inner product ranks higher values first.",
256    },
257    GpuOptionDefinition {
258        name: "n_lists",
259        required: true,
260        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: None },
261        description: "Number of inverted lists; must not exceed the evaluated dataset row count.",
262    },
263    GpuOptionDefinition {
264        name: "n_probes",
265        required: false,
266        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: Some(20) },
267        description: "Lists probed per query; native default 20, and must not exceed n_lists.",
268    },
269    GpuOptionDefinition {
270        name: "kmeans_n_iters",
271        required: false,
272        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: Some(20) },
273        description: "Iterations for training list centers; native default 20.",
274    },
275    GpuOptionDefinition {
276        name: "kmeans_trainset_fraction",
277        required: false,
278        kind: GpuOptionKind::Fraction { default: 0.5 },
279        description: "Fraction of the dataset used for training centers; native default 0.5.",
280    },
281];
282const IVF_PQ_CODEBOOK_KINDS: &[&str] = &["per_subspace", "per_cluster"];
283const IVF_PQ_OPTIONS: &[GpuOptionDefinition] = &[
284    IVF_FLAT_OPTIONS[0],
285    IVF_FLAT_OPTIONS[1],
286    IVF_FLAT_OPTIONS[2],
287    IVF_FLAT_OPTIONS[3],
288    IVF_FLAT_OPTIONS[4],
289    IVF_FLAT_OPTIONS[5],
290    IVF_FLAT_OPTIONS[6],
291    IVF_FLAT_OPTIONS[7],
292    GpuOptionDefinition {
293        name: "pq_dim",
294        required: false,
295        kind: GpuOptionKind::Integer { minimum: 0, maximum: u32::MAX as u64, default: Some(0) },
296        description: "Compressed dimensions; native default 0 selects the dimension automatically. pq_dim * pq_bits must be divisible by 8.",
297    },
298    GpuOptionDefinition {
299        name: "pq_bits",
300        required: false,
301        kind: GpuOptionKind::Integer { minimum: 4, maximum: 8, default: Some(8) },
302        description: "Bits per PQ code; native default 8.",
303    },
304    GpuOptionDefinition {
305        name: "codebook_kind",
306        required: false,
307        kind: GpuOptionKind::Enum { values: IVF_PQ_CODEBOOK_KINDS, default: Some("per_subspace") },
308        description: "PQ codebook training scope; native default per_subspace.",
309    },
310];
311const PAIRWISE_OPTIONS: &[GpuOptionDefinition] = &[
312    GpuOptionDefinition {
313        name: "x",
314        required: true,
315        kind: GpuOptionKind::RelationBinding,
316        description: "First dense-vector relation.",
317    },
318    GpuOptionDefinition {
319        name: "y",
320        required: true,
321        kind: GpuOptionKind::RelationBinding,
322        description: "Second dense-vector relation.",
323    },
324    GpuOptionDefinition {
325        name: "metric",
326        required: true,
327        kind: GpuOptionKind::Enum { values: KNN_METRICS, default: None },
328        description: "Distance metric; inner_product is a similarity where larger is closer.",
329    },
330];
331const TOP_K_OPTIONS: &[GpuOptionDefinition] = &[
332    GpuOptionDefinition {
333        name: "input",
334        required: true,
335        kind: GpuOptionKind::RelationBinding,
336        description: "Dense-vector rows to select from.",
337    },
338    GpuOptionDefinition {
339        name: "k",
340        required: true,
341        kind: GpuOptionKind::Integer { minimum: 1, maximum: u32::MAX as u64, default: None },
342        description: "Number of values selected from each row.",
343    },
344    GpuOptionDefinition {
345        name: "select",
346        required: true,
347        kind: GpuOptionKind::Enum { values: TOP_K_SELECTIONS, default: None },
348        description: "Select the minimum or maximum values.",
349    },
350];
351const DENSE_INPUT_OPTION: GpuOptionDefinition = GpuOptionDefinition {
352    name: "input",
353    required: true,
354    kind: GpuOptionKind::RelationBinding,
355    description: "Names the input ID column and one dense-vector binding shape.",
356};
357const KMEANS_OPTIONS: &[GpuOptionDefinition] = &[
358    DENSE_INPUT_OPTION,
359    GpuOptionDefinition {
360        name: "n_clusters",
361        required: true,
362        kind: GpuOptionKind::Integer { minimum: 1, maximum: i32::MAX as u64, default: None },
363        description: "Number of clusters to fit.",
364    },
365    GpuOptionDefinition {
366        name: "metric",
367        required: false,
368        kind: GpuOptionKind::Enum { values: KMEANS_METRICS, default: Some("l2_expanded") },
369        description: "L2 distance form used while fitting and assigning the current input relation.",
370    },
371    GpuOptionDefinition {
372        name: "max_iter",
373        required: false,
374        kind: GpuOptionKind::Integer { minimum: 1, maximum: i32::MAX as u64, default: Some(100) },
375        description: "Maximum fitting iterations for this invocation.",
376    },
377    GpuOptionDefinition {
378        name: "tol",
379        required: false,
380        kind: GpuOptionKind::NonnegativeNumber { default: 0.0001 },
381        description: "Non-negative convergence tolerance used by the fitted model.",
382    },
383    GpuOptionDefinition {
384        name: "n_init",
385        required: false,
386        kind: GpuOptionKind::Integer { minimum: 1, maximum: i32::MAX as u64, default: Some(1) },
387        description: "Number of initialization attempts performed within this call.",
388    },
389    GpuOptionDefinition {
390        name: "init",
391        required: false,
392        kind: GpuOptionKind::Enum { values: KMEANS_INIT, default: Some("kmeans++") },
393        description: "Centroid initialization strategy for this fitted model.",
394    },
395];
396const PCA_OPTIONS: &[GpuOptionDefinition] = &[
397    DENSE_INPUT_OPTION,
398    GpuOptionDefinition {
399        name: "n_components",
400        required: true,
401        kind: GpuOptionKind::Integer { minimum: 1, maximum: i32::MAX as u64, default: None },
402        description: "Number of principal-component columns emitted as pc_0 through pc_<n_components - 1>.",
403    },
404    GpuOptionDefinition {
405        name: "whiten",
406        required: false,
407        kind: GpuOptionKind::Boolean { default: Some(false) },
408        description: "Whether the query-local transformed components are whitened.",
409    },
410    GpuOptionDefinition {
411        name: "solver",
412        required: false,
413        kind: GpuOptionKind::Enum { values: PCA_SOLVERS, default: Some("cov_eig_dq") },
414        description: "Covariance eigensolver used for this fit-transform call.",
415    },
416    GpuOptionDefinition {
417        name: "tol",
418        required: false,
419        kind: GpuOptionKind::NonnegativeNumber { default: 0.0 },
420        description: "Non-negative eigensolver tolerance.",
421    },
422    GpuOptionDefinition {
423        name: "n_iterations",
424        required: false,
425        kind: GpuOptionKind::Integer { minimum: 1, maximum: i32::MAX as u64, default: Some(15) },
426        description: "Maximum eigensolver iterations for this invocation.",
427    },
428    GpuOptionDefinition {
429        name: "flip_signs_based_on_u",
430        required: false,
431        kind: GpuOptionKind::Boolean { default: Some(false) },
432        description: "Whether cuVS applies its U-based component-sign convention.",
433    },
434];
435
436const KNN_OUTPUT_FIELDS: &[GpuOutputFieldDefinition] = &[
437    GpuOutputFieldDefinition {
438        name: "query_ordinal",
439        data_type: "UInt64",
440        nullable: false,
441        repeats: None,
442        description: "Zero-based ordinal of the evaluated query row; it disambiguates duplicate query IDs.",
443    },
444    GpuOutputFieldDefinition {
445        name: "query_id",
446        data_type: "same_as_queries.id",
447        nullable: false,
448        repeats: None,
449        description: "Logical ID copied from the queries relation.",
450    },
451    GpuOutputFieldDefinition {
452        name: "neighbor_ordinal",
453        data_type: "UInt64",
454        nullable: false,
455        repeats: None,
456        description: "Zero-based ordinal of the matched dataset row; it disambiguates duplicate dataset IDs.",
457    },
458    GpuOutputFieldDefinition {
459        name: "neighbor_id",
460        data_type: "same_as_dataset.id",
461        nullable: false,
462        repeats: None,
463        description: "Logical ID copied from the matched dataset row.",
464    },
465    GpuOutputFieldDefinition {
466        name: "rank",
467        data_type: "UInt32",
468        nullable: false,
469        repeats: None,
470        description: "One-based neighbor rank within a query. Order consumers explicitly by query_ordinal, rank.",
471    },
472    GpuOutputFieldDefinition {
473        name: "distance",
474        data_type: "Float32",
475        nullable: false,
476        repeats: None,
477        description: "Metric value; smaller is better for distance metrics, while inner_product prefers larger values.",
478    },
479];
480
481const PAIRWISE_OUTPUT_FIELDS: &[GpuOutputFieldDefinition] = &[
482    GpuOutputFieldDefinition {
483        name: "x_ordinal",
484        data_type: "UInt64",
485        nullable: false,
486        repeats: None,
487        description: "Zero-based physical ordinal of the evaluated x row; it disambiguates duplicate x IDs.",
488    },
489    GpuOutputFieldDefinition {
490        name: "x_id",
491        data_type: "same_as_x.id",
492        nullable: false,
493        repeats: None,
494        description: "Logical ID copied from the x relation.",
495    },
496    GpuOutputFieldDefinition {
497        name: "y_ordinal",
498        data_type: "UInt64",
499        nullable: false,
500        repeats: None,
501        description: "Zero-based physical ordinal of the evaluated y row; it disambiguates duplicate y IDs.",
502    },
503    GpuOutputFieldDefinition {
504        name: "y_id",
505        data_type: "same_as_y.id",
506        nullable: false,
507        repeats: None,
508        description: "Logical ID copied from the y relation.",
509    },
510    GpuOutputFieldDefinition {
511        name: "distance",
512        data_type: "same_as_x.vector_element",
513        nullable: false,
514        repeats: None,
515        description: "Float32 for Float32 input, Float64 for Float64 input. Smaller is closer for distance metrics; inner_product prefers larger values.",
516    },
517];
518const TOP_K_OUTPUT_FIELDS: &[GpuOutputFieldDefinition] = &[
519    GpuOutputFieldDefinition {
520        name: "row_ordinal",
521        data_type: "UInt64",
522        nullable: false,
523        repeats: None,
524        description: "Zero-based physical ordinal of the evaluated input row; it disambiguates duplicate input IDs. Follows the input's top-level ORDER BY on output columns; unspecified without one.",
525    },
526    GpuOutputFieldDefinition {
527        name: "id",
528        data_type: "same_as_input.id",
529        nullable: false,
530        repeats: None,
531        description: "Logical ID copied from the input relation.",
532    },
533    GpuOutputFieldDefinition {
534        name: "rank",
535        data_type: "UInt32",
536        nullable: false,
537        repeats: None,
538        description: "One-based value rank within an input row. Order consumers explicitly by row_ordinal, rank.",
539    },
540    GpuOutputFieldDefinition {
541        name: "position",
542        data_type: "Int64",
543        nullable: false,
544        repeats: None,
545        description: "Zero-based dimension index of the selected value within the input vector.",
546    },
547    GpuOutputFieldDefinition {
548        name: "value",
549        data_type: "Float32",
550        nullable: false,
551        repeats: None,
552        description: "Selected Float32 input vector value; rank follows ascending values for min and descending values for max.",
553    },
554];
555const KMEANS_OUTPUT_FIELDS: &[GpuOutputFieldDefinition] = &[
556    GpuOutputFieldDefinition {
557        name: "row_ordinal",
558        data_type: "UInt64",
559        nullable: false,
560        repeats: None,
561        description: "Zero-based ordinal of the evaluated input row; it disambiguates duplicate IDs. Follows the input's top-level ORDER BY on output columns; unspecified without one.",
562    },
563    GpuOutputFieldDefinition {
564        name: "id",
565        data_type: "same_as_input.id",
566        nullable: false,
567        repeats: None,
568        description: "Logical ID copied from the input relation.",
569    },
570    GpuOutputFieldDefinition {
571        name: "cluster_id",
572        data_type: "Int32",
573        nullable: false,
574        repeats: None,
575        description: "Query-local numeric assignment label, not a stable business or topic identifier.",
576    },
577];
578
579const PCA_OUTPUT_FIELDS: &[GpuOutputFieldDefinition] = &[
580    GpuOutputFieldDefinition {
581        name: "row_ordinal",
582        data_type: "UInt64",
583        nullable: false,
584        repeats: None,
585        description: "Zero-based ordinal of the evaluated input row; it disambiguates duplicate IDs. Follows the input's top-level ORDER BY on output columns; unspecified without one.",
586    },
587    GpuOutputFieldDefinition {
588        name: "id",
589        data_type: "same_as_input.id",
590        nullable: false,
591        repeats: None,
592        description: "Logical ID copied from the input relation.",
593    },
594    GpuOutputFieldDefinition {
595        name: "pc_<index>",
596        data_type: "Float32",
597        nullable: false,
598        repeats: Some("n_components"),
599        description: "One transformed principal-component value; the field count is set by the literal n_components option.",
600    },
601];
602
603/// Operation identity used by native IR and the SQL adapter.
604#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
605pub enum CuvsOperation {
606    /// Exact neighbor search over the evaluated dataset.
607    BruteForceKnn,
608    /// Approximate neighbor search with inverted lists storing full vectors.
609    IvfFlat,
610    /// Approximate neighbor search with product-quantized inverted lists.
611    IvfPq,
612    /// Approximate neighbor search with a query-local graph index.
613    Cagra,
614    /// Evaluates every pair across two dense relations.
615    PairwiseDistance,
616    /// Selects the smallest or largest values within each dense row.
617    TopK,
618    /// Fits clusters and assigns the input rows.
619    KMeans,
620    /// Fits principal components and transforms the input rows.
621    Pca,
622}
623
624impl CuvsOperation {
625    /// Iterates supported operations in the canonical registry order.
626    pub fn all() -> impl ExactSizeIterator<Item = Self> + Clone {
627        CUVS_OPERATION_REGISTRY.iter().map(|descriptor| descriptor.operation)
628    }
629
630    /// Returns this operation's static options, element support, and output contract.
631    pub fn descriptor(self) -> &'static CuvsOperationDescriptor {
632        CUVS_OPERATION_REGISTRY
633            .iter()
634            .find(|descriptor| descriptor.operation == self)
635            .expect("cuVS operation registry is missing an operation variant")
636    }
637
638    /// Looks up an exact canonical SQL name; returns `None` for an unknown name.
639    ///
640    /// This lookup does not normalize case.
641    pub fn from_sql_function_name(name: &str) -> Option<Self> {
642        CUVS_OPERATION_REGISTRY.iter().find(|descriptor| descriptor.sql_function_name == name).map(|d| d.operation)
643    }
644}
645
646/// Supported cuVS operations in their canonical catalog order.
647pub const CUVS_OPERATION_REGISTRY: &[CuvsOperationDescriptor] = &[
648    CuvsOperationDescriptor {
649        operation: CuvsOperation::BruteForceKnn,
650        sql_function_name: CUVS_BRUTE_FORCE_KNN_FUNCTION,
651        display_name: "Brute-force kNN",
652        summary: "Exact brute-force nearest-neighbor search over dense vectors.",
653        use_this_when: "Use this when SQL can first filter a candidate relation and you need exact batch reranking, not a persistent ANN index.",
654        relation_roles: KNN_RELATIONS,
655        vector_element_support: BRUTE_FORCE_SUPPORT,
656        output_contract: CuvsOutputContract::KnnNeighbors,
657        options: KNN_OPTIONS,
658        output_fields: KNN_OUTPUT_FIELDS,
659        quickstart_sql: r#"SELECT *
660FROM cuvs_brute_force_knn(
661  dataset => (SELECT id, d0, d1 FROM dataset_vectors),
662  queries => (SELECT id, d0, d1 FROM query_vectors),
663  k => 8,
664  metric => 'l2_expanded'
665);"#,
666        limitations: &[
667            "Builds an exact query-local index over the evaluated dataset relation; it does not persist an ANN index or replace a vector database.",
668            "The evaluated dataset must be non-empty and k must not exceed its row count. An empty query relation may return an empty result with the stable schema.",
669            "Dataset and query vector dimensions must match after both relation children are evaluated.",
670            "Distances repeat between runs; neighbors that tie at the k-th distance can differ, because candidate merge order is not fixed.",
671        ],
672    },
673    CuvsOperationDescriptor {
674        operation: CuvsOperation::IvfFlat,
675        sql_function_name: CUVS_IVF_FLAT_FUNCTION,
676        display_name: "IVF-Flat",
677        summary: "Query-local inverted-file approximate nearest-neighbor search.",
678        use_this_when: "Use this when a query-local approximate index reduces search work over a large evaluated dataset.",
679        relation_roles: KNN_RELATIONS,
680        vector_element_support: IVF_SUPPORT,
681        output_contract: CuvsOutputContract::KnnNeighbors,
682        options: IVF_FLAT_OPTIONS,
683        output_fields: KNN_OUTPUT_FIELDS,
684        quickstart_sql: r#"SELECT *
685FROM cuvs_ivf_flat(
686  dataset => (SELECT id, d0, d1 FROM dataset_vectors),
687  queries => (SELECT id, d0, d1 FROM query_vectors),
688  k => 8,
689  metric => 'l2_expanded',
690  n_lists => 16,
691  n_probes => 16
692)
693ORDER BY query_ordinal, rank;"#,
694        limitations: &[
695            "Builds and frees its index during each query execution; executions never share an index.",
696            "Results are approximate unless all lists are probed. Fewer than k neighbors may be returned for a query; missing (query, rank) rows are omitted and valid ranks remain contiguous.",
697            "The evaluated dataset must be non-empty, k and n_lists must not exceed its row count, and n_probes must not exceed n_lists.",
698            "Dataset and query dimensions and element types must match. Cosine requires at least two dimensions.",
699            "Neighbors and distances can differ slightly between runs, because index construction accumulates k-means centroid sums with floating-point atomics; there is no seed option.",
700        ],
701    },
702    CuvsOperationDescriptor {
703        operation: CuvsOperation::IvfPq,
704        sql_function_name: CUVS_IVF_PQ_FUNCTION,
705        display_name: "IVF-PQ",
706        summary: "Query-local product-quantized inverted-file nearest-neighbor search.",
707        use_this_when: "Use this when a query-local compressed index can reduce vector search memory, subject to lossy recall.",
708        relation_roles: KNN_RELATIONS,
709        vector_element_support: IVF_PQ_SUPPORT,
710        output_contract: CuvsOutputContract::KnnNeighbors,
711        options: IVF_PQ_OPTIONS,
712        output_fields: KNN_OUTPUT_FIELDS,
713        quickstart_sql: r#"SELECT *
714FROM cuvs_ivf_pq(
715  dataset => (SELECT id, d0, d1 FROM dataset_vectors),
716  queries => (SELECT id, d0, d1 FROM query_vectors),
717  k => 8,
718  metric => 'l2_expanded',
719  n_lists => 16,
720  n_probes => 16,
721  pq_dim => 2
722)
723ORDER BY query_ordinal, rank;"#,
724        limitations: &[
725            "Builds and frees its index during each query execution; executions never share an index.",
726            "Product quantization is lossy even when all lists are probed. Missing neighbors are omitted and valid ranks remain contiguous.",
727            "The dataset must be non-empty; k and n_lists must not exceed its row count, and n_probes must not exceed n_lists.",
728            "Dataset and query dimensions and element types must match. Cosine requires at least two dimensions. Nonzero pq_dim * pq_bits must be divisible by 8; dimensions are padded when needed.",
729            "Neighbors and distances can differ slightly between runs, because index training accumulates k-means centroid sums with floating-point atomics; there is no seed option.",
730        ],
731    },
732    CuvsOperationDescriptor {
733        operation: CuvsOperation::Cagra,
734        sql_function_name: CUVS_CAGRA_FUNCTION,
735        display_name: "CAGRA",
736        summary: "Query-local graph approximate nearest-neighbor search.",
737        use_this_when: "Use this for graph-based approximate search over an evaluated dataset.",
738        relation_roles: KNN_RELATIONS,
739        vector_element_support: CAGRA_SUPPORT,
740        output_contract: CuvsOutputContract::KnnNeighbors,
741        options: CAGRA_OPTIONS,
742        output_fields: KNN_OUTPUT_FIELDS,
743        quickstart_sql: r#"SELECT *
744FROM cuvs_cagra(
745  dataset => (SELECT id, d0, d1 FROM dataset_vectors),
746  queries => (SELECT id, d0, d1 FROM query_vectors),
747  k => 8,
748  metric => 'l2_expanded',
749  graph_degree => 32,
750  intermediate_graph_degree => 64
751)
752ORDER BY query_ordinal, rank;"#,
753        limitations: &[
754            "Builds and frees its graph and padded dataset during each query execution; executions never share an index.",
755            "Approximate search may return fewer than k neighbors; missing (query, rank) rows are omitted and valid ranks remain contiguous.",
756            "The dataset must be non-empty, k and intermediate_graph_degree must be less than or equal to its row count, and graph_degree must be less than intermediate_graph_degree.",
757            "k must not exceed itopk_size. single_cta requires itopk_size no greater than 1024. cosine cannot use nn_descent.",
758            "Dataset and query dimensions and element types must match.",
759            "Neighbors and distances can differ slightly between runs, because graph construction orders edges by atomic arrival; there is no seed option.",
760        ],
761    },
762    CuvsOperationDescriptor {
763        operation: CuvsOperation::PairwiseDistance,
764        sql_function_name: CUVS_PAIRWISE_DISTANCE_FUNCTION,
765        display_name: "Pairwise Distance",
766        summary: "Pairwise distance or similarity for every x and y vector row.",
767        use_this_when: "Use this when every pair between two evaluated relations needs a metric value.",
768        relation_roles: PAIRWISE_RELATIONS,
769        vector_element_support: PAIRWISE_SUPPORT,
770        output_contract: CuvsOutputContract::PairwiseDistances,
771        options: PAIRWISE_OPTIONS,
772        output_fields: PAIRWISE_OUTPUT_FIELDS,
773        quickstart_sql: r#"SELECT *
774FROM cuvs_pairwise_distance(
775  x => (SELECT id, d0, d1 FROM x_vectors),
776  y => (SELECT id, d0, d1 FROM y_vectors),
777  metric => 'l2_expanded'
778)
779ORDER BY x_ordinal, y_ordinal;"#,
780        limitations: &[
781            "Output rows equal x rows multiplied by y rows; the query allocation cap limits the result size.",
782            "Empty x or y returns an empty result with the stable schema.",
783            "x and y must share the same vector dimension and element type. inner_product is a similarity where larger is closer.",
784        ],
785    },
786    CuvsOperationDescriptor {
787        operation: CuvsOperation::TopK,
788        sql_function_name: CUVS_TOP_K_FUNCTION,
789        display_name: "Vector Top-K",
790        summary: "Select the smallest or largest values in every dense vector row.",
791        use_this_when: "Use this when each row needs its own ranked scalar values and vector positions.",
792        relation_roles: INPUT_RELATION,
793        vector_element_support: FLOAT32_VECTOR_SUPPORT,
794        output_contract: CuvsOutputContract::TopKValues,
795        options: TOP_K_OPTIONS,
796        output_fields: TOP_K_OUTPUT_FIELDS,
797        quickstart_sql: r#"SELECT *
798FROM cuvs_top_k(
799  input => (SELECT id, d0, d1 FROM input_vectors),
800  k => 1,
801  select => 'max'
802)
803ORDER BY row_ordinal, rank;"#,
804        limitations: &[
805            "k must not exceed the vector dimension. Empty input returns an empty result with the stable schema.",
806            "Equal selected values are ordered by position; when ties straddle rank k, the selected positions among them are unspecified.",
807        ],
808    },
809    CuvsOperationDescriptor {
810        operation: CuvsOperation::KMeans,
811        sql_function_name: CUVS_KMEANS_FUNCTION,
812        display_name: "KMeans",
813        summary: "K-means clustering assignments for dense vector rows.",
814        use_this_when: "Use this when one query should fit clusters over its evaluated relation and immediately join or aggregate the assignments.",
815        relation_roles: INPUT_RELATION,
816        vector_element_support: KMEANS_SUPPORT,
817        output_contract: CuvsOutputContract::KMeansAssignments,
818        options: KMEANS_OPTIONS,
819        output_fields: KMEANS_OUTPUT_FIELDS,
820        quickstart_sql: r#"SELECT id, cluster_id
821FROM cuvs_kmeans(
822  input => (SELECT id, d0, d1 FROM input_vectors),
823  n_clusters => 8
824)
825ORDER BY row_ordinal;"#,
826        limitations: &[
827            "Fits a model and returns assignments within one call; it does not return a reusable model, centroids, or inertia.",
828            "The evaluated input must be non-empty and n_clusters must not exceed its row count.",
829            "cluster_id values are query-local labels. Do not attach permanent business meaning to their numeric values.",
830            "Cluster sizes can differ slightly between runs, because centroid sums accumulate with floating-point atomics; the native seed is fixed at 0 and cannot be set.",
831        ],
832    },
833    CuvsOperationDescriptor {
834        operation: CuvsOperation::Pca,
835        sql_function_name: CUVS_PCA_FUNCTION,
836        display_name: "PCA",
837        summary: "Principal-component analysis transform for dense vector rows.",
838        use_this_when: "Use this when one query should fit and transform its evaluated input relation for relational export or downstream visualization.",
839        relation_roles: INPUT_RELATION,
840        vector_element_support: FLOAT32_VECTOR_SUPPORT,
841        output_contract: CuvsOutputContract::PcaComponents,
842        options: PCA_OPTIONS,
843        output_fields: PCA_OUTPUT_FIELDS,
844        quickstart_sql: r#"SELECT id, pc_0, pc_1
845FROM cuvs_pca(
846  input => (SELECT id, d0, d1 FROM input_vectors),
847  n_components => 2
848)
849ORDER BY row_ordinal;"#,
850        limitations: &[
851            "Fits and transforms the evaluated input in one call; it does not return a reusable transform, components, or explained-variance metadata.",
852            "The evaluated input needs at least two rows and two dimensions, and n_components must not exceed the evaluated dimension.",
853            "Component sign orientation is not a semantic label. The output schema changes with the literal n_components option.",
854        ],
855    },
856];
857
858impl CuvsOperationDescriptor {
859    /// Returns supported metrics for `element`, or `None` when the element type is unsupported.
860    pub fn element_support(&self, element: VectorElementType) -> Option<&CuvsVectorElementSupport> {
861        self.vector_element_support.iter().find(|support| support.element == element)
862    }
863
864    /// Looks up an exact option name; returns `None` if this operation has no such option.
865    pub fn option(&self, name: &str) -> Option<&'static GpuOptionDefinition> {
866        self.options.iter().find(|option| option.name == name)
867    }
868
869    /// Builds a draft 2020-12 JSON Schema for this operation's options.
870    ///
871    /// The schema includes defaults, supported relation element types, and required
872    /// properties. Unknown properties are rejected; supported metric combinations
873    /// are also included in the metric description.
874    pub fn options_json_schema(&self) -> Value {
875        let mut properties = serde_json::Map::new();
876        let mut required = Vec::new();
877        for option in self.options {
878            let mut schema = gpu_option_schema(
879                option.kind,
880                &dense_relation_binding_schema(
881                    &self.vector_element_support.iter().map(|support| support.element).collect::<Vec<_>>(),
882                ),
883            );
884            schema["description"] = Value::String(option.description.to_string());
885            properties.insert(option.name.to_string(), schema);
886            if option.required {
887                required.push(Value::String(option.name.to_string()));
888            }
889        }
890        if let Some(metric) = properties.get_mut("metric") {
891            let combinations = self
892                .vector_element_support
893                .iter()
894                .map(|support| format!("{}: {}", support.element.as_str(), support.metrics.join(", ")))
895                .collect::<Vec<_>>()
896                .join("; ");
897            metric["description"] = json!(format!(
898                "{} Supported element/metric combinations: {combinations}.",
899                self.option("metric").expect("metric option").description
900            ));
901        }
902        json!({
903            "$schema": "https://json-schema.org/draft/2020-12/schema",
904            "title": format!("{} options", self.sql_function_name),
905            "type": "object",
906            "additionalProperties": false,
907            "required": required,
908            "properties": properties,
909        })
910    }
911}
912
913#[cfg(test)]
914fn option_schema(kind: GpuOptionKind, element_support: &[CuvsVectorElementSupport]) -> Value {
915    gpu_option_schema(
916        kind,
917        &dense_relation_binding_schema(&element_support.iter().map(|support| support.element).collect::<Vec<_>>()),
918    )
919}
920
921/// Validate and canonicalize the scalar portion of an options JSON object.
922/// Relation binding field names/types are checked by the SQL adapter against
923/// its child schema, but presence and object shape are enforced here.
924pub fn canonicalize_cuvs_options(
925    operation: CuvsOperation,
926    raw: &str,
927) -> Result<Map<String, Value>, GpuOptionValidationError> {
928    let descriptor = operation.descriptor();
929    let object = canonicalize_options(descriptor.options, raw, "cuVS", |option, value| {
930        validate_dense_relation_binding_value("cuVS", option, value, false)
931    })?;
932    if matches!(operation, CuvsOperation::IvfFlat | CuvsOperation::IvfPq) {
933        let n_probes = object.get("n_probes").and_then(Value::as_u64).expect("validated n_probes");
934        let n_lists = object.get("n_lists").and_then(Value::as_u64).expect("validated n_lists");
935        if n_probes > n_lists {
936            return Err(
937                option_error("missing_or_invalid_option", "n_probes must not exceed n_lists").for_argument("n_probes")
938            );
939        }
940    }
941    if operation == CuvsOperation::Cagra {
942        let integer = |name| object.get(name).and_then(Value::as_u64).expect("validated CAGRA integer");
943        if integer("intermediate_graph_degree") <= integer("graph_degree") {
944            return Err(option_error(
945                "missing_or_invalid_option",
946                "intermediate_graph_degree must exceed graph_degree",
947            )
948            .for_argument("intermediate_graph_degree"));
949        }
950        if integer("k") > integer("itopk_size") {
951            return Err(
952                option_error("missing_or_invalid_option", "k must not exceed itopk_size").for_argument("itopk_size")
953            );
954        }
955        if object["search_algo"] == "single_cta" && integer("itopk_size") > 1024 {
956            return Err(option_error("missing_or_invalid_option", "single_cta requires itopk_size <= 1024")
957                .for_argument("itopk_size"));
958        }
959        if object["metric"] == "cosine" && object["build_algo"] == "nn_descent" {
960            return Err(option_error("missing_or_invalid_option", "nn_descent cannot build a cosine graph")
961                .for_argument("build_algo"));
962        }
963    }
964    if operation == CuvsOperation::IvfPq {
965        let pq_dim = object.get("pq_dim").and_then(Value::as_u64).expect("validated pq_dim");
966        let pq_bits = object.get("pq_bits").and_then(Value::as_u64).expect("validated pq_bits");
967        if pq_dim != 0 && (pq_dim * pq_bits) % 8 != 0 {
968            return Err(option_error("missing_or_invalid_option", "pq_dim * pq_bits must be divisible by 8")
969                .for_argument("pq_dim"));
970        }
971    }
972    Ok(object)
973}
974
975fn option_error(code: &'static str, message: impl Into<String>) -> GpuOptionValidationError {
976    GpuOptionValidationError::new(code, message)
977}
978
979/// Vector algorithm identity carried by the native IR.
980///
981/// Ref: <https://docs.nvidia.com/cuvs/api-reference/c-api-documentation>
982#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
983pub enum VectorAlgorithm {
984    /// Exact brute-force k-nearest-neighbor search.
985    /// <https://docs.nvidia.com/cuvs/api-reference/c-api-neighbors-brute-force>
986    BruteForceKnn,
987    /// Query-local IVF-Flat approximate nearest-neighbor search.
988    /// <https://docs.nvidia.com/cuvs/api-reference/c-api-neighbors-ivf-flat>
989    IvfFlat,
990    /// Query-local IVF-PQ compressed nearest-neighbor search.
991    /// <https://docs.nvidia.com/cuvs/api-reference/c-api-neighbors-ivf-pq>
992    IvfPq,
993    /// Query-local graph approximate nearest-neighbor search.
994    /// <https://docs.nvidia.com/cuvs/api-reference/c-api-neighbors-cagra>
995    Cagra,
996    /// Pairwise metric values between two relations.
997    /// <https://docs.nvidia.com/cuvs/api-reference/c-api-distance-pairwise-distance>
998    PairwiseDistance,
999    /// Ranked values selected within each input vector.
1000    /// <https://docs.nvidia.com/cuvs/api-reference/cpp-api-selection-select-k>
1001    TopK,
1002    /// K-means fit followed by row-label prediction.
1003    /// <https://docs.nvidia.com/cuvs/api-reference/c-api-cluster-kmeans>
1004    KMeans,
1005    /// Principal-component fit and transform.
1006    /// <https://docs.nvidia.com/cuvs/api-reference/c-api-preprocessing-pca>
1007    Pca,
1008}
1009
1010impl VectorAlgorithm {
1011    /// Stable diagnostic label.
1012    pub const fn as_str(self) -> &'static str {
1013        match self {
1014            Self::BruteForceKnn => "brute_force_knn",
1015            Self::IvfFlat => "ivf_flat",
1016            Self::IvfPq => "ivf_pq",
1017            Self::Cagra => "cagra",
1018            Self::PairwiseDistance => "pairwise_distance",
1019            Self::TopK => "top_k",
1020            Self::KMeans => "kmeans",
1021            Self::Pca => "pca",
1022        }
1023    }
1024}
1025
1026impl fmt::Display for VectorAlgorithm {
1027    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1028        f.write_str(self.as_str())
1029    }
1030}
1031
1032#[cfg(test)]
1033mod tests {
1034    use serde_json::json;
1035
1036    use super::{
1037        CuvsOperation, GpuOptionDefinition, GpuOptionKind, canonicalize_cuvs_options, option_schema,
1038        validate_dense_relation_binding_value,
1039    };
1040    use crate::plan::{
1041        GpuOptionValidationError, gpu_function_options::validate_option_value as gpu_validate_option_value,
1042    };
1043
1044    fn validate_option_value(
1045        option: &GpuOptionDefinition,
1046        value: &serde_json::Value,
1047    ) -> Result<(), GpuOptionValidationError> {
1048        gpu_validate_option_value(option, value, "cuVS", &|option, value| {
1049            validate_dense_relation_binding_value("cuVS", option, value, false)
1050        })
1051    }
1052
1053    #[test]
1054    fn ivf_flat_rejects_probes_above_lists_with_argument_attribution() {
1055        let raw = json!({
1056            "dataset": {"id": "id", "vector": {"columns": ["d0", "d1"]}},
1057            "queries": {"id": "id", "vector": {"columns": ["d0", "d1"]}},
1058            "k": 1, "metric": "l2_expanded", "n_lists": 2, "n_probes": 3
1059        });
1060        let error = canonicalize_cuvs_options(CuvsOperation::IvfFlat, &raw.to_string()).unwrap_err();
1061        assert_eq!(error.argument(), Some("n_probes"));
1062        assert!(error.message().contains("n_probes must not exceed n_lists"));
1063        for fraction in [0.0, 1.01] {
1064            let mut invalid = raw.clone();
1065            invalid["n_probes"] = json!(2);
1066            invalid["kmeans_trainset_fraction"] = json!(fraction);
1067            let error = canonicalize_cuvs_options(CuvsOperation::IvfFlat, &invalid.to_string()).unwrap_err();
1068            assert_eq!(error.argument(), Some("kmeans_trainset_fraction"));
1069        }
1070    }
1071
1072    #[test]
1073    fn ivf_pq_rejects_invalid_bit_packing_and_probe_count_with_argument_attribution() {
1074        let mut raw = json!({
1075            "dataset": {"id": "id", "vector": {"columns": ["d0", "d1"]}},
1076            "queries": {"id": "id", "vector": {"columns": ["d0", "d1"]}},
1077            "k": 1, "metric": "l2_expanded", "n_lists": 2,
1078            "n_probes": 2, "pq_dim": 1, "pq_bits": 5
1079        });
1080        let error = canonicalize_cuvs_options(CuvsOperation::IvfPq, &raw.to_string()).unwrap_err();
1081        assert_eq!(error.argument(), Some("pq_dim"));
1082        assert!(error.message().contains("divisible by 8"));
1083        raw["pq_dim"] = json!(8);
1084        raw["n_probes"] = json!(3);
1085        assert_eq!(
1086            canonicalize_cuvs_options(CuvsOperation::IvfPq, &raw.to_string()).unwrap_err().argument(),
1087            Some("n_probes")
1088        );
1089        raw["n_probes"] = json!(2);
1090        raw["pq_bits"] = json!(9);
1091        assert_eq!(
1092            canonicalize_cuvs_options(CuvsOperation::IvfPq, &raw.to_string()).unwrap_err().argument(),
1093            Some("pq_bits")
1094        );
1095    }
1096
1097    #[test]
1098    fn cagra_rejects_incompatible_graph_and_search_parameters() {
1099        let mut raw = json!({
1100            "dataset": {"id": "id", "vector": {"columns": ["d0", "d1"]}},
1101            "queries": {"id": "id", "vector": {"columns": ["d0", "d1"]}},
1102            "k": 8, "metric": "l2_expanded", "graph_degree": 32,
1103            "intermediate_graph_degree": 16
1104        });
1105        assert_eq!(
1106            canonicalize_cuvs_options(CuvsOperation::Cagra, &raw.to_string()).unwrap_err().argument(),
1107            Some("intermediate_graph_degree")
1108        );
1109        raw["intermediate_graph_degree"] = json!(64);
1110        raw["itopk_size"] = json!(7);
1111        assert_eq!(
1112            canonicalize_cuvs_options(CuvsOperation::Cagra, &raw.to_string()).unwrap_err().argument(),
1113            Some("itopk_size")
1114        );
1115        raw["itopk_size"] = json!(8);
1116        raw["metric"] = json!("cosine");
1117        raw["build_algo"] = json!("nn_descent");
1118        assert_eq!(
1119            canonicalize_cuvs_options(CuvsOperation::Cagra, &raw.to_string()).unwrap_err().argument(),
1120            Some("build_algo")
1121        );
1122        raw["metric"] = json!("inner_product");
1123        let normalized = canonicalize_cuvs_options(CuvsOperation::Cagra, &raw.to_string()).unwrap();
1124        assert_eq!(normalized["search_algo"], "auto");
1125    }
1126
1127    #[test]
1128    fn integer_options_honor_zero_and_bounded_minima() {
1129        for (minimum, maximum, accepted, rejected) in [(0, 10, [0, 10], [11, 12]), (4, 8, [4, 8], [0, 9])] {
1130            let option = GpuOptionDefinition {
1131                name: "integer",
1132                required: false,
1133                kind: GpuOptionKind::Integer { minimum, maximum, default: Some(minimum) },
1134                description: "Integer range.",
1135            };
1136            assert_eq!(option_schema(option.kind, &[])["minimum"], json!(minimum));
1137            for value in accepted {
1138                validate_option_value(&option, &json!(value)).expect("value inside range");
1139            }
1140            for value in rejected {
1141                assert!(validate_option_value(&option, &json!(value)).is_err());
1142            }
1143        }
1144    }
1145
1146    #[test]
1147    fn fraction_options_exclude_zero_and_values_above_one() {
1148        let option = GpuOptionDefinition {
1149            name: "fraction",
1150            required: false,
1151            kind: GpuOptionKind::Fraction { default: 0.5 },
1152            description: "Training fraction.",
1153        };
1154        assert_eq!(option_schema(option.kind, &[])["exclusiveMinimum"], json!(0));
1155        for value in [0.5, 1.0] {
1156            validate_option_value(&option, &json!(value)).expect("fraction inside range");
1157        }
1158        for value in [0.0, -0.1, 1.1] {
1159            assert!(validate_option_value(&option, &json!(value)).is_err());
1160        }
1161    }
1162}
query_planning/plan/vector_algorithm/definition.rs

query_planning/plan/vector_algorithm/
definition.rs
