jayzhan211 commented on code in PR #12264:
URL: https://github.com/apache/datafusion/pull/12264#discussion_r1740355186
##########
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 doubt whether this is helpful, a little over abstraction IMO 🤔
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
