datafusion_federation/
plan_node.rs

1use core::fmt;
2use std::{
3    fmt::Debug,
4    hash::{Hash, Hasher},
5    sync::Arc,
6};
7
8use async_trait::async_trait;
9use datafusion::{
10    catalog::Session,
11    common::DFSchemaRef,
12    error::{DataFusionError, Result},
13    execution::context::QueryPlanner,
14    logical_expr::{
15        physical_planning_context::PhysicalPlanningContext, Expr, LogicalPlan,
16        UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
17    },
18    physical_plan::ExecutionPlan,
19    physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner},
20};
21
22pub struct FederatedPlanNode {
23    pub plan: LogicalPlan,
24    pub planner: Arc<dyn FederationPlanner>,
25}
26
27impl FederatedPlanNode {
28    pub fn new(plan: LogicalPlan, planner: Arc<dyn FederationPlanner>) -> Self {
29        Self { plan, planner }
30    }
31
32    pub fn plan(&self) -> &LogicalPlan {
33        &self.plan
34    }
35}
36
37impl Debug for FederatedPlanNode {
38    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
39        UserDefinedLogicalNodeCore::fmt_for_explain(self, f)
40    }
41}
42
43impl UserDefinedLogicalNodeCore for FederatedPlanNode {
44    fn name(&self) -> &str {
45        "Federated"
46    }
47
48    fn inputs(&self) -> Vec<&LogicalPlan> {
49        Vec::new()
50    }
51
52    fn schema(&self) -> &DFSchemaRef {
53        self.plan.schema()
54    }
55
56    fn expressions(&self) -> Vec<Expr> {
57        Vec::new()
58    }
59
60    fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result {
61        write!(f, "Federated\n {}", self.plan)
62    }
63
64    fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> Result<Self> {
65        if !inputs.is_empty() {
66            return Err(DataFusionError::Plan("input size inconsistent".into()));
67        }
68        if !exprs.is_empty() {
69            return Err(DataFusionError::Plan("expression size inconsistent".into()));
70        }
71
72        Ok(Self {
73            plan: self.plan.clone(),
74            planner: self.planner.clone(),
75        })
76    }
77}
78
79#[derive(Default, Debug)]
80pub struct FederatedQueryPlanner {}
81
82impl FederatedQueryPlanner {
83    pub fn new() -> Self {
84        Self::default()
85    }
86}
87
88#[async_trait]
89impl QueryPlanner for FederatedQueryPlanner {
90    async fn create_physical_plan(
91        &self,
92        logical_plan: &LogicalPlan,
93        session_state: &dyn Session,
94    ) -> Result<Arc<dyn ExecutionPlan>> {
95        // Get provider here?
96
97        let physical_planner =
98            DefaultPhysicalPlanner::with_extension_planners(vec![
99                Arc::new(FederatedPlanner::new()),
100            ]);
101        physical_planner
102            .create_physical_plan(logical_plan, session_state)
103            .await
104    }
105}
106
107#[async_trait]
108pub trait FederationPlanner: Send + Sync {
109    async fn plan_federation(
110        &self,
111        node: &FederatedPlanNode,
112        session_state: &dyn Session,
113    ) -> Result<Arc<dyn ExecutionPlan>>;
114}
115
116impl std::fmt::Debug for dyn FederationPlanner {
117    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
118        write!(f, "FederationPlanner")
119    }
120}
121
122impl PartialEq<FederatedPlanNode> for FederatedPlanNode {
123    /// Comparing name, args and return_type
124    fn eq(&self, other: &FederatedPlanNode) -> bool {
125        self.plan == other.plan
126    }
127}
128
129impl PartialOrd<FederatedPlanNode> for FederatedPlanNode {
130    fn partial_cmp(&self, other: &FederatedPlanNode) -> Option<std::cmp::Ordering> {
131        self.plan.partial_cmp(&other.plan)
132    }
133}
134
135impl Eq for FederatedPlanNode {}
136
137impl Hash for FederatedPlanNode {
138    fn hash<H: Hasher>(&self, state: &mut H) {
139        self.plan.hash(state);
140    }
141}
142
143#[derive(Default)]
144pub struct FederatedPlanner {}
145
146impl FederatedPlanner {
147    pub fn new() -> Self {
148        Self::default()
149    }
150}
151
152#[async_trait]
153impl ExtensionPlanner for FederatedPlanner {
154    async fn plan_extension(
155        &self,
156        _planner: &dyn PhysicalPlanner,
157        node: &dyn UserDefinedLogicalNode,
158        logical_inputs: &[&LogicalPlan],
159        physical_inputs: &[Arc<dyn ExecutionPlan>],
160        session_state: &dyn Session,
161        _planning_ctx: &PhysicalPlanningContext,
162    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
163        let dc_node = node.as_any().downcast_ref::<FederatedPlanNode>();
164        if let Some(fed_node) = dc_node {
165            if !logical_inputs.is_empty() || !physical_inputs.is_empty() {
166                return Err(DataFusionError::Plan(
167                    "Inconsistent number of inputs".into(),
168                ));
169            }
170
171            let fed_planner = Arc::clone(&fed_node.planner);
172            let exec_plan = fed_planner.plan_federation(fed_node, session_state).await?;
173            return Ok(Some(exec_plan));
174        }
175        Ok(None)
176    }
177}
datafusion_federation/plan_node.rs

datafusion_federation/
plan_node.rs
