datafusion-fedeartion: Remove assert_eq macros and handle errors prop…

…erly (datafusion-contrib#49)
spiceai · Aug 26, 2024 · 48dbea5 · 48dbea5
1 parent f4babee
commit 48dbea5
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/datafusion-federation/src/plan_node.rs b/datafusion-federation/src/plan_node.rs
@@ -8,7 +8,7 @@ use std::{
 use async_trait::async_trait;
 use datafusion::{
     common::DFSchemaRef,
-    error::Result,
+    error::{DataFusionError, Result},
     execution::context::{QueryPlanner, SessionState},
     logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore},
     physical_plan::ExecutionPlan,
@@ -57,17 +57,14 @@ impl UserDefinedLogicalNodeCore for FederatedPlanNode {
         write!(f, "Federated\n {:?}", self.plan)
     }
 
-    fn from_template(&self, exprs: &[Expr], inputs: &[LogicalPlan]) -> Self {
-        assert_eq!(inputs.len(), 0, "input size inconsistent");
-        assert_eq!(exprs.len(), 0, "expression size inconsistent");
-        Self {
-            plan: self.plan.clone(),
-            planner: self.planner.clone(),
+    fn with_exprs_and_inputs(&self, exprs: Vec<Expr>, inputs: Vec<LogicalPlan>) -> Result<Self> {
+        if !inputs.is_empty() {
+            return Err(DataFusionError::Plan("input size inconsistent".into()));
+        }
+        if !exprs.is_empty() {
+            return Err(DataFusionError::Plan("expression size inconsistent".into()));
         }
-    }
 
-    /// XXX should consider something else here ?
-    fn with_exprs_and_inputs(&self, _exprs: Vec<Expr>, _inputs: Vec<LogicalPlan>) -> Result<Self> {
         Ok(Self {
             plan: self.plan.clone(),
             planner: self.planner.clone(),
@@ -149,8 +146,11 @@ impl ExtensionPlanner for FederatedPlanner {
     ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
         let dc_node = node.as_any().downcast_ref::<FederatedPlanNode>();
         if let Some(fed_node) = dc_node {
-            assert_eq!(logical_inputs.len(), 0, "Inconsistent number of inputs");
-            assert_eq!(physical_inputs.len(), 0, "Inconsistent number of inputs");
+            if !logical_inputs.is_empty() || !physical_inputs.is_empty() {
+                return Err(DataFusionError::Plan(
+                    "Inconsistent number of inputs".into(),
+                ));
+            }
 
             let fed_planner = fed_node.planner.clone();
             let exec_plan = fed_planner.plan_federation(fed_node, session_state).await?;