feat: rewrite count(non-nullable column) to count(1)

datafusion/core/tests/datasource/object_store_access.rs

@@ -965,8 +965,8 @@ async fn query_single_parquet_file() {
     RequestCountingObjectStore()
     Total Requests: 3
     - GET  (opts) path=parquet_table.parquet head=true
-    - GET  (ranges) path=parquet_table.parquet ranges=4-534,534-1064
-    - GET  (ranges) path=parquet_table.parquet ranges=1064-1594,1594-2124
+    - GET  (ranges) path=parquet_table.parquet ranges=4-534
+    - GET  (ranges) path=parquet_table.parquet ranges=1064-1594
     "
     );
 }

datafusion/core/tests/sql/explain_analyze.rs

@@ -1137,7 +1137,7 @@ async fn explain_analyze_aggregate_metrics_map_indices_to_expressions() {
         .to_string();
     assert_contains!(
         normal.as_str(),
-        "aggr=[sum(aggregate_test_100.c5), sum(aggregate_test_100.c6), count(aggregate_test_100.c7)]"
+        "aggr=[sum(aggregate_test_100.c5), sum(aggregate_test_100.c6), count(1) as count(aggregate_test_100.c7)]"
     );
     assert_contains!(normal.as_str(), "agg_expr_0_arguments_time");
     assert_contains!(normal.as_str(), "agg_expr_1_arguments_time");
@@ -1159,7 +1159,7 @@ async fn explain_analyze_aggregate_metrics_map_indices_to_expressions() {
     );
     assert_contains!(
         verbose.as_str(),
-        "agg_expr_2_arguments_time{partition=0, aggregate=count(aggregate_test_100.c7)}"
+        "agg_expr_2_arguments_time{partition=0, aggregate=count(1) as count(aggregate_test_100.c7)}"
     );
 }
 

datafusion/core/tests/sql/unparser.rs

@@ -567,9 +567,8 @@ async fn optimized_duckdb_unparse_preserves_nested_aggregate_scope() -> Result<(
 
     assert!(
         sql.contains(concat!(
-            r#"FROM (SELECT sum("total_revenue") AS "alias2", "#,
-            r#"date_part('year', "signup_date") AS "group_alias_0", "#,
-            r#""customer_id" AS "alias1" "#
+            r#"FROM (SELECT date_part('year', "signup_date") AS "group_alias_0", "#,
+            r#"sum("total_revenue") AS "alias2" "#
         )),
         "inner aggregate should define the aliases before the outer aggregate uses them: {sql}",
     );

datafusion/functions-aggregate/src/count.rs

@@ -39,7 +39,8 @@ use datafusion_expr::{
     GroupsAccumulator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs,
     TypeSignature, Volatility, WindowFunctionDefinition,
     expr::WindowFunction,
-    function::{AccumulatorArgs, StateFieldsArgs},
+    function::{AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs},
+    simplify::SimplifyContext,
     utils::{AggregateOrderSensitivity, format_state_name},
 };
 use datafusion_functions_aggregate_common::aggregate::count_distinct::PrimitiveDistinctCountGroupsAccumulator;
@@ -400,6 +401,18 @@ impl AggregateUDFImpl for Count {
         AggregateOrderSensitivity::Insensitive
     }
 
+    fn simplify(&self) -> Option<AggregateFunctionSimplification> {
+        Some(Box::new(|mut aggregate_function, info| {
+            if !aggregate_function.params.distinct
+                && can_rewrite_count_args(&aggregate_function.params.args, info)
+            {
+                aggregate_function.params.args =
+                    vec![Expr::Literal(COUNT_STAR_EXPANSION, None)];
+            }
+            Ok(Expr::AggregateFunction(aggregate_function))
+        }))
+    }
+
     fn default_value(&self, _data_type: &DataType) -> Result<ScalarValue> {
         Ok(ScalarValue::Int64(Some(0)))
     }
@@ -470,6 +483,20 @@ impl AggregateUDFImpl for Count {
     }
 }
 
+fn can_rewrite_count_args(args: &[Expr], info: &SimplifyContext) -> bool {
+    let mut has_non_nullable_column = false;
+    for arg in args {
+        match arg {
+            Expr::Column(_) if matches!(info.nullable(arg), Ok(false)) => {
+                has_non_nullable_column = true;
+            }
+            Expr::Literal(value, _) if !value.is_null() => {}
+            _ => return false,
+        }
+    }
+    has_non_nullable_column
+}
+
 #[cold]
 fn create_distinct_count_groups_accumulator(
     args: &AccumulatorArgs,
@@ -965,7 +992,9 @@ mod tests {
         array::{DictionaryArray, Int32Array, Int64Array, NullArray, StringArray},
         datatypes::{DataType, Field, Int32Type, Schema},
     };
+    use datafusion_common::DFSchema;
     use datafusion_expr::function::AccumulatorArgs;
+    use datafusion_expr::{col, lit};
     use datafusion_physical_expr::{PhysicalExpr, expressions::Column};
     use std::sync::Arc;
     /// Helper function to create a dictionary array with non-null keys but some null values
@@ -1007,6 +1036,47 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn simplify_count_non_nullable_columns() -> Result<()> {
+        let schema = DFSchema::try_from(Schema::new(vec![
+            Field::new("a", DataType::Int32, false),
+            Field::new("b", DataType::Int32, true),
+        ]))?;
+        let info = SimplifyContext::builder()
+            .with_schema(Arc::new(schema))
+            .build();
+        let simplify = Count::new().simplify().unwrap();
+        let simplified_args = |args: Vec<Expr>, distinct: bool| -> Result<Vec<Expr>> {
+            let aggregate_function = datafusion_expr::expr::AggregateFunction::new_udf(
+                count_udaf(),
+                args,
+                distinct,
+                None,
+                vec![],
+                None,
+            );
+            match simplify(aggregate_function, &info)? {
+                Expr::AggregateFunction(f) => Ok(f.params.args),
+                other => internal_err!("unexpected expression {other}"),
+            }
+        };
+        let count_one = vec![Expr::Literal(COUNT_STAR_EXPANSION, None)];
+
+        assert_eq!(simplified_args(vec![col("a")], false)?, count_one);
+        assert_eq!(simplified_args(vec![col("a"), lit(2)], false)?, count_one);
+
+        for args in [
+            vec![lit(1)],
+            vec![col("b")],
+            vec![col("a"), col("b")],
+            vec![col("a") + lit(1)],
+        ] {
+            assert_eq!(simplified_args(args.clone(), false)?, args);
+        }
+        assert_eq!(simplified_args(vec![col("a")], true)?, vec![col("a")]);
+        Ok(())
+    }
+
     #[test]
     fn count_groups_preserving_reads() -> Result<()> {
         let mut accumulator = CountGroupsAccumulator::new();

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -1679,7 +1679,13 @@ impl TreeNodeRewriter for Simplifier<'_> {
                 ..
             }) => match (func.simplify(), expr) {
                 (Some(simplify_function), Expr::AggregateFunction(af)) => {
-                    Transformed::yes(simplify_function(af, info)?)
+                    let original = Expr::AggregateFunction(af.clone());
+                    let simplified = simplify_function(af, info)?;
+                    if simplified == original {
+                        Transformed::no(simplified)
+                    } else {
+                        Transformed::yes(simplified)
+                    }
                 }
                 (_, expr) => Transformed::no(expr),
             },
@@ -5628,23 +5634,47 @@ mod tests {
 
         let expected = aggregate_function_expr.clone();
         assert_eq!(simplify(aggregate_function_expr), expected);
+
+        let udaf = AggregateUDF::new_from_impl(SimplifyMockUdaf::new_with_noop());
+        let aggregate_function_expr =
+            Expr::AggregateFunction(expr::AggregateFunction::new_udf(
+                udaf.into(),
+                vec![],
+                false,
+                None,
+                vec![],
+                None,
+            ));
+
+        let expected = aggregate_function_expr.clone();
+        let (actual, cycles) = simplify_with_cycle_count(aggregate_function_expr);
+        assert_eq!(actual, expected);
+        assert_eq!(cycles, 1);
     }
 
     /// A Mock UDAF which defines `simplify` to be used in tests
     /// related to UDAF simplification
     #[derive(Debug, Clone, PartialEq, Eq, Hash)]
     struct SimplifyMockUdaf {
-        simplify: bool,
+        simplify: Option<bool>,
     }
 
     impl SimplifyMockUdaf {
         /// make simplify method return new expression
         fn new_with_simplify() -> Self {
-            Self { simplify: true }
+            Self {
+                simplify: Some(true),
+            }
         }
         /// make simplify method return no change
         fn new_without_simplify() -> Self {
-            Self { simplify: false }
+            Self { simplify: None }
+        }
+        /// make simplify method return the original expression
+        fn new_with_noop() -> Self {
+            Self {
+                simplify: Some(false),
+            }
         }
     }
 
@@ -5680,10 +5710,12 @@ mod tests {
         }
 
         fn simplify(&self) -> Option<AggregateFunctionSimplification> {
-            if self.simplify {
-                Some(Box::new(|_, _| Ok(col("result_column"))))
-            } else {
-                None
+            match self.simplify {
+                Some(true) => Some(Box::new(|_, _| Ok(col("result_column")))),
+                Some(false) => Some(Box::new(|aggregate_function, _| {
+                    Ok(Expr::AggregateFunction(aggregate_function))
+                })),
+                None => None,
             }
         }
     }

datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs

@@ -396,7 +396,7 @@ mod tests {
             plan,
             @r"
         Projection: sum(test.a) + Int64(2) * CAST(count(test.a) AS Int64) AS sum(test.a + Int64(2)), sum(test.a) + Int64(3) * CAST(count(test.a) AS Int64) AS sum(test.a + Int64(3))
-          Aggregate: groupBy=[[]], aggr=[[sum(test.a), count(test.a)]]
+          Aggregate: groupBy=[[]], aggr=[[sum(test.a), count(Int64(1)) AS count(test.a)]]
             TableScan: test
         "
         )?;

datafusion/optimizer/src/single_distinct_to_groupby.rs

@@ -104,6 +104,32 @@ struct CountRollup {
     sum: Arc<AggregateUDF>,
 }
 
+fn unalias_top(mut expr: &Expr) -> &Expr {
+    while let Expr::Alias(alias) = expr
+        && alias
+            .metadata
+            .as_ref()
+            .is_none_or(|metadata| metadata.is_empty())
+    {
+        expr = &alias.expr;
+    }
+    expr
+}
+
+fn into_unaliased_top(expr: Expr) -> Expr {
+    match expr {
+        Expr::Alias(alias)
+            if alias
+                .metadata
+                .as_ref()
+                .is_none_or(|metadata| metadata.is_empty()) =>
+        {
+            into_unaliased_top(*alias.expr)
+        }
+        expr => expr,
+    }
+}
+
 impl CountRollup {
     fn try_new(config: &dyn OptimizerConfig) -> Option<Self> {
         let registry = config.function_registry()?;
@@ -131,6 +157,7 @@ fn is_single_distinct_agg(
     let mut distinct_aggs = vec![];
     let mut has_count_rollup = false;
     for expr in aggr_expr {
+        let expr = unalias_top(expr);
         if let Expr::AggregateFunction(AggregateFunction {
             func,
             params:
@@ -301,7 +328,7 @@ impl OptimizerRule for SingleDistinctToGroupBy {
                 // zero that `sum` reports as NULL over an empty input.
                 let (outer_aggr_exprs, outer_proj_exprs): (Vec<Expr>, Vec<Expr>) = aggr_expr
                     .into_iter()
-                    .map(|aggr_expr| match aggr_expr {
+                    .map(|aggr_expr| match into_unaliased_top(aggr_expr) {
                         Expr::AggregateFunction(AggregateFunction {
                             func,
                             params:
@@ -381,7 +408,7 @@ impl OptimizerRule for SingleDistinctToGroupBy {
                                 Ok((outer, proj))
                             }
                         }
-                        _ => Ok((aggr_expr.clone(), aggr_expr)),
+                        aggr_expr => Ok((aggr_expr.clone(), aggr_expr)),
                     })
                     .collect::<Result<Vec<_>>>()?
                     .into_iter()
@@ -1157,6 +1184,33 @@ mod tests {
         )
     }
 
+    #[test]
+    fn aliased_count_and_distinct_without_groupby() -> Result<()> {
+        let table_scan = test_table_scan_utf8_b()?;
+
+        // Simplifying `count(a)` to `count(1)` preserves its old name with an
+        // alias before this rule runs.
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                Vec::<Expr>::new(),
+                vec![
+                    count_udaf().call(vec![lit(1_i64)]).alias("count(test.a)"),
+                    count_distinct(col("b")),
+                ],
+            )?
+            .build()?;
+
+        assert_optimized_plan_equal!(
+            plan,
+            @r"
+        Projection: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(test.a), count(alias1) AS count(DISTINCT test.b) [count(test.a):Int64, count(DISTINCT test.b):Int64]
+          Aggregate: groupBy=[[]], aggr=[[sum(alias2), count(alias1)]] [sum(alias2):Int64;N, count(alias1):Int64]
+            Aggregate: groupBy=[[test.b AS alias1]], aggr=[[count(Int64(1)) AS alias2]] [alias1:Utf8, alias2:Int64]
+              TableScan: test [a:UInt32, b:Utf8, c:UInt32]
+        "
+        )
+    }
+
     #[test]
     fn count_star_min_max_sum_and_distinct_with_groupby() -> Result<()> {
         let table_scan = test_table_scan_utf8_b()?;

datafusion/sqllogictest/test_files/aggregate_memory_spill.slt

@@ -113,7 +113,7 @@ FROM (
 )
 ----
 <slt:ignore>
-04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
+04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(1) as count(alias1)], metrics=[<slt:ignore>spill_count=7,<slt:ignore>]
 <slt:ignore>
 
 # --- Case D: multiple aggregates (sum/min/max) under memory limit ---

datafusion/sqllogictest/test_files/group_by.slt

@@ -4532,16 +4532,16 @@ EXPLAIN SELECT c1, count(distinct c2), min(distinct c2), sum(c3), max(c4) FROM a
 logical_plan
 01)Sort: aggregate_test_100.c1 ASC NULLS LAST
 02)--Projection: aggregate_test_100.c1, count(alias1) AS count(DISTINCT aggregate_test_100.c2), min(alias1) AS min(DISTINCT aggregate_test_100.c2), sum(alias2) AS sum(aggregate_test_100.c3), max(alias3) AS max(aggregate_test_100.c4)
-03)----Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[count(alias1), min(alias1), sum(alias2), max(alias3)]]
+03)----Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[count(Int64(1)) AS count(alias1), min(alias1), sum(alias2), max(alias3)]]
 04)------Aggregate: groupBy=[[aggregate_test_100.c1, aggregate_test_100.c2 AS alias1]], aggr=[[sum(CAST(aggregate_test_100.c3 AS Int64)) AS alias2, max(aggregate_test_100.c4) AS alias3]]
 05)--------TableScan: aggregate_test_100 projection=[c1, c2, c3, c4]
 physical_plan
 01)SortPreservingMergeExec: [c1@0 ASC NULLS LAST]
 02)--ProjectionExec: expr=[c1@0 as c1, count(alias1)@1 as count(DISTINCT aggregate_test_100.c2), min(alias1)@2 as min(DISTINCT aggregate_test_100.c2), sum(alias2)@3 as sum(aggregate_test_100.c3), max(alias3)@4 as max(aggregate_test_100.c4)]
 03)----SortExec: expr=[c1@0 ASC NULLS LAST], preserve_partitioning=[true]
-04)------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[count(alias1), min(alias1), sum(alias2), max(alias3)]
+04)------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[count(1) as count(alias1), min(alias1), sum(alias2), max(alias3)]
 05)--------RepartitionExec: partitioning=Hash([c1@0], 8), input_partitions=8
-06)----------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[count(alias1), min(alias1), sum(alias2), max(alias3)]
+06)----------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[count(1) as count(alias1), min(alias1), sum(alias2), max(alias3)]
 07)------------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1, alias1@1 as alias1], aggr=[sum(aggregate_test_100.c3) as alias2, max(aggregate_test_100.c4) as alias3]
 08)--------------RepartitionExec: partitioning=Hash([c1@0, alias1@1], 8), input_partitions=8
 09)----------------AggregateExec: mode=Partial, gby=[c1@0 as c1, c2@1 as alias1], aggr=[sum(aggregate_test_100.c3) as alias2, max(aggregate_test_100.c4) as alias3]
@@ -5235,15 +5235,16 @@ GROUP BY
 ----
 logical_plan
 01)Projection: date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"),keywords_stream.ts,Utf8("2000-01-01")) AS ts_chunk, count(keywords_stream.keyword) AS alert_keyword_count
-02)--Aggregate: groupBy=[[date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"), keywords_stream.ts, TimestampNanosecond(946684800000000000, None)) AS date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"),keywords_stream.ts,Utf8("2000-01-01"))]], aggr=[[count(keywords_stream.keyword)]]
-03)----LeftSemi Join: keywords_stream.keyword = __correlated_sq_1.keyword
-04)------TableScan: keywords_stream projection=[ts, keyword]
-05)------SubqueryAlias: __correlated_sq_1
-06)--------TableScan: alert_keywords projection=[keyword]
+02)--Aggregate: groupBy=[[date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"), keywords_stream.ts, TimestampNanosecond(946684800000000000, None)) AS date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"),keywords_stream.ts,Utf8("2000-01-01"))]], aggr=[[count(Int64(1)) AS count(keywords_stream.keyword)]]
+03)----Projection: keywords_stream.ts
+04)------LeftSemi Join: keywords_stream.keyword = __correlated_sq_1.keyword
+05)--------TableScan: keywords_stream projection=[ts, keyword]
+06)--------SubqueryAlias: __correlated_sq_1
+07)----------TableScan: alert_keywords projection=[keyword]
 physical_plan
 01)ProjectionExec: expr=[date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"),keywords_stream.ts,Utf8("2000-01-01"))@0 as ts_chunk, count(keywords_stream.keyword)@1 as alert_keyword_count]
-02)--AggregateExec: mode=Single, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }, ts@0, 946684800000000000) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"),keywords_stream.ts,Utf8("2000-01-01"))], aggr=[count(keywords_stream.keyword)]
-03)----HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(keyword@0, keyword@1)]
+02)--AggregateExec: mode=Single, gby=[date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }, ts@0, 946684800000000000) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 120000000000 }"),keywords_stream.ts,Utf8("2000-01-01"))], aggr=[count(1) as count(keywords_stream.keyword)]
+03)----HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(keyword@0, keyword@1)], projection=[ts@0]
 04)------DataSourceExec: partitions=1, partition_sizes=[3]
 05)------DataSourceExec: partitions=1, partition_sizes=[3]