peter-toth commented on code in PR #7942:
URL: https://github.com/apache/arrow-datafusion/pull/7942#discussion_r1432829464
##########
datafusion/expr/src/logical_plan/plan.rs:
##########
@@ -311,60 +309,63 @@ impl LogicalPlan {
exprs
}
- /// Calls `f` on all expressions (non-recursively) in the current
- /// logical plan node. This does not include expressions in any
- /// children.
- pub fn inspect_expressions<F, E>(self: &LogicalPlan, mut f: F) ->
Result<(), E>
+ /// Apply `f` on expressions of the plan node.
+ /// `f` is not allowed to return [`TreeNodeRecursion::Prune`].
+ pub fn apply_expressions<F>(&self, f: &mut F) -> Result<TreeNodeRecursion>
where
- F: FnMut(&Expr) -> Result<(), E>,
+ F: FnMut(&Expr) -> Result<TreeNodeRecursion>,
{
+ let f = &mut |e: &Expr| f(e)?.fail_on_prune();
+
match self {
LogicalPlan::Projection(Projection { expr, .. }) => {
- expr.iter().try_for_each(f)
+ expr.iter().for_each_till_continue(f)
}
LogicalPlan::Values(Values { values, .. }) => {
- values.iter().flatten().try_for_each(f)
+ values.iter().flatten().for_each_till_continue(f)
}
LogicalPlan::Filter(Filter { predicate, .. }) => f(predicate),
LogicalPlan::Repartition(Repartition {
partitioning_scheme,
..
}) => match partitioning_scheme {
- Partitioning::Hash(expr, _) => expr.iter().try_for_each(f),
- Partitioning::DistributeBy(expr) =>
expr.iter().try_for_each(f),
- Partitioning::RoundRobinBatch(_) => Ok(()),
+ Partitioning::Hash(expr, _) =>
expr.iter().for_each_till_continue(f),
+ Partitioning::DistributeBy(expr) =>
expr.iter().for_each_till_continue(f),
+ Partitioning::RoundRobinBatch(_) =>
Ok(TreeNodeRecursion::Continue),
},
LogicalPlan::Window(Window { window_expr, .. }) => {
- window_expr.iter().try_for_each(f)
+ window_expr.iter().for_each_till_continue(f)
}
LogicalPlan::Aggregate(Aggregate {
group_expr,
aggr_expr,
..
- }) => group_expr.iter().chain(aggr_expr.iter()).try_for_each(f),
+ }) => group_expr
+ .iter()
+ .chain(aggr_expr.iter())
+ .for_each_till_continue(f),
// There are two part of expression for join, equijoin(on) and
non-equijoin(filter).
// 1. the first part is `on.len()` equijoin expressions, and the
struct of each expr is `left-on = right-on`.
// 2. the second part is non-equijoin(filter).
LogicalPlan::Join(Join { on, filter, .. }) => {
on.iter()
// it not ideal to create an expr here to analyze them,
but could cache it on the Join itself
.map(|(l, r)| Expr::eq(l.clone(), r.clone()))
- .try_for_each(|e| f(&e))?;
-
- if let Some(filter) = filter.as_ref() {
- f(filter)
- } else {
- Ok(())
- }
+ .for_each_till_continue(&mut |e| f(&e))?
+ .and_then_on_continue(||
filter.iter().for_each_till_continue(f))
}
- LogicalPlan::Sort(Sort { expr, .. }) =>
expr.iter().try_for_each(f),
+ LogicalPlan::Sort(Sort { expr, .. }) =>
expr.iter().for_each_till_continue(f),
LogicalPlan::Extension(extension) => {
// would be nice to avoid this copy -- maybe can
// update extension to just observer Exprs
- extension.node.expressions().iter().try_for_each(f)
+ extension
+ .node
+ .expressions()
+ .iter()
+ .for_each_till_continue(f)
Review Comment:
I was trying define clear APIs on `TreeNode`s. After
https://github.com/apache/arrow-datafusion/pull/7942/commits/f4d28e0d7a340c037e3e9729cd788a28149cfb1e
we have
- `visit()`, `visit_down()`,
- `transform()`, `transform_down()`, `transform_up()`,
- `transform_with_payload()`, `transform_down_with_payload()` and
`transform_up_with_payload()`
functions on `TreeNode` and all can be controlled with `TreeNodeRecursion`.
--
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]
