algeon_datafusion/cuvs_sql/
options.rs

1use std::{
2    cmp::Ordering,
3    hash::{Hash, Hasher},
4};
5
6use query_planning::plan::{
7    VectorAlgorithmParams, VectorBruteForceKnnParams, VectorCagraBuildAlgo, VectorCagraParams, VectorCagraSearchAlgo,
8    VectorDistanceMetric, VectorIvfFlatParams, VectorIvfPqCodebookKind, VectorIvfPqParams, VectorKMeansInit,
9    VectorKMeansMetric, VectorKMeansParams, VectorPairwiseDistanceParams, VectorPcaParams, VectorPcaSolver,
10    VectorRelationBinding, VectorRowIdBinding, VectorTopKParams,
11};
12use serde_json::{Map, Value};
13
14use super::{
15    CuvsOperation, CuvsRelationRole, DenseVectorBinding, canonicalize_cuvs_options,
16    error::{ERRORS, invalid_option},
17};
18use crate::{
19    error::Result,
20    gpu_functions::dense_relation::{CanonicalOptions, DenseRelationArgument, DenseRelationFunction, LabelColumn},
21};
22
23#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
24/// Validated columns bound to one named relation role in a cuVS SQL call.
25pub struct CuvsVectorRelationBinding {
26    relation_role: CuvsRelationRole,
27    id: String,
28    vector: DenseVectorBinding,
29}
30
31impl CuvsVectorRelationBinding {
32    /// Returns the semantic input role associated with this relation.
33    pub fn relation_role(&self) -> CuvsRelationRole {
34        self.relation_role
35    }
36
37    /// Returns the bound row-id column used to correlate results with input rows.
38    pub fn id(&self) -> &str {
39        &self.id
40    }
41
42    /// Returns the validated dense-vector representation and column binding.
43    pub fn vector(&self) -> &DenseVectorBinding {
44        &self.vector
45    }
46}
47
48/// Options of one cuVS call, validated and parsed once. The canonical JSON
49/// is the call's identity and serialized form; validation and lowering read
50/// the typed parameters.
51#[derive(Debug, Clone)]
52pub struct CuvsValidatedOptions {
53    canonical_json: String,
54    value_options: Map<String, Value>,
55    bindings: Vec<CuvsVectorRelationBinding>,
56    params: VectorAlgorithmParams,
57}
58
59impl CuvsValidatedOptions {
60    /// Returns normalized option JSON used as this call's identity and serialized representation.
61    pub fn canonical_json(&self) -> &str {
62        &self.canonical_json
63    }
64
65    /// Returns validated relation bindings in the operation descriptor's role order.
66    pub fn bindings(&self) -> &[CuvsVectorRelationBinding] {
67        &self.bindings
68    }
69
70    /// The native vector algorithm parameters. Their relation `n` reads the
71    /// `n`th binding.
72    pub fn params(&self) -> &VectorAlgorithmParams {
73        &self.params
74    }
75
76    /// The value options by name, without relation bindings.
77    pub(super) fn value_options(&self) -> &Map<String, Value> {
78        &self.value_options
79    }
80}
81
82// The canonical JSON determines every other field, so it is the identity.
83impl PartialEq for CuvsValidatedOptions {
84    fn eq(&self, other: &Self) -> bool {
85        self.canonical_json == other.canonical_json
86    }
87}
88
89impl Eq for CuvsValidatedOptions {}
90
91impl Hash for CuvsValidatedOptions {
92    fn hash<H: Hasher>(&self, state: &mut H) {
93        self.canonical_json.hash(state);
94    }
95}
96
97impl PartialOrd for CuvsValidatedOptions {
98    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
99        Some(self.cmp(other))
100    }
101}
102
103impl Ord for CuvsValidatedOptions {
104    fn cmp(&self, other: &Self) -> Ordering {
105        self.canonical_json.cmp(&other.canonical_json)
106    }
107}
108
109pub(super) fn validate_options(operation: CuvsOperation, raw: &str) -> Result<CuvsValidatedOptions> {
110    let canonical = canonicalize_cuvs_options(operation, raw)
111        .map_err(|error| invalid_option(operation, error.code(), error.argument(), error.message()))?;
112    let options = dense_relation_function(operation).canonical_options(&canonical);
113    let bindings = operation
114        .descriptor()
115        .relation_roles
116        .iter()
117        .map(|&relation_role| {
118            let binding = options.relation_binding(relation_role.as_str())?;
119            Ok(CuvsVectorRelationBinding { relation_role, id: binding.id, vector: binding.vector })
120        })
121        .collect::<Result<Vec<_>>>()?;
122    let params = algorithm_params(operation, &options, &bindings)?;
123    let value_options = options.value_options(operation.descriptor().options);
124    Ok(CuvsValidatedOptions { canonical_json: Value::Object(canonical).to_string(), value_options, bindings, params })
125}
126
127const DISTANCE_METRICS: [(&str, VectorDistanceMetric); 4] = [
128    ("l2_expanded", VectorDistanceMetric::L2Expanded),
129    ("l2_sqrt_expanded", VectorDistanceMetric::L2SqrtExpanded),
130    ("cosine", VectorDistanceMetric::CosineExpanded),
131    ("inner_product", VectorDistanceMetric::InnerProduct),
132];
133
134fn algorithm_params(
135    operation: CuvsOperation,
136    options: &CanonicalOptions<'_>,
137    bindings: &[CuvsVectorRelationBinding],
138) -> Result<VectorAlgorithmParams> {
139    let relation = |ordinal: usize| {
140        let binding = &bindings[ordinal];
141        let row_id = VectorRowIdBinding::column(binding.id());
142        match binding.vector() {
143            DenseVectorBinding::Columns(columns) => {
144                VectorRelationBinding::new(ordinal, row_id, columns.iter().cloned())
145            },
146            DenseVectorBinding::VectorColumn(column) => {
147                VectorRelationBinding::new_vector_column(ordinal, row_id, column.clone())
148            },
149        }
150    };
151    Ok(match operation {
152        CuvsOperation::BruteForceKnn => VectorAlgorithmParams::BruteForceKnn(VectorBruteForceKnnParams::new(
153            relation(0),
154            relation(1),
155            options.usize("k")?,
156            options.choice("metric", &DISTANCE_METRICS)?,
157        )),
158        CuvsOperation::IvfFlat => {
159            let mut params = VectorIvfFlatParams::new(
160                relation(0),
161                relation(1),
162                options.usize("k")?,
163                options.choice("metric", &DISTANCE_METRICS)?,
164                options.u32("n_lists")?,
165            );
166            params.n_probes = options.u32("n_probes")?;
167            params.kmeans_n_iters = options.u32("kmeans_n_iters")?;
168            params.kmeans_trainset_fraction = options.f64("kmeans_trainset_fraction")?;
169            VectorAlgorithmParams::IvfFlat(params)
170        },
171        CuvsOperation::IvfPq => {
172            let mut params = VectorIvfPqParams::new(
173                relation(0),
174                relation(1),
175                options.usize("k")?,
176                options.choice("metric", &DISTANCE_METRICS)?,
177                options.u32("n_lists")?,
178            );
179            params.n_probes = options.u32("n_probes")?;
180            params.kmeans_n_iters = options.u32("kmeans_n_iters")?;
181            params.kmeans_trainset_fraction = options.f64("kmeans_trainset_fraction")?;
182            params.pq_dim = options.u32("pq_dim")?;
183            params.pq_bits = options.u32("pq_bits")?;
184            params.codebook_kind = options.choice(
185                "codebook_kind",
186                &[
187                    ("per_subspace", VectorIvfPqCodebookKind::PerSubspace),
188                    ("per_cluster", VectorIvfPqCodebookKind::PerCluster),
189                ],
190            )?;
191            VectorAlgorithmParams::IvfPq(params)
192        },
193        CuvsOperation::Cagra => {
194            let mut params = VectorCagraParams::new(
195                relation(0),
196                relation(1),
197                options.usize("k")?,
198                options.choice("metric", &DISTANCE_METRICS)?,
199            );
200            params.graph_degree = options.u32("graph_degree")?;
201            params.intermediate_graph_degree = options.u32("intermediate_graph_degree")?;
202            params.build_algo = options.choice(
203                "build_algo",
204                &[("ivf_pq", VectorCagraBuildAlgo::IvfPq), ("nn_descent", VectorCagraBuildAlgo::NnDescent)],
205            )?;
206            params.itopk_size = options.u32("itopk_size")?;
207            params.search_width = options.u32("search_width")?;
208            params.max_iterations = options.u32("max_iterations")?;
209            params.search_algo = options.choice(
210                "search_algo",
211                &[
212                    ("auto", VectorCagraSearchAlgo::Auto),
213                    ("single_cta", VectorCagraSearchAlgo::SingleCta),
214                    ("multi_cta", VectorCagraSearchAlgo::MultiCta),
215                    ("multi_kernel", VectorCagraSearchAlgo::MultiKernel),
216                ],
217            )?;
218            VectorAlgorithmParams::Cagra(params)
219        },
220        CuvsOperation::PairwiseDistance => VectorAlgorithmParams::PairwiseDistance(VectorPairwiseDistanceParams {
221            x: relation(0),
222            y: relation(1),
223            metric: options.choice("metric", &DISTANCE_METRICS)?,
224        }),
225        CuvsOperation::TopK => VectorAlgorithmParams::TopK(VectorTopKParams {
226            input: relation(0),
227            k: options.usize("k")?,
228            select_min: options.choice("select", &[("min", true), ("max", false)])?,
229        }),
230        CuvsOperation::KMeans => {
231            let mut params = VectorKMeansParams::new(relation(0), options.usize("n_clusters")?);
232            params.metric = options.choice(
233                "metric",
234                &[
235                    ("l2_expanded", VectorKMeansMetric::L2Expanded),
236                    ("l2_sqrt_expanded", VectorKMeansMetric::L2SqrtExpanded),
237                ],
238            )?;
239            params.max_iterations = options.u32("max_iter")?;
240            params.tolerance = options.f64("tol")?;
241            params.n_init = options.u32("n_init")?;
242            params.init = options.choice(
243                "init",
244                &[("kmeans++", VectorKMeansInit::KMeansPlusPlus), ("random", VectorKMeansInit::Random)],
245            )?;
246            VectorAlgorithmParams::KMeans(params)
247        },
248        CuvsOperation::Pca => {
249            let mut params = VectorPcaParams::new(relation(0), options.usize("n_components")?);
250            params.whiten = options.bool("whiten")?;
251            params.solver = options.choice(
252                "solver",
253                &[("cov_eig_dq", VectorPcaSolver::CovEigDq), ("cov_eig_jacobi", VectorPcaSolver::CovEigJacobi)],
254            )?;
255            params.tolerance = options.f64("tol")?;
256            params.n_iterations = options.u32("n_iterations")?;
257            params.flip_signs_based_on_u = options.bool("flip_signs_based_on_u")?;
258            VectorAlgorithmParams::Pca(params)
259        },
260    })
261}
262
263/// The cuVS operation as a dense-relation function. No cuVS relation trains
264/// on a label, so a `label` column is a feature.
265pub(super) fn dense_relation_function(operation: CuvsOperation) -> DenseRelationFunction {
266    let descriptor = operation.descriptor();
267    DenseRelationFunction {
268        errors: &ERRORS,
269        sql_function_name: descriptor.sql_function_name,
270        relations: descriptor
271            .relation_roles
272            .iter()
273            .map(|role| DenseRelationArgument { role: role.as_str(), label: LabelColumn::Feature })
274            .collect(),
275        options: descriptor.options,
276    }
277}
algeon_datafusion/cuvs_sql/options.rs

algeon_datafusion/cuvs_sql/
options.rs
