findepi commented on code in PR #12264:
URL: https://github.com/apache/datafusion/pull/12264#discussion_r1741825845


##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -1013,92 +1046,157 @@ impl LogicalPlan {
             }
             LogicalPlan::Distinct(distinct) => {
                 let distinct = match distinct {
-                    Distinct::All(_) => 
Distinct::All(Arc::new(inputs.swap_remove(0))),
+                    Distinct::All(_) => {
+                        self.check_no_expressions(expr);
+                        let input = self.only_input(inputs);
+                        Distinct::All(Arc::new(input))
+                    }
                     Distinct::On(DistinctOn {
                         on_expr,
                         select_expr,
                         ..
                     }) => {
+                        let input = self.only_input(inputs);
                         let sort_expr = expr.split_off(on_expr.len() + 
select_expr.len());
                         let select_expr = expr.split_off(on_expr.len());
                         assert!(sort_expr.is_empty(), "with_new_exprs for 
Distinct does not support sort expressions");
                         Distinct::On(DistinctOn::try_new(
                             expr,
                             select_expr,
                             None, // no sort expressions accepted
-                            Arc::new(inputs.swap_remove(0)),
+                            Arc::new(input),
                         )?)
                     }
                 };
                 Ok(LogicalPlan::Distinct(distinct))
             }
             LogicalPlan::RecursiveQuery(RecursiveQuery {
                 name, is_distinct, ..
-            }) => Ok(LogicalPlan::RecursiveQuery(RecursiveQuery {
-                name: name.clone(),
-                static_term: Arc::new(inputs.swap_remove(0)),
-                recursive_term: Arc::new(inputs.swap_remove(0)),
-                is_distinct: *is_distinct,
-            })),
+            }) => {
+                self.check_no_expressions(expr);
+                let (static_term, recursive_term) = 
self.only_two_inputs(inputs);
+                Ok(LogicalPlan::RecursiveQuery(RecursiveQuery {
+                    name: name.clone(),
+                    static_term: Arc::new(static_term),
+                    recursive_term: Arc::new(recursive_term),
+                    is_distinct: *is_distinct,
+                }))
+            }
             LogicalPlan::Analyze(a) => {
-                assert!(expr.is_empty());
-                assert_eq!(inputs.len(), 1);
+                self.check_no_expressions(expr);
+                let input = self.only_input(inputs);
                 Ok(LogicalPlan::Analyze(Analyze {
                     verbose: a.verbose,
                     schema: Arc::clone(&a.schema),
-                    input: Arc::new(inputs.swap_remove(0)),
+                    input: Arc::new(input),
                 }))
             }
             LogicalPlan::Explain(e) => {
-                assert!(
-                    expr.is_empty(),
-                    "Invalid EXPLAIN command. Expression should empty"
-                );
-                assert_eq!(inputs.len(), 1, "Invalid EXPLAIN command. Inputs 
are empty");
+                self.check_no_expressions(expr);
+                let input = self.only_input(inputs);
                 Ok(LogicalPlan::Explain(Explain {
                     verbose: e.verbose,
-                    plan: Arc::new(inputs.swap_remove(0)),
+                    plan: Arc::new(input),
                     stringified_plans: e.stringified_plans.clone(),
                     schema: Arc::clone(&e.schema),
                     logical_optimization_succeeded: 
e.logical_optimization_succeeded,
                 }))
             }
             LogicalPlan::Prepare(Prepare {
                 name, data_types, ..
-            }) => Ok(LogicalPlan::Prepare(Prepare {
-                name: name.clone(),
-                data_types: data_types.clone(),
-                input: Arc::new(inputs.swap_remove(0)),
-            })),
+            }) => {
+                self.check_no_expressions(expr);
+                let input = self.only_input(inputs);
+                Ok(LogicalPlan::Prepare(Prepare {
+                    name: name.clone(),
+                    data_types: data_types.clone(),
+                    input: Arc::new(input),
+                }))
+            }
             LogicalPlan::TableScan(ts) => {
-                assert!(inputs.is_empty(), "{self:?}  should have no inputs");
+                self.check_no_inputs(inputs);
                 Ok(LogicalPlan::TableScan(TableScan {
                     filters: expr,
                     ..ts.clone()
                 }))
             }
             LogicalPlan::EmptyRelation(_)
             | LogicalPlan::Ddl(_)
-            | LogicalPlan::Statement(_) => {
+            | LogicalPlan::Statement(_)
+            | LogicalPlan::DescribeTable(_) => {
                 // All of these plan types have no inputs / exprs so should 
not be called
-                assert!(expr.is_empty(), "{self:?} should have no exprs");
-                assert!(inputs.is_empty(), "{self:?}  should have no inputs");
+                self.check_no_expressions(expr);
+                self.check_no_inputs(inputs);
                 Ok(self.clone())
             }
-            LogicalPlan::DescribeTable(_) => Ok(self.clone()),
             LogicalPlan::Unnest(Unnest {
                 exec_columns: columns,
                 options,
                 ..
             }) => {
+                self.check_no_expressions(expr);
+                let input = self.only_input(inputs);
                 // Update schema with unnested column type.
-                let input = inputs.swap_remove(0);
                 let new_plan =
                     unnest_with_options(input, columns.clone(), 
options.clone())?;
                 Ok(new_plan)
             }
         }
     }
+
+    #[allow(clippy::needless_pass_by_value)] // expr is moved intentionally to 
ensure it's not used again
+    fn check_no_expressions(&self, expr: Vec<Expr>) {
+        assert!(
+            expr.is_empty(),
+            "{self:?} should have no exprs, got {:?}",
+            expr
+        );
+    }
+
+    #[allow(clippy::needless_pass_by_value)] // inputs is moved intentionally 
to ensure it's not used again
+    fn check_no_inputs(&self, inputs: Vec<LogicalPlan>) {
+        assert!(
+            inputs.is_empty(),
+            "{self:?} should have no inputs, got: {:?}",
+            inputs
+        );
+    }
+
+    fn only_expr(&self, mut expr: Vec<Expr>) -> Expr {

Review Comment:
   i had this initially inline but then the check code would take more space, 
and it was making it harder to see the important part of the logic (obviously 
this is subjective)
   another subjective bonus is there is type system help -- i cannot check 
there are no expressions or that there is only one expression and still use 
expressions vector to pull something later from it.
   
   if you would prefer to have these functions inlined back, just let me know, 
happy to do so
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org

Reply via email to