Fix aggregate type coercion bug

datafusion/core/tests/sql/subqueries.rs

@@ -147,8 +147,8 @@ order by s_acctbal desc, n_name, s_name, p_partkey;"#;
           Inner Join: #supplier.s_nationkey = #nation.n_nationkey
             Inner Join: #partsupp.ps_suppkey = #supplier.s_suppkey
               Inner Join: #part.p_partkey = #partsupp.ps_partkey
-                Filter: #part.p_size = Int32(15) AND #part.p_type LIKE Utf8("%BRASS")
-                  TableScan: part projection=[p_partkey, p_mfgr, p_type, p_size], partial_filters=[#part.p_size = Int32(15), #part.p_type LIKE Utf8("%BRASS")]
+                Filter: CAST(#part.p_size AS Int64) = Int64(15) AND #part.p_type LIKE Utf8("%BRASS")
+                  TableScan: part projection=[p_partkey, p_mfgr, p_type, p_size], partial_filters=[CAST(#part.p_size AS Int64) = Int64(15), #part.p_type LIKE Utf8("%BRASS")]
                 TableScan: partsupp projection=[ps_partkey, ps_suppkey, ps_supplycost]
               TableScan: supplier projection=[s_suppkey, s_name, s_address, s_nationkey, s_phone, s_acctbal, s_comment]
             TableScan: nation projection=[n_nationkey, n_name, n_regionkey]

datafusion/optimizer/src/optimizer.rs

@@ -178,16 +178,15 @@ impl Optimizer {
         F: FnMut(&LogicalPlan, &dyn OptimizerRule),
     {
         let mut new_plan = plan.clone();
-        debug!("Input logical plan:\n{}\n", plan.display_indent());
-        trace!("Full input logical plan:\n{:?}", plan);
+        log_plan("Optimizer input", plan);
+
         for rule in &self.rules {
             let result = rule.optimize(&new_plan, optimizer_config);
             match result {
                 Ok(plan) => {
                     new_plan = plan;
                     observer(&new_plan, rule.as_ref());
-                    debug!("After apply {} rule:\n", rule.name());
-                    debug!("Optimized logical plan:\n{}\n", new_plan.display_indent());
+                    log_plan(rule.name(), &new_plan);
                 }
                 Err(ref e) => {
                     if optimizer_config.skip_failing_rules {
@@ -209,12 +208,17 @@ impl Optimizer {
                 }
             }
         }
-        debug!("Optimized logical plan:\n{}\n", new_plan.display_indent());
-        trace!("Full Optimized logical plan:\n {:?}", new_plan);
+        log_plan("Optimized plan", &new_plan);
         Ok(new_plan)
     }
 }
 
+/// Log the plan in debug/tracing mode after some part of the optimizer runs
+fn log_plan(description: &str, plan: &LogicalPlan) {
+    debug!("{description}:\n{}\n", plan.display_indent());
+    trace!("{description}::\n{}\n", plan.display_indent_schema());
+}
+
 #[cfg(test)]
 mod tests {
     use crate::optimizer::Optimizer;

datafusion/optimizer/src/unwrap_cast_in_comparison.rs

@@ -97,7 +97,18 @@ fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
     let new_exprs = plan
         .expressions()
         .into_iter()
-        .map(|expr| expr.rewrite(&mut expr_rewriter))
+        .map(|expr| {
+            let original_name = expr.name()?;
+            let expr = expr.rewrite(&mut expr_rewriter)?;
+
+            // Ensure this rewrite doesn't change the name
+            // https://github.com/apache/arrow-datafusion/issues/3704
+            if expr.name()? != original_name {
+                Ok(expr.alias(&original_name))
+            } else {
+                Ok(expr)
+            }
+        })
         .collect::<Result<Vec<_>>>()?;
 
     from_plan(plan, new_exprs.as_slice(), new_inputs.as_slice())

datafusion/optimizer/tests/integration-test.rs

@@ -29,6 +29,12 @@ use std::any::Any;
 use std::collections::HashMap;
 use std::sync::Arc;
 
+#[cfg(test)]
+#[ctor::ctor]
+fn init() {
+    let _ = env_logger::try_init();
+}
+
 #[test]
 fn case_when() -> Result<()> {
     let sql = "SELECT CASE WHEN col_int32 > 0 THEN 1 ELSE 0 END FROM test";
@@ -45,6 +51,17 @@ fn case_when() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn case_when_aggregate() -> Result<()> {
+    let sql = "SELECT col_utf8, SUM(CASE WHEN col_int32 > 0 THEN 1 ELSE 0 END) AS n FROM test GROUP BY col_utf8";
+    let plan = test_sql(sql)?;
+    let expected = "Projection: #test.col_utf8, #SUM(CASE WHEN test.col_int32 > Int64(0) THEN Int64(1) ELSE Int64(0) END) AS n\
+    \n  Aggregate: groupBy=[[#test.col_utf8]], aggr=[[SUM(CASE WHEN #test.col_int32 > Int32(0) THEN Int64(1) ELSE Int64(0) END) AS SUM(CASE WHEN test.col_int32 > Int64(0) THEN Int64(1) ELSE Int64(0) END)]]\
+    \n    TableScan: test projection=[col_int32, col_utf8]";
+    assert_eq!(expected, format!("{:?}", plan));
+    Ok(())
+}
+
 #[test]
 fn unsigned_target_type() -> Result<()> {
     let sql = "SELECT * FROM test WHERE col_uint32 > 0";