(arrow-datafusion) branch main updated: use `Expr::apply()` instead of `inspect_expr_pre()` (#9984)

[alamb]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 7acc8f16cf use `Expr::apply()` instead of `inspect_expr_pre()` (#9984)
7acc8f16cf is described below

commit 7acc8f16cf0776a4112a5e62214a44ad20c4c673
Author: Peter Toth <peter.t...@gmail.com>
AuthorDate: Sun Apr 7 18:03:27 2024 +0200

    use `Expr::apply()` instead of `inspect_expr_pre()` (#9984)
---
 datafusion/expr/src/utils.rs                     | 13 +++++++------
 datafusion/optimizer/src/optimize_projections.rs |  6 +++---
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs
index a93282574e..0d99d0b502 100644
--- a/datafusion/expr/src/utils.rs
+++ b/datafusion/expr/src/utils.rs
@@ -264,7 +264,7 @@ pub fn grouping_set_to_exprlist(group_expr: &[Expr]) -> 
Result<Vec<Expr>> {
 /// Recursively walk an expression tree, collecting the unique set of columns
 /// referenced in the expression
 pub fn expr_to_columns(expr: &Expr, accum: &mut HashSet<Column>) -> Result<()> 
{
-    inspect_expr_pre(expr, |expr| {
+    expr.apply(&mut |expr| {
         match expr {
             Expr::Column(qc) => {
                 accum.insert(qc.clone());
@@ -307,8 +307,9 @@ pub fn expr_to_columns(expr: &Expr, accum: &mut 
HashSet<Column>) -> Result<()> {
             | Expr::Placeholder(_)
             | Expr::OuterReferenceColumn { .. } => {}
         }
-        Ok(())
+        Ok(TreeNodeRecursion::Continue)
     })
+    .map(|_| ())
 }
 
 /// Find excluded columns in the schema, if any
@@ -838,11 +839,11 @@ pub fn find_column_exprs(exprs: &[Expr]) -> Vec<Expr> {
 
 pub(crate) fn find_columns_referenced_by_expr(e: &Expr) -> Vec<Column> {
     let mut exprs = vec![];
-    inspect_expr_pre(e, |expr| {
+    e.apply(&mut |expr| {
         if let Expr::Column(c) = expr {
             exprs.push(c.clone())
         }
-        Ok(()) as Result<()>
+        Ok(TreeNodeRecursion::Continue)
     })
     // As the closure always returns Ok, this "can't" error
     .expect("Unexpected error");
@@ -867,7 +868,7 @@ pub(crate) fn find_column_indexes_referenced_by_expr(
     schema: &DFSchemaRef,
 ) -> Vec<usize> {
     let mut indexes = vec![];
-    inspect_expr_pre(e, |expr| {
+    e.apply(&mut |expr| {
         match expr {
             Expr::Column(qc) => {
                 if let Ok(idx) = schema.index_of_column(qc) {
@@ -879,7 +880,7 @@ pub(crate) fn find_column_indexes_referenced_by_expr(
             }
             _ => {}
         }
-        Ok(()) as Result<()>
+        Ok(TreeNodeRecursion::Continue)
     })
     .unwrap();
     indexes
diff --git a/datafusion/optimizer/src/optimize_projections.rs 
b/datafusion/optimizer/src/optimize_projections.rs
index 147702cc04..69905c990a 100644
--- a/datafusion/optimizer/src/optimize_projections.rs
+++ b/datafusion/optimizer/src/optimize_projections.rs
@@ -34,7 +34,7 @@ use datafusion_expr::{
     Expr, Projection, TableScan, Window,
 };
 
-use datafusion_expr::utils::inspect_expr_pre;
+use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use hashbrown::HashMap;
 use itertools::{izip, Itertools};
 
@@ -613,7 +613,7 @@ fn rewrite_expr(expr: &Expr, input: &Projection) -> 
Result<Option<Expr>> {
 ///   columns are collected.
 fn outer_columns(expr: &Expr, columns: &mut HashSet<Column>) {
     // inspect_expr_pre doesn't handle subquery references, so find them 
explicitly
-    inspect_expr_pre(expr, |expr| {
+    expr.apply(&mut |expr| {
         match expr {
             Expr::OuterReferenceColumn(_, col) => {
                 columns.insert(col.clone());
@@ -632,7 +632,7 @@ fn outer_columns(expr: &Expr, columns: &mut 
HashSet<Column>) {
             }
             _ => {}
         };
-        Ok(()) as Result<()>
+        Ok(TreeNodeRecursion::Continue)
     })
     // unwrap: closure above never returns Err, so can not be Err here
     .unwrap();