alamb commented on a change in pull request #9278:
URL: https://github.com/apache/arrow/pull/9278#discussion_r561366255
##########
File path: rust/datafusion/src/logical_plan/expr.rs
##########
@@ -422,6 +422,136 @@ impl Expr {
nulls_first,
}
}
+
+ /// Performs a depth first depth first walk of an expression and
+ /// its children, calling `visitor.pre_visit` and
+ /// `visitor.post_visit`.
+ ///
+ /// Implements the [visitor
pattern](https://en.wikipedia.org/wiki/Visitor_pattern) to
+ /// separate expression algorithms from the structure of the
+ /// `Expr` tree and make it easier to add new types of expressions
+ /// and algorithms that walk the tree.
+ ///
+ /// For an expression tree such as
+ /// BinaryExpr (GT)
+ /// left: Column("foo")
+ /// right: Column("bar")
+ ///
+ /// The nodes are visited using the following order
+ /// ```text
+ /// pre_visit(BinaryExpr(GT))
+ /// pre_visit(Column("foo"))
+ /// pre_visit(Column("bar"))
+ /// post_visit(Column("bar"))
+ /// post_visit(Column("bar"))
+ /// post_visit(BinaryExpr(GT))
+ /// ```
+ ///
+ /// If an Err result is returned, recursion is stopped immediately
+ ///
+ /// If `Recursion::Stop` is returned on a call to pre_visit, no
+ /// children of that expression are visited, nor is post_visit
+ /// called on that expression
+ ///
+ pub fn accept<V: ExpressionVisitor>(&self, visitor: V) -> Result<V> {
+ let visitor = match visitor.pre_visit(self)? {
+ Recursion::Continue(visitor) => visitor,
+ // If the recursion should stop, do not visit children
+ Recursion::Stop(visitor) => return Ok(visitor),
+ };
+
+ // recurse (and cover all expression types)
+ let visitor = match self {
+ Expr::Alias(expr, _) => expr.accept(visitor),
+ Expr::Column(..) => Ok(visitor),
+ Expr::ScalarVariable(..) => Ok(visitor),
+ Expr::Literal(..) => Ok(visitor),
+ Expr::BinaryExpr { left, right, .. } => {
+ let visitor = left.accept(visitor)?;
+ right.accept(visitor)
+ }
+ Expr::Not(expr) => expr.accept(visitor),
+ Expr::IsNotNull(expr) => expr.accept(visitor),
+ Expr::IsNull(expr) => expr.accept(visitor),
+ Expr::Negative(expr) => expr.accept(visitor),
+ Expr::Between {
+ expr, low, high, ..
+ } => {
+ let visitor = expr.accept(visitor)?;
+ let visitor = low.accept(visitor)?;
+ high.accept(visitor)
+ }
+ Expr::Case {
+ expr,
+ when_then_expr,
+ else_expr,
+ } => {
+ let visitor = if let Some(expr) = expr.as_ref() {
+ expr.accept(visitor)
+ } else {
+ Ok(visitor)
+ }?;
+ let visitor = when_then_expr.iter().try_fold(
+ visitor,
+ |visitor, (when, then)| {
+ let visitor = when.accept(visitor)?;
+ then.accept(visitor)
+ },
+ )?;
+ if let Some(else_expr) = else_expr.as_ref() {
+ else_expr.accept(visitor)
+ } else {
+ Ok(visitor)
+ }
+ }
+ Expr::Cast { expr, .. } => expr.accept(visitor),
+ Expr::Sort { expr, .. } => expr.accept(visitor),
+ Expr::ScalarFunction { args, .. } => args
+ .iter()
+ .try_fold(visitor, |visitor, arg| arg.accept(visitor)),
+ Expr::ScalarUDF { args, .. } => args
+ .iter()
+ .try_fold(visitor, |visitor, arg| arg.accept(visitor)),
+ Expr::AggregateFunction { args, .. } => args
+ .iter()
+ .try_fold(visitor, |visitor, arg| arg.accept(visitor)),
+ Expr::AggregateUDF { args, .. } => args
+ .iter()
+ .try_fold(visitor, |visitor, arg| arg.accept(visitor)),
+ Expr::InList { expr, list, .. } => {
+ let visitor = expr.accept(visitor)?;
+ list.iter()
+ .try_fold(visitor, |visitor, arg| arg.accept(visitor))
+ }
+ Expr::Wildcard => Ok(visitor),
+ }?;
+
+ visitor.post_visit(self)
+ }
+}
+
+/// Controls how the visitor recursion should proceed.
+pub enum Recursion<V: ExpressionVisitor> {
+ /// Attempt to visit all the children, recursively, of this expression.
+ Continue(V),
+ /// Do not visit the children of this expression, though the walk
+ /// of parents of this expression will not be affected
+ Stop(V),
+}
+
+/// Encode the traversal of an expression tree. When passed to
+/// `visit_expression`, `ExpressionVisitor::visit` is invoked
Review comment:
```suggestion
/// `Expr::accept`, `ExpressionVisitor::visit` is invoked
```
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
