algeon_datafusion/cuvs_sql/
extension_planner.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use datafusion::{
5    catalog::Session,
6    common::Result,
7    logical_expr::{LogicalPlan, UserDefinedLogicalNode, physical_planning_context::PhysicalPlanningContext},
8    physical_expr::{LexOrdering, create_physical_sort_exprs},
9    physical_plan::ExecutionPlan,
10    physical_planner::{ExtensionPlanner, PhysicalPlanner},
11};
12
13use super::{CuvsAlgoExec, CuvsLogicalNode};
14use crate::session::planning_native_execution_config;
15
16#[derive(Debug)]
17/// DataFusion physical extension planner for validated cuVS logical nodes.
18///
19/// Calls outside this family remain available to other extension planners.
20pub struct CuvsExtensionPlanner;
21
22#[async_trait]
23impl ExtensionPlanner for CuvsExtensionPlanner {
24    async fn plan_extension(
25        &self,
26        _planner: &dyn PhysicalPlanner,
27        node: &dyn UserDefinedLogicalNode,
28        logical_inputs: &[&LogicalPlan],
29        physical_inputs: &[Arc<dyn ExecutionPlan>],
30        session: &dyn Session,
31        planning_ctx: &PhysicalPlanningContext,
32    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
33        let Some(node) = node.as_any().downcast_ref::<CuvsLogicalNode>() else {
34            return Ok(None);
35        };
36        let native_execution_config = planning_native_execution_config(session);
37        let input_orderings = node
38            .input_orderings()
39            .iter()
40            .zip(logical_inputs)
41            .map(|(ordering, input)| {
42                create_physical_sort_exprs(ordering, input.schema(), session.execution_props(), planning_ctx)
43                    .map(LexOrdering::new)
44            })
45            .collect::<Result<Vec<_>>>()?;
46        CuvsAlgoExec::try_new(
47            node.operation(),
48            node.validated_options().clone(),
49            physical_inputs.to_vec(),
50            input_orderings,
51            node.output_schema(),
52            native_execution_config,
53        )
54        .map(|exec| Some(Arc::new(exec) as Arc<dyn ExecutionPlan>))
55    }
56}
algeon_datafusion/cuvs_sql/extension_planner.rs

algeon_datafusion/cuvs_sql/
extension_planner.rs
