fix COUNT(UInt8(1))

Author: jackwener

datafusion/core/tests/sql/window.rs

@@ -686,14 +686,15 @@ async fn test_remove_unnecessary_sort_in_sub_query() -> Result<()> {
             "    CoalescePartitionsExec",
             "      AggregateExec: mode=Partial, gby=[], aggr=[COUNT(UInt8(1))]",
             "        RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=8",
-            "          AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[]",
-            "            CoalesceBatchesExec: target_batch_size=4096",
-            "              RepartitionExec: partitioning=Hash([Column { name: \"c1\", index: 0 }], 8), input_partitions=8",
-            "                AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[]",
-            "                  ProjectionExec: expr=[c1@0 as c1]",
-            "                    CoalesceBatchesExec: target_batch_size=4096",
-            "                      FilterExec: c13@1 != C2GT5KVyOPZpgKVl110TyZO0NcJ434",
-            "                        RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1",
+            "          ProjectionExec: expr=[c1@0 as c1]",
+            "            AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[COUNT(UInt8(1))]",
+            "              CoalesceBatchesExec: target_batch_size=4096",
+            "                RepartitionExec: partitioning=Hash([Column { name: \"c1\", index: 0 }], 8), input_partitions=8",
+            "                  AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[COUNT(UInt8(1))]",
+            "                    ProjectionExec: expr=[c1@0 as c1]",
+            "                      CoalesceBatchesExec: target_batch_size=4096",
+            "                        FilterExec: c13@1 != C2GT5KVyOPZpgKVl110TyZO0NcJ434",
+            "                          RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1",
         ]
     };
 

datafusion/optimizer/src/push_down_projection.rs

@@ -23,9 +23,11 @@ use crate::optimizer::ApplyOrder;
 use crate::push_down_filter::replace_cols_by_name;
 use crate::{OptimizerConfig, OptimizerRule};
 use arrow::error::Result as ArrowResult;
+use datafusion_common::ScalarValue::UInt8;
 use datafusion_common::{
     Column, DFField, DFSchema, DFSchemaRef, DataFusionError, Result, ToDFSchema,
 };
+use datafusion_expr::expr::AggregateFunction;
 use datafusion_expr::utils::exprlist_to_fields;
 use datafusion_expr::{
     logical_plan::{Aggregate, LogicalPlan, Projection, TableScan, Union},
@@ -262,7 +264,6 @@ impl OptimizerRule for PushDownProjection {
                 generate_plan!(projection_is_empty, plan, new_alias)
             }
             LogicalPlan::Aggregate(agg) => {
-                // remove any aggregate expression that is not required
                 let mut required_columns = HashSet::new();
                 exprlist_to_columns(&projection.expr, &mut required_columns)?;
                 // Gather all columns needed for expressions in this Aggregate
@@ -274,6 +275,21 @@ impl OptimizerRule for PushDownProjection {
                     }
                 }
 
+                // if new_aggr_expr emtpy and aggr is COUNT(UInt8(1)), push it
+                if new_aggr_expr.is_empty() && agg.aggr_expr.len() == 1 {
+                    if let Expr::AggregateFunction(AggregateFunction {
+                        fun, args, ..
+                    }) = &agg.aggr_expr[0]
+                    {
+                        if matches!(fun, datafusion_expr::AggregateFunction::Count)
+                            && args.len() == 1
+                            && args[0] == Expr::Literal(UInt8(Some(1)))
+                        {
+                            new_aggr_expr.push(agg.aggr_expr[0].clone());
+                        }
+                    }
+                }
+
                 let new_agg = LogicalPlan::Aggregate(Aggregate::try_new(
                     agg.input.clone(),
                     agg.group_expr.clone(),