add sum(count) case in PushDownProjection

cnosdb · Oct 24, 2024 · ce4ee92 · ce4ee92
1 parent f6c50a7
commit ce4ee92
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/datafusion/optimizer/src/push_down_projection.rs b/datafusion/optimizer/src/push_down_projection.rs
@@ -264,7 +264,13 @@ impl OptimizerRule for PushDownProjection {
                         fun, args, ..
                     }) = &agg.aggr_expr[0]
                     {
-                        if matches!(fun, datafusion_expr::AggregateFunction::Count)
+                        if matches!(fun, datafusion_expr::AggregateFunction::Sum) && args.len() == 1 {
+                            if let Expr::Column(Column { relation, name }) = &args[0] {
+                                if relation.is_none() && name == "COUNT(UInt8(1))" {
+                                    new_aggr_expr.push(agg.aggr_expr[0].clone());
+                                }
+                            }
+                        } else if matches!(fun, datafusion_expr::AggregateFunction::Count)
                             && args.len() == 1
                             && args[0] == Expr::Literal(UInt8(Some(1)))
                         {