query_planning/plan/
vector_algorithm.rs

1//! Native vector algorithm IR.
2//!
3//! The contracts in this module describe relational inputs, algorithm parameters, and result schemas without
4//! depending on cuVS. The optional cuVS executor is selected later by the runtime.
5
6use std::sync::Arc;
7
8use arrow::datatypes::{DataType, Field, Schema};
9
10use super::{NodeId, SchemaRef};
11
12mod brute_force;
13mod cagra;
14mod common;
15mod definition;
16mod ivf_flat;
17mod ivf_pq;
18mod kmeans;
19mod pairwise;
20mod pca;
21mod top_k;
22
23pub use brute_force::VectorBruteForceKnnParams;
24pub use cagra::{VectorCagraBuildAlgo, VectorCagraParams, VectorCagraSearchAlgo};
25pub use common::{VectorDistanceMetric, VectorRelationBinding, VectorRowIdBinding};
26pub use definition::{
27    CUVS_BRUTE_FORCE_KNN_FUNCTION, CUVS_CAGRA_FUNCTION, CUVS_IVF_FLAT_FUNCTION, CUVS_IVF_PQ_FUNCTION,
28    CUVS_KMEANS_FUNCTION, CUVS_OPERATION_REGISTRY, CUVS_PAIRWISE_DISTANCE_FUNCTION, CUVS_PCA_FUNCTION,
29    CUVS_TOP_K_FUNCTION, CuvsOperation, CuvsOperationDescriptor, CuvsOutputContract, CuvsRelationRole,
30    CuvsVectorElementSupport, VectorAlgorithm, canonicalize_cuvs_options,
31};
32pub use ivf_flat::VectorIvfFlatParams;
33pub use ivf_pq::{VectorIvfPqCodebookKind, VectorIvfPqParams};
34pub use kmeans::{VectorKMeansInit, VectorKMeansMetric, VectorKMeansParams};
35pub use pairwise::VectorPairwiseDistanceParams;
36pub use pca::{VectorPcaParams, VectorPcaSolver};
37pub use top_k::VectorTopKParams;
38
39/// Algorithm-specific vector parameter payload.
40#[derive(Debug, Clone, PartialEq)]
41pub enum VectorAlgorithmParams {
42    /// Exact brute-force kNN parameters.
43    BruteForceKnn(VectorBruteForceKnnParams),
44    /// Query-local IVF-Flat search parameters.
45    IvfFlat(VectorIvfFlatParams),
46    /// Query-local IVF-PQ search parameters.
47    IvfPq(VectorIvfPqParams),
48    /// Query-local graph search parameters.
49    Cagra(VectorCagraParams),
50    /// Pairwise metric values between two relations.
51    PairwiseDistance(VectorPairwiseDistanceParams),
52    /// Ranked values selected within each input vector.
53    TopK(VectorTopKParams),
54    /// K-means fit/predict parameters.
55    KMeans(VectorKMeansParams),
56    /// PCA fit-transform parameters.
57    Pca(VectorPcaParams),
58}
59
60impl VectorAlgorithmParams {
61    /// Whether this payload matches the declared algorithm.
62    pub const fn supports_algorithm(&self, algorithm: VectorAlgorithm) -> bool {
63        matches!(
64            (self, algorithm),
65            (Self::BruteForceKnn(_), VectorAlgorithm::BruteForceKnn)
66                | (Self::IvfFlat(_), VectorAlgorithm::IvfFlat)
67                | (Self::IvfPq(_), VectorAlgorithm::IvfPq)
68                | (Self::Cagra(_), VectorAlgorithm::Cagra)
69                | (Self::PairwiseDistance(_), VectorAlgorithm::PairwiseDistance)
70                | (Self::TopK(_), VectorAlgorithm::TopK)
71                | (Self::KMeans(_), VectorAlgorithm::KMeans)
72                | (Self::Pca(_), VectorAlgorithm::Pca)
73        )
74    }
75
76    /// Visit relation bindings in stable semantic execution order.
77    pub fn relation_bindings(&self) -> Box<dyn Iterator<Item = &VectorRelationBinding> + '_> {
78        match self {
79            Self::BruteForceKnn(params) => Box::new([&params.dataset, &params.queries].into_iter()),
80            Self::IvfFlat(params) => Box::new([&params.dataset, &params.queries].into_iter()),
81            Self::IvfPq(params) => Box::new([&params.dataset, &params.queries].into_iter()),
82            Self::Cagra(params) => Box::new([&params.dataset, &params.queries].into_iter()),
83            Self::PairwiseDistance(params) => Box::new([&params.x, &params.y].into_iter()),
84            Self::TopK(params) => Box::new(std::iter::once(&params.input)),
85            Self::KMeans(params) => Box::new(std::iter::once(&params.input).chain(params.initial_centroids.iter())),
86            Self::Pca(params) => Box::new(std::iter::once(&params.input)),
87        }
88    }
89}
90
91/// Native vector algorithm specification.
92#[derive(Debug, Clone, PartialEq)]
93pub struct VectorAlgorithmSpec {
94    /// Algorithm to execute.
95    pub algorithm: VectorAlgorithm,
96    /// Algorithm-specific relation and scalar parameters.
97    pub params: VectorAlgorithmParams,
98    /// Explicit relational output schema.
99    pub output_schema: SchemaRef,
100}
101
102impl VectorAlgorithmSpec {
103    /// Construct a vector algorithm specification.
104    pub fn new(algorithm: VectorAlgorithm, params: VectorAlgorithmParams, output_schema: SchemaRef) -> Self {
105        assert!(params.supports_algorithm(algorithm), "vector algorithm and parameter variant must match");
106        Self { algorithm, params, output_schema }
107    }
108}
109
110/// Native vector algorithm operation over one or more relation inputs.
111#[derive(Debug, Clone, PartialEq)]
112pub struct VectorAlgorithmNode {
113    /// Relation inputs; each relation binding names its relation by its
114    /// position here.
115    pub inputs: Vec<NodeId>,
116    /// Typed algorithm specification.
117    pub spec: VectorAlgorithmSpec,
118}
119
120impl VectorAlgorithmNode {
121    /// Output schema declared by the vector algorithm specification.
122    pub fn output_schema(&self) -> &SchemaRef {
123        &self.spec.output_schema
124    }
125}
126
127/// Stable relational result schema for nearest-neighbor search.
128pub fn vector_knn_output_schema(query_id_type: DataType, neighbor_id_type: DataType) -> SchemaRef {
129    Arc::new(Schema::new(vec![
130        Field::new("query_ordinal", DataType::UInt64, false),
131        Field::new("query_id", query_id_type, false),
132        Field::new("neighbor_ordinal", DataType::UInt64, false),
133        Field::new("neighbor_id", neighbor_id_type, false),
134        Field::new("rank", DataType::UInt32, false),
135        Field::new("distance", DataType::Float32, false),
136    ]))
137}
138
139/// Stable long-form row-major pairwise metric result.
140pub fn vector_pairwise_distance_output_schema(
141    x_id_type: DataType,
142    y_id_type: DataType,
143    distance_type: DataType,
144) -> SchemaRef {
145    Arc::new(Schema::new(vec![
146        Field::new("x_ordinal", DataType::UInt64, false),
147        Field::new("x_id", x_id_type, false),
148        Field::new("y_ordinal", DataType::UInt64, false),
149        Field::new("y_id", y_id_type, false),
150        Field::new("distance", distance_type, false),
151    ]))
152}
153
154/// Stable ranked selection result schema.
155pub fn vector_top_k_output_schema(id_type: DataType) -> SchemaRef {
156    Arc::new(Schema::new(vec![
157        Field::new("row_ordinal", DataType::UInt64, false),
158        Field::new("id", id_type, false),
159        Field::new("rank", DataType::UInt32, false),
160        Field::new("position", DataType::Int64, false),
161        Field::new("value", DataType::Float32, false),
162    ]))
163}
164
165/// Stable relational result schema for K-means row assignments.
166pub fn vector_kmeans_output_schema(row_id_type: DataType) -> SchemaRef {
167    Arc::new(Schema::new(vec![
168        Field::new("row_ordinal", DataType::UInt64, false),
169        Field::new("id", row_id_type, false),
170        Field::new("cluster_id", DataType::Int32, false),
171    ]))
172}
173
174/// Stable relational result schema for PCA transformed rows.
175pub fn vector_pca_output_schema(row_id_type: DataType, n_components: usize) -> SchemaRef {
176    let mut fields = Vec::with_capacity(n_components.saturating_add(2));
177    fields.push(Field::new("row_ordinal", DataType::UInt64, false));
178    fields.push(Field::new("id", row_id_type, false));
179    fields.extend((0..n_components).map(|index| Field::new(format!("pc_{index}"), DataType::Float32, false)));
180    Arc::new(Schema::new(fields))
181}
query_planning/plan/vector_algorithm.rs

query_planning/plan/
vector_algorithm.rs
