algeon_datafusion/cuvs_sql/exec/
node.rs

1use std::{
2    collections::{HashMap, HashSet},
3    sync::{
4        Arc,
5        atomic::{AtomicU64, Ordering},
6    },
7};
8
9use arrow::datatypes::SchemaRef;
10use datafusion::{common::Result as DataFusionResult, physical_expr::LexOrdering, physical_plan::ExecutionPlan};
11use query_planning::plan::QueryPlan;
12
13use super::{
14    super::{
15        CuvsOperation, CuvsValidatedOptions, CuvsVectorRelationBinding, DenseVectorBinding,
16        error::{feature_disabled, invalid_call},
17    },
18    lowering::lower_cuvs_query_plan,
19};
20use crate::{
21    error::{Error, Result},
22    native_exec::{
23        AlgeonNativeExecutionConfig,
24        child_inputs::NativeChildBinding,
25        relation_function::{
26            RelationFunctionExec, RelationFunctionMessages, RelationFunctionOperations, RelationFunctionSpec,
27            RelationFunctionText,
28        },
29    },
30};
31
32static NEXT_CUVS_GPU_RESULT_INPUT_ID: AtomicU64 = AtomicU64::new(0);
33
34#[derive(Clone, Debug)]
35/// Validated operation settings retained by the specialized physical execution wrapper.
36pub struct CuvsRelationFunctionSpec {
37    operation: CuvsOperation,
38    options: CuvsValidatedOptions,
39}
40
41/// Physical execution wrapper for one cuVS call with validated input relations.
42///
43/// Native execution stages the declared inputs and exposes the operation's Arrow result schema.
44pub type CuvsAlgoExec = RelationFunctionExec<CuvsRelationFunctionSpec>;
45
46impl RelationFunctionExec<CuvsRelationFunctionSpec> {
47    pub(crate) fn try_new(
48        operation: CuvsOperation,
49        options: CuvsValidatedOptions,
50        inputs: Vec<Arc<dyn ExecutionPlan>>,
51        input_orderings: Vec<Option<LexOrdering>>,
52        output_schema: SchemaRef,
53        native_execution_config: AlgeonNativeExecutionConfig,
54    ) -> DataFusionResult<Self> {
55        let id = NEXT_CUVS_GPU_RESULT_INPUT_ID.fetch_add(1, Ordering::Relaxed);
56        let input_names =
57            (0..inputs.len()).map(|index| format!("__algeon_datafusion_cuvs_input_{id}_{index}")).collect();
58        Self::try_new_with_names(
59            CuvsRelationFunctionSpec { operation, options },
60            inputs,
61            output_schema,
62            native_execution_config,
63            input_names,
64            input_orderings,
65            format!("__algeon_datafusion_cuvs_result_{id}"),
66        )
67    }
68
69    /// Returns the registered operation executed by this physical wrapper.
70    pub fn operation(&self) -> CuvsOperation {
71        self.spec().operation
72    }
73}
74
75impl RelationFunctionSpec for CuvsRelationFunctionSpec {
76    const EXEC_NAME: &'static str = "CuvsAlgoExec";
77    const HOST_OUTPUT_MEMORY_CONSUMER: &'static str = "CuvsAlgoExec[host-output-staging]";
78    fn operation_name(&self) -> &'static str {
79        self.operation.descriptor().sql_function_name
80    }
81    fn binding_count(&self) -> usize {
82        self.options.bindings().len()
83    }
84    fn roles(&self) -> Vec<&'static str> {
85        self.operation.descriptor().relation_roles.iter().map(|role| role.as_str()).collect()
86    }
87    fn required_projection(&self, index: usize, schema: SchemaRef) -> Result<Vec<usize>> {
88        required_projection(self.operation, schema, &self.options.bindings()[index])
89    }
90    fn lower(&self, bindings: &[NativeChildBinding], schema: SchemaRef) -> Result<QueryPlan> {
91        lower_cuvs_query_plan(self.operation, &self.options, bindings, schema)
92    }
93    fn feature_available(&self) -> bool {
94        cfg!(feature = "cuvs")
95    }
96    fn feature_disabled(&self) -> Error {
97        feature_disabled(self.operation)
98    }
99    fn invalid_child_count(&self) -> Error {
100        invalid_call(
101            self.operation,
102            "physical_child_count_mismatch",
103            "cuVS physical inputs do not match validated relation bindings",
104        )
105    }
106    const TEXT: RelationFunctionText = RelationFunctionText {
107        operations: RelationFunctionOperations {
108            execute: "algeon_datafusion.cuvs_sql.execute",
109            execute_gpu: "algeon_datafusion.cuvs_sql.execute_gpu",
110            gpu_execution_pool: "algeon_datafusion.cuvs_sql.gpu_execution_pool",
111            host_batch_cursor: "algeon_datafusion.cuvs_sql.host_batch_cursor",
112            preflight: "algeon_datafusion.cuvs_sql.preflight",
113            with_new_children: "algeon_datafusion.cuvs_sql.with_new_children",
114        },
115        messages: RelationFunctionMessages {
116            preflight_failed: "native cuVS query-plan preflight failed",
117            staged_inputs_missing: "cuVS GPU-resident producer inputs were not prepared before native context acquisition",
118            staged_inputs_poisoned: "cuVS staged-input ownership was poisoned",
119            staged_inputs_consumed: "cuVS staged inputs were already consumed",
120            execution_failed: "native cuVS vector query execution failed",
121            gpu_partition: "cuVS GPU-resident producer exposes exactly one output partition",
122            unexpected_host_result: "native cuVS producer materialized an unexpected host result",
123            child_count: "cuVS execution received an unexpected child count",
124            host_partition: "cuVS execution exposes exactly one output partition",
125            gpu_pool_closed: "query service closed before cuVS acquired its GPU execution pool",
126            deferred_result: "deferred cuVS host execution did not produce a GPU-resident cursor input",
127            cursor_failed: "failed to bind the deferred cuVS host-output cursor",
128        },
129    };
130}
131
132fn required_projection(
133    operation: CuvsOperation,
134    schema: SchemaRef,
135    binding: &CuvsVectorRelationBinding,
136) -> Result<Vec<usize>> {
137    let mut names = Vec::new();
138    names.push(binding.id());
139    match binding.vector() {
140        DenseVectorBinding::Columns(columns) => names.extend(columns.iter().map(String::as_str)),
141        DenseVectorBinding::VectorColumn(column) => names.push(column),
142    }
143    let mut field_indexes = HashMap::with_capacity(schema.fields().len());
144    for (index, field) in schema.fields().iter().enumerate() {
145        field_indexes.entry(field.name().as_str()).or_insert(index);
146    }
147    let mut seen = HashSet::new();
148    names
149        .into_iter()
150        .filter(|name| seen.insert(*name))
151        .map(|name| {
152            field_indexes.get(name).copied().ok_or_else(|| {
153                invalid_call(
154                    operation,
155                    "physical_projection_field_missing",
156                    format!("cuVS physical input has no field named '{name}'"),
157                )
158            })
159        })
160        .collect()
161}
algeon_datafusion/cuvs_sql/exec/node.rs

algeon_datafusion/cuvs_sql/exec/
node.rs
