[arrow-datafusion] branch master updated: Bug fix: expr_visitor was not visiting aggregate filter expressions (#3548)

[agrove]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new 1261741af Bug fix: expr_visitor was not visiting aggregate filter 
expressions (#3548)
1261741af is described below

commit 1261741af2a5e142fa0c7916e759859cc18ea59a
Author: Andy Grove <andygrov...@gmail.com>
AuthorDate: Tue Sep 20 12:10:07 2022 -0600

    Bug fix: expr_visitor was not visiting aggregate filter expressions (#3548)
---
 datafusion/expr/src/expr_visitor.rs              | 18 ++++++++++----
 datafusion/optimizer/src/projection_push_down.rs | 30 ++++++++++++++++++++++--
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/datafusion/expr/src/expr_visitor.rs 
b/datafusion/expr/src/expr_visitor.rs
index 0f2f5077a..3885456cc 100644
--- a/datafusion/expr/src/expr_visitor.rs
+++ b/datafusion/expr/src/expr_visitor.rs
@@ -176,12 +176,22 @@ impl ExprVisitable for Expr {
                     Ok(visitor)
                 }
             }
-            Expr::ScalarFunction { args, .. }
-            | Expr::ScalarUDF { args, .. }
-            | Expr::AggregateFunction { args, .. }
-            | Expr::AggregateUDF { args, .. } => args
+            Expr::ScalarFunction { args, .. } | Expr::ScalarUDF { args, .. } 
=> args
                 .iter()
                 .try_fold(visitor, |visitor, arg| arg.accept(visitor)),
+            Expr::AggregateFunction { args, filter, .. }
+            | Expr::AggregateUDF { args, filter, .. } => {
+                if let Some(f) = filter {
+                    let mut aggr_exprs = args.clone();
+                    aggr_exprs.push(f.as_ref().clone());
+                    aggr_exprs
+                        .iter()
+                        .try_fold(visitor, |visitor, arg| arg.accept(visitor))
+                } else {
+                    args.iter()
+                        .try_fold(visitor, |visitor, arg| arg.accept(visitor))
+                }
+            }
             Expr::WindowFunction {
                 args,
                 partition_by,
diff --git a/datafusion/optimizer/src/projection_push_down.rs 
b/datafusion/optimizer/src/projection_push_down.rs
index 5a0b38b57..86f73c061 100644
--- a/datafusion/optimizer/src/projection_push_down.rs
+++ b/datafusion/optimizer/src/projection_push_down.rs
@@ -532,9 +532,9 @@ mod tests {
     use crate::test::*;
     use arrow::datatypes::DataType;
     use datafusion_expr::{
-        col, lit,
+        col, count, lit,
         logical_plan::{builder::LogicalPlanBuilder, JoinType},
-        max, min, Expr,
+        max, min, AggregateFunction, Expr,
     };
     use std::collections::HashMap;
 
@@ -990,6 +990,32 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn aggregate_filter_pushdown() -> Result<()> {
+        let table_scan = test_table_scan()?;
+
+        let aggr_with_filter = Expr::AggregateFunction {
+            fun: AggregateFunction::Count,
+            args: vec![col("b")],
+            distinct: false,
+            filter: Some(Box::new(col("c").gt(lit(42)))),
+        };
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                vec![col("a")],
+                vec![count(col("b")), aggr_with_filter.alias("count2")],
+            )?
+            .build()?;
+
+        let expected = "Aggregate: groupBy=[[#test.a]], aggr=[[COUNT(#test.b), 
COUNT(#test.b) FILTER (WHERE #c > Int32(42)) AS count2]]\
+        \n  TableScan: test projection=[a, b, c]";
+
+        assert_optimized_plan_eq(&plan, expected);
+
+        Ok(())
+    }
+
     fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
         let optimized_plan = optimize(plan).expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);