algeon_datafusion/cuvs_sql/
relation_planner.rs

1use std::sync::Arc;
2
3use datafusion::{
4    common::Result as DataFusionResult,
5    logical_expr::{
6        Extension, LogicalPlan,
7        planner::{PlannedRelation, RelationPlanner, RelationPlannerContext, RelationPlanning},
8    },
9    sql::sqlparser::ast::TableFactor,
10};
11
12use super::{CuvsLogicalNode, CuvsOperation, options::dense_relation_function};
13use crate::gpu_functions::dense_relation::table_function_name;
14
15#[derive(Debug, Default)]
16/// DataFusion SQL relation planner that binds named cuVS function arguments.
17///
18/// Unrecognized relations retain their original DataFusion planning path.
19pub struct CuvsRelationPlanner;
20
21/// Creates a shared cuVS relation planner for installation on a session builder.
22pub fn cuvs_relation_planner() -> Arc<dyn RelationPlanner> {
23    Arc::new(CuvsRelationPlanner)
24}
25
26impl RelationPlanner for CuvsRelationPlanner {
27    fn plan_relation(
28        &self,
29        relation: TableFactor,
30        context: &mut dyn RelationPlannerContext,
31    ) -> DataFusionResult<RelationPlanning> {
32        let Some(operation) =
33            table_function_name(&relation, context).and_then(|name| CuvsOperation::from_sql_function_name(&name))
34        else {
35            return Ok(RelationPlanning::Original(Box::new(relation)));
36        };
37        let call = dense_relation_function(operation).plan_call(relation, context)?;
38        let node = CuvsLogicalNode::try_new(operation, call.inputs, &call.options_json)?
39            .with_input_orderings(call.input_orderings)?;
40        let plan = LogicalPlan::Extension(Extension { node: Arc::new(node) });
41        Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(plan, call.alias))))
42    }
43}
algeon_datafusion/cuvs_sql/relation_planner.rs

algeon_datafusion/cuvs_sql/
relation_planner.rs
