algeon_datafusion/cuvs_sql/
logical.rs

1use std::{cmp::Ordering, fmt, hash::Hash, sync::Arc};
2
3use arrow::datatypes::{DataType, Field, FieldRef, SchemaRef};
4use datafusion::{
5    common::{DFSchema, DFSchemaRef, DataFusionError, Result as DataFusionResult},
6    logical_expr::{Expr, LogicalPlan, SortExpr, UserDefinedLogicalNodeCore},
7};
8use query_planning::plan::{
9    RelationIdType, VectorAlgorithmParams, VectorDistanceMetric, VectorElementType, VectorIvfFlatParams,
10    VectorIvfPqParams, vector_kmeans_output_schema, vector_knn_output_schema, vector_pairwise_distance_output_schema,
11    vector_pca_output_schema, vector_top_k_output_schema,
12};
13use serde_json::Value;
14
15use super::{
16    CuvsOperation, CuvsValidatedOptions, CuvsVectorRelationBinding, DenseVectorBinding,
17    error::{field_not_found, invalid_call, unsupported_type},
18    options::{dense_relation_function, validate_options},
19};
20use crate::{
21    error::{Error, Result},
22    gpu_functions::{
23        dense_relation::known_row_count, fmt_input_orderings, input_ordering_expressions,
24        input_orderings_with_expressions,
25    },
26};
27
28#[derive(Debug, Clone)]
29/// DataFusion logical extension for a statically validated cuVS relation call.
30///
31/// Input schemas and options determine its output schema before GPU execution.
32pub struct CuvsLogicalNode {
33    operation: CuvsOperation,
34    inputs: Vec<LogicalPlan>,
35    /// Top-level `ORDER BY` of each relation argument, in relation-role order;
36    /// empty when the input declares no order.
37    input_orderings: Vec<Vec<SortExpr>>,
38    options: CuvsValidatedOptions,
39    schema: DFSchemaRef,
40}
41
42impl CuvsLogicalNode {
43    /// Validates relation count, option JSON, input schemas, and known cardinalities.
44    ///
45    /// Inputs follow the operation descriptor's relation-role order. Invalid options,
46    /// missing columns, unsupported types, or incompatible cardinalities return typed
47    /// Algeon errors through the DataFusion error boundary.
48    pub fn try_new(operation: CuvsOperation, inputs: Vec<LogicalPlan>, options_json: &str) -> DataFusionResult<Self> {
49        let descriptor = operation.descriptor();
50        if inputs.len() != descriptor.relation_roles.len() {
51            return Err(invalid_call(
52                operation,
53                "invalid_relation_count",
54                format!(
55                    "{} requires {} relation argument(s), received {}",
56                    descriptor.sql_function_name,
57                    descriptor.relation_roles.len(),
58                    inputs.len()
59                ),
60            )
61            .into());
62        }
63        let options = validate_options(operation, options_json)?;
64        let (dimensions, element) = validate_input_schemas(operation, &inputs, &options)?;
65        validate_known_cardinalities(operation, &inputs, &options)?;
66        let schema = derive_output_schema(operation, &inputs, &options, &dimensions, element)?;
67        let input_orderings = vec![Vec::new(); inputs.len()];
68        Ok(Self { operation, inputs, input_orderings, options, schema })
69    }
70
71    /// Declares each relation argument's order, which fixes its `row_ordinal`
72    /// numbering; `input_orderings` has one entry per relation role.
73    pub(crate) fn with_input_orderings(mut self, input_orderings: Vec<Vec<SortExpr>>) -> Result<Self> {
74        if input_orderings.len() != self.inputs.len() {
75            return Err(invalid_call(
76                self.operation,
77                "input_ordering_count_mismatch",
78                "cuVS input orderings do not match the relation arguments",
79            ));
80        }
81        self.input_orderings = input_orderings;
82        Ok(self)
83    }
84
85    pub(crate) fn input_orderings(&self) -> &[Vec<SortExpr>] {
86        &self.input_orderings
87    }
88
89    /// Returns the registered operation represented by this logical node.
90    pub fn operation(&self) -> CuvsOperation {
91        self.operation
92    }
93
94    /// Borrows options parsed and validated during node construction.
95    pub fn validated_options(&self) -> &CuvsValidatedOptions {
96        &self.options
97    }
98
99    /// Returns the Arrow result schema established by call validation.
100    pub fn output_schema(&self) -> SchemaRef {
101        Arc::clone(self.schema.inner())
102    }
103}
104
105impl PartialEq for CuvsLogicalNode {
106    fn eq(&self, other: &Self) -> bool {
107        self.operation == other.operation
108            && self.inputs == other.inputs
109            && self.input_orderings == other.input_orderings
110            && self.options == other.options
111    }
112}
113
114impl Eq for CuvsLogicalNode {}
115
116impl Hash for CuvsLogicalNode {
117    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
118        self.operation.hash(state);
119        self.inputs.hash(state);
120        self.input_orderings.hash(state);
121        self.options.hash(state);
122    }
123}
124
125impl PartialOrd for CuvsLogicalNode {
126    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
127        (&self.operation, &self.inputs, &self.input_orderings, &self.options).partial_cmp(&(
128            &other.operation,
129            &other.inputs,
130            &other.input_orderings,
131            &other.options,
132        ))
133    }
134}
135
136impl UserDefinedLogicalNodeCore for CuvsLogicalNode {
137    fn name(&self) -> &str {
138        "CuvsLogicalNode"
139    }
140
141    fn inputs(&self) -> Vec<&LogicalPlan> {
142        self.inputs.iter().collect()
143    }
144
145    fn schema(&self) -> &DFSchemaRef {
146        &self.schema
147    }
148
149    /// The input sort keys, flattened in relation-role order.
150    fn expressions(&self) -> Vec<Expr> {
151        input_ordering_expressions(&self.input_orderings)
152    }
153
154    fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
155        write!(
156            f,
157            "CuvsLogicalNode: operation={}, relation_count={}, options={}",
158            self.operation.descriptor().sql_function_name,
159            self.inputs.len(),
160            self.options.canonical_json()
161        )?;
162        fmt_input_orderings(
163            f,
164            self.operation.descriptor().relation_roles.iter().map(|role| role.as_str()),
165            &self.input_orderings,
166        )
167    }
168
169    fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> DataFusionResult<Self> {
170        let Some(input_orderings) = input_orderings_with_expressions(&self.input_orderings, exprs) else {
171            return Err(invalid_call(
172                self.operation,
173                "unexpected_logical_expressions",
174                "cuVS logical node expressions must be its input sort keys",
175            )
176            .into());
177        };
178        let options_json =
179            dense_relation_function(self.operation).options_json(&inputs, self.options.value_options().clone())?;
180        Self::try_new(self.operation, inputs, &options_json)?
181            .with_input_orderings(input_orderings)
182            .map_err(DataFusionError::from)
183    }
184
185    fn supports_limit_pushdown(&self) -> bool {
186        false
187    }
188}
189
190fn validate_input_schemas(
191    operation: CuvsOperation,
192    inputs: &[LogicalPlan],
193    options: &CuvsValidatedOptions,
194) -> Result<(Vec<Option<usize>>, VectorElementType)> {
195    let mut dimensions = Vec::with_capacity(inputs.len());
196    let mut elements = Vec::with_capacity(inputs.len());
197    for ((input, role), binding) in
198        inputs.iter().zip(operation.descriptor().relation_roles.iter().copied()).zip(options.bindings())
199    {
200        if binding.relation_role() != role {
201            return Err(invalid_call(
202                operation,
203                "binding_relation_mismatch",
204                "cuVS options binding order does not match the relation arguments",
205            ));
206        }
207        let (dimension, element) = validate_relation_binding(operation, input.schema(), binding)?;
208        dimensions.push(dimension);
209        elements.push(element);
210    }
211    if elements.len() == 2 && elements[0] != elements[1] {
212        return Err(invalid_call(
213            operation,
214            "vector_element_type_mismatch",
215            format!(
216                "{} element {} does not match {} element {}",
217                operation.descriptor().relation_roles[0].as_str(),
218                elements[0].as_str(),
219                operation.descriptor().relation_roles[1].as_str(),
220                elements[1].as_str()
221            ),
222        ));
223    }
224    let element = elements[0];
225    if let Some(metric) = options.value_options().get("metric").and_then(Value::as_str)
226        && let Some(supported) = operation.descriptor().element_support(element)
227        && !supported.metrics.contains(&metric)
228    {
229        return Err(invalid_call(
230            operation,
231            "unsupported_vector_element_metric",
232            format!(
233                "{} {} does not support metric '{metric}'; valid metrics: {}",
234                operation.descriptor().sql_function_name,
235                element.as_str(),
236                supported.metrics.join(", ")
237            ),
238        ));
239    }
240    if matches!(
241        operation,
242        CuvsOperation::BruteForceKnn
243            | CuvsOperation::IvfFlat
244            | CuvsOperation::IvfPq
245            | CuvsOperation::Cagra
246            | CuvsOperation::PairwiseDistance
247    ) && let [Some(first), Some(second)] = dimensions.as_slice()
248        && first != second
249    {
250        return Err(invalid_call(
251            operation,
252            "vector_dimension_mismatch",
253            format!(
254                "{} dimension {first} does not match {} dimension {second}",
255                operation.descriptor().relation_roles[0].as_str(),
256                operation.descriptor().relation_roles[1].as_str()
257            ),
258        ));
259    }
260    let ivf_metric = match options.params() {
261        VectorAlgorithmParams::IvfFlat(params) => Some(params.metric),
262        VectorAlgorithmParams::IvfPq(params) => Some(params.metric),
263        _ => None,
264    };
265    if ivf_metric == Some(VectorDistanceMetric::CosineExpanded)
266        && let Some(Some(dimension)) = dimensions.first()
267        && *dimension < 2
268    {
269        return Err(invalid_call(
270            operation,
271            "insufficient_feature_dimension",
272            "cosine IVF search requires at least two dimensions",
273        ));
274    }
275    if let VectorAlgorithmParams::TopK(params) = options.params()
276        && let Some(Some(dimension)) = dimensions.first()
277        && params.k > *dimension
278    {
279        return Err(invalid_call(operation, "k_exceeds_dimension", "Top-K k exceeds the vector dimension"));
280    }
281    if operation == CuvsOperation::Pca
282        && let Some(Some(dimension)) = dimensions.first()
283        && *dimension < 2
284    {
285        return Err(invalid_call(
286            operation,
287            "insufficient_feature_dimension",
288            "cuVS PCA requires at least two input feature columns",
289        ));
290    }
291    Ok((dimensions, element))
292}
293
294fn validate_relation_binding(
295    operation: CuvsOperation,
296    schema: &DFSchema,
297    binding: &CuvsVectorRelationBinding,
298) -> Result<(Option<usize>, VectorElementType)> {
299    let role = binding.relation_role().as_str();
300    let id = resolve_field(operation, schema, role, binding.id())?;
301    if RelationIdType::from_data_type(id.data_type()).is_none() {
302        return Err(unsupported_type(
303            operation,
304            role,
305            binding.id(),
306            id.data_type(),
307            "Int32, Int64, Utf8, LargeUtf8, or Utf8View",
308        ));
309    }
310
311    match binding.vector() {
312        DenseVectorBinding::Columns(columns) => {
313            let mut element = None;
314            for column in columns {
315                let field = resolve_field(operation, schema, role, column)?;
316                let current = require_supported_vector_element(operation, role, column, field)?;
317                if element.is_some_and(|element| element != current) {
318                    return Err(invalid_call(
319                        operation,
320                        "vector_element_type_mismatch",
321                        format!(
322                            "cuVS {role} feature field '{column}' must have the same element type as the other feature fields"
323                        ),
324                    ));
325                }
326                element = Some(current);
327            }
328            Ok((Some(columns.len()), element.expect("validated non-empty columns")))
329        },
330        DenseVectorBinding::VectorColumn(column) => {
331            let field = resolve_field(operation, schema, role, column)?;
332            match field.data_type() {
333                DataType::FixedSizeList(element, dimension) if *dimension > 0 => {
334                    let element = require_supported_vector_element(operation, role, column, element)?;
335                    Ok((Some(*dimension as usize), element))
336                },
337                DataType::List(element) | DataType::LargeList(element) => {
338                    let element = require_supported_vector_element(operation, role, column, element)?;
339                    Ok((None, element))
340                },
341                actual => Err(unsupported_type(
342                    operation,
343                    role,
344                    column,
345                    actual,
346                    format!("FixedSizeList, List, or LargeList with {} elements", supported_vector_elements(operation)),
347                )),
348            }
349        },
350    }
351}
352
353fn resolve_field<'a>(
354    operation: CuvsOperation,
355    schema: &'a DFSchema,
356    role: &'static str,
357    name: &str,
358) -> Result<&'a FieldRef> {
359    schema.field_with_unqualified_name(name).map_err(|_| field_not_found(operation, role, name))
360}
361
362fn supported_vector_elements(operation: CuvsOperation) -> String {
363    operation
364        .descriptor()
365        .vector_element_support
366        .iter()
367        .map(|support| support.element.as_str())
368        .collect::<Vec<_>>()
369        .join(", ")
370}
371
372fn require_supported_vector_element(
373    operation: CuvsOperation,
374    role: &'static str,
375    name: &str,
376    field: &FieldRef,
377) -> Result<VectorElementType> {
378    let element = VectorElementType::from_data_type(field.data_type())
379        .filter(|element| operation.descriptor().element_support(*element).is_some())
380        .ok_or_else(|| {
381            unsupported_type(operation, role, name, field.data_type(), supported_vector_elements(operation))
382        })?;
383    Ok(element)
384}
385
386fn validate_known_cardinalities(
387    operation: CuvsOperation,
388    inputs: &[LogicalPlan],
389    options: &CuvsValidatedOptions,
390) -> Result<()> {
391    let Some(input_rows) = known_row_count(&inputs[0]) else {
392        return Ok(());
393    };
394    match options.params() {
395        VectorAlgorithmParams::BruteForceKnn(params) => {
396            if input_rows == 0 {
397                return Err(invalid_call(
398                    operation,
399                    "empty_dataset",
400                    "cuvs_brute_force_knn requires a non-empty dataset relation",
401                ));
402            }
403            if params.k > input_rows {
404                return Err(invalid_call(
405                    operation,
406                    "k_exceeds_dataset_rows",
407                    format!("k ({}) exceeds the statically known dataset row count ({input_rows})", params.k),
408                ));
409            }
410        },
411        VectorAlgorithmParams::IvfFlat(VectorIvfFlatParams { k, n_lists, .. })
412        | VectorAlgorithmParams::IvfPq(VectorIvfPqParams { k, n_lists, .. }) => {
413            if input_rows == 0 || *k > input_rows || *n_lists as usize > input_rows {
414                return Err(invalid_ann_dataset_cardinality(operation));
415            }
416        },
417        VectorAlgorithmParams::Cagra(params) => {
418            if input_rows == 0 || params.k > input_rows || params.intermediate_graph_degree as usize >= input_rows {
419                return Err(invalid_ann_dataset_cardinality(operation));
420            }
421        },
422        VectorAlgorithmParams::PairwiseDistance(_) | VectorAlgorithmParams::TopK(_) => {},
423        VectorAlgorithmParams::KMeans(params) => {
424            if input_rows == 0 {
425                return Err(invalid_call(operation, "empty_input", "cuvs_kmeans requires a non-empty input relation"));
426            }
427            if params.n_clusters > input_rows {
428                return Err(invalid_call(
429                    operation,
430                    "cluster_count_exceeds_input_rows",
431                    format!(
432                        "n_clusters ({}) exceeds the statically known input row count ({input_rows})",
433                        params.n_clusters
434                    ),
435                ));
436            }
437        },
438        VectorAlgorithmParams::Pca(_) => {
439            if input_rows < 2 {
440                return Err(invalid_call(
441                    operation,
442                    "insufficient_input_rows",
443                    "cuvs_pca requires at least two input rows",
444                ));
445            }
446        },
447    }
448    Ok(())
449}
450
451fn invalid_ann_dataset_cardinality(operation: CuvsOperation) -> Error {
452    invalid_call(
453        operation,
454        "invalid_dataset_cardinality",
455        "ANN dataset must be non-empty; k must not exceed rows and graph degrees must be less than rows",
456    )
457}
458
459fn derive_output_schema(
460    operation: CuvsOperation,
461    inputs: &[LogicalPlan],
462    options: &CuvsValidatedOptions,
463    dimensions: &[Option<usize>],
464    element: VectorElementType,
465) -> DataFusionResult<DFSchemaRef> {
466    let id_field = |ordinal: usize, output_name| {
467        output_id_field(inputs[ordinal].schema(), options.bindings()[ordinal].id(), output_name)
468            .map(|field| field.data_type().clone())
469    };
470    let schema = match options.params() {
471        VectorAlgorithmParams::BruteForceKnn(_)
472        | VectorAlgorithmParams::IvfFlat(_)
473        | VectorAlgorithmParams::IvfPq(_)
474        | VectorAlgorithmParams::Cagra(_) => {
475            vector_knn_output_schema(id_field(1, "query_id")?, id_field(0, "neighbor_id")?)
476        },
477        VectorAlgorithmParams::PairwiseDistance(_) => vector_pairwise_distance_output_schema(
478            id_field(0, "x_id")?,
479            id_field(1, "y_id")?,
480            match element {
481                VectorElementType::Float64 => DataType::Float64,
482                _ => DataType::Float32,
483            },
484        ),
485        VectorAlgorithmParams::TopK(_) => vector_top_k_output_schema(id_field(0, "id")?),
486        VectorAlgorithmParams::KMeans(_) => vector_kmeans_output_schema(id_field(0, "id")?),
487        VectorAlgorithmParams::Pca(params) => {
488            if let Some(dimension) = dimensions[0]
489                && params.n_components > dimension
490            {
491                return Err(invalid_call(
492                    operation,
493                    "component_count_exceeds_dimension",
494                    format!(
495                        "n_components ({}) exceeds the statically known feature dimension ({dimension})",
496                        params.n_components
497                    ),
498                )
499                .into());
500            }
501            vector_pca_output_schema(id_field(0, "id")?, params.n_components)
502        },
503    };
504    Ok(Arc::new(DFSchema::try_from(schema.as_ref().clone())?))
505}
506
507fn output_id_field(schema: &DFSchema, input_name: &str, output_name: &str) -> DataFusionResult<Field> {
508    let input = schema.field_with_unqualified_name(input_name)?;
509    Ok(Field::new(output_name, input.data_type().clone(), false))
510}
algeon_datafusion/cuvs_sql/logical.rs

algeon_datafusion/cuvs_sql/
logical.rs
