perf: Remove an aggregate whose GROUP BY covers a unique key

datafusion/expr/src/logical_plan/plan.rs

@@ -4896,8 +4896,15 @@ impl Unnest {
 
         let metadata = input_schema.metadata().clone();
         let df_schema = DFSchema::new_with_metadata(fields, metadata)?;
-        // We can use the existing functional dependencies:
-        let deps = input_schema.functional_dependencies().clone();
+        // Unnesting a list turns one input row into several, so a determinant
+        // that occurred once in the input can now occur many times. It still
+        // determines the same columns, so downgrade the dependency instead of
+        // dropping it. Unnesting a struct keeps one row per input row, and so
+        // keeps the dependencies as they are.
+        let mut deps = input_schema.functional_dependencies().clone();
+        if !list_columns.is_empty() {
+            deps = deps.with_dependency(Dependency::Multi);
+        }
         let schema = Arc::new(df_schema.with_functional_dependencies(deps)?);
 
         Ok(Unnest {

datafusion/optimizer/src/eliminate_group_by_constant.rs

@@ -15,20 +15,44 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! [`EliminateGroupByConstant`] removes constant and functionally redundant
-//! expressions from `GROUP BY` clause
+//! [`EliminateGroupByConstant`] simplifies a `GROUP BY` clause, and removes the
+//! aggregate altogether when the grouping makes it redundant.
 use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
 
 use std::collections::HashSet;
+use std::sync::Arc;
 
-use datafusion_common::Result;
+use arrow::datatypes::DataType;
 use datafusion_common::tree_node::Transformed;
-use datafusion_expr::{Aggregate, Expr, LogicalPlan, LogicalPlanBuilder, Volatility};
-
-/// Optimizer rule that removes constant expressions from `GROUP BY` clause
-/// and places additional projection on top of aggregation, to preserve
-/// original schema
+use datafusion_common::{DFSchema, Dependency, Result, ScalarValue};
+use datafusion_expr::expr::AggregateFunction;
+use datafusion_expr::{
+    Aggregate, Expr, ExprSchemable, LogicalPlan, LogicalPlanBuilder, Volatility, cast,
+    lit, when,
+};
+
+/// Optimizer rule that simplifies aggregation based on its `GROUP BY` clause.
+///
+/// Constant expressions are removed from the `GROUP BY`, with a projection on
+/// top to preserve the original schema.
+///
+/// When the `GROUP BY` covers a unique key of the input, every group holds
+/// exactly one row, so the aggregate produces one output row per input row like
+/// a projection, and each aggregate function returns that row's value:
+///
+/// ```text
+/// -- lineitem is keyed by (l_orderkey, l_linenumber)
+/// SELECT l_orderkey, l_linenumber, sum(l_quantity)
+/// FROM lineitem GROUP BY l_orderkey, l_linenumber
+///
+/// -- becomes
+/// SELECT l_orderkey, l_linenumber, CAST(l_quantity AS Decimal128(38, 2))
+/// FROM lineitem
+/// ```
+///
+/// That also removes a `SELECT DISTINCT` over a unique key, which the planner
+/// turns into an aggregate with no aggregate expressions.
 #[derive(Default, Debug)]
 pub struct EliminateGroupByConstant {}
 
@@ -50,6 +74,10 @@ impl OptimizerRule for EliminateGroupByConstant {
     ) -> Result<Transformed<LogicalPlan>> {
         match plan {
             LogicalPlan::Aggregate(aggregate) => {
+                if let Some(projection) = eliminate_aggregate(&aggregate)? {
+                    return Ok(Transformed::yes(projection));
+                }
+
                 // Collect bare column references in GROUP BY
                 let group_by_columns: HashSet<&datafusion_common::Column> = aggregate
                     .group_expr
@@ -103,6 +131,125 @@ impl OptimizerRule for EliminateGroupByConstant {
     }
 }
 
+/// Replaces an aggregate whose `GROUP BY` covers a unique key of its input with
+/// a projection, or returns `None` when that cannot be done.
+fn eliminate_aggregate(aggregate: &Aggregate) -> Result<Option<LogicalPlan>> {
+    // An empty GROUP BY produces one row even for an empty input, which a
+    // projection would not. Grouping sets do not group by every expression at
+    // once, so a key among them proves nothing.
+    if aggregate.group_expr.is_empty()
+        || aggregate
+            .group_expr
+            .iter()
+            .any(|expr| matches!(expr, Expr::GroupingSet(_)))
+        || !group_by_covers_unique_key(&aggregate.input, &aggregate.group_expr)
+    {
+        return Ok(None);
+    }
+
+    // The output schema has to survive untouched, so each aggregate is replaced
+    // by an expression of the same type aliased to the same name.
+    let output_fields = aggregate.schema.fields();
+    let mut projection = aggregate.group_expr.clone();
+    for (index, aggr_expr) in aggregate.aggr_expr.iter().enumerate() {
+        let field = &output_fields[aggregate.group_expr.len() + index];
+        let Some(value) =
+            single_row_value(aggr_expr, field.data_type(), aggregate.input.schema())?
+        else {
+            return Ok(None);
+        };
+        projection.push(value.alias(field.name()));
+    }
+
+    let input = Arc::clone(&aggregate.input);
+    Ok(Some(
+        LogicalPlanBuilder::from(Arc::unwrap_or_clone(input))
+            .project(projection)?
+            .build()?,
+    ))
+}
+
+/// Returns true when the GROUP BY expressions include a unique key of the
+/// input, so that every group holds exactly one row.
+///
+/// Extra grouping expressions beyond the key are harmless: they can only split
+/// groups further, and the groups are already singletons.
+fn group_by_covers_unique_key(input: &LogicalPlan, group_expr: &[Expr]) -> bool {
+    let schema = input.schema();
+    let grouped: HashSet<usize> = group_expr
+        .iter()
+        .filter_map(|expr| match expr {
+            Expr::Alias(alias) => alias.expr.as_ref().try_as_col(),
+            _ => expr.try_as_col(),
+        })
+        .filter_map(|column| schema.maybe_index_of_column(column))
+        .collect();
+    if grouped.is_empty() {
+        return false;
+    }
+
+    schema.functional_dependencies().iter().any(|dependency| {
+        // A nullable key does not identify rows: two rows can both be NULL and
+        // are then not distinguished by the grouping either.
+        let nullable = dependency.nullable
+            && dependency
+                .source_indices
+                .iter()
+                .any(|&index| schema.field(index).is_nullable());
+        !nullable
+            && dependency.mode == Dependency::Single
+            // The dependency has to determine the whole row, not just part of it.
+            && dependency.target_indices.len() == schema.fields().len()
+            && dependency
+                .source_indices
+                .iter()
+                .all(|index| grouped.contains(index))
+    })
+}
+
+/// The value an aggregate takes over a group of exactly one row, as an
+/// expression over that row, or `None` when it cannot be expressed.
+fn single_row_value(
+    aggr_expr: &Expr,
+    output_type: &DataType,
+    input_schema: &DFSchema,
+) -> Result<Option<Expr>> {
+    let Expr::AggregateFunction(AggregateFunction { func, params }) = aggr_expr else {
+        return Ok(None);
+    };
+    // A FILTER can exclude the single row, leaving the aggregate with no input
+    // at all, which is a different value for every function.
+    if params.filter.is_some() {
+        return Ok(None);
+    }
+    // DISTINCT, ORDER BY and IGNORE NULLS all make no difference to a group of
+    // one row: there is nothing to deduplicate or order, and a single NULL is
+    // skipped by the functions below anyway.
+    let [arg] = params.args.as_slice() else {
+        return Ok(None);
+    };
+
+    Ok(match func.name() {
+        // Over one row these all return that row's value, in the type the
+        // aggregate would have returned.
+        "min" | "max" | "sum" | "avg" | "first_value" | "last_value" => {
+            Some(if &arg.get_type(input_schema)? == output_type {
+                arg.clone()
+            } else {
+                cast(arg.clone(), output_type.clone())
+            })
+        }
+        // COUNT ignores NULLs, so it is 1 unless the single value is NULL.
+        "count" => Some(if arg.nullable(input_schema)? {
+            when(arg.clone().is_null(), lit(ScalarValue::Int64(Some(0))))
+                .otherwise(lit(ScalarValue::Int64(Some(1))))?
+        } else {
+            lit(ScalarValue::Int64(Some(1)))
+        }),
+        _ => None,
+    })
+}
+
 /// Checks if a GROUP BY expression is redundant (can be removed without
 /// changing grouping semantics). An expression is redundant if it is a
 /// deterministic function of constants and columns already present as bare

datafusion/sqllogictest/test_files/aggregate.slt

@@ -7221,17 +7221,15 @@ EXPLAIN SELECT DISTINCT c3, min(c1) FROM aggregate_test_100 group by c3 limit 5;
 ----
 logical_plan
 01)Limit: skip=0, fetch=5
-02)--Aggregate: groupBy=[[aggregate_test_100.c3, min(aggregate_test_100.c1)]], aggr=[[]]
-03)----Aggregate: groupBy=[[aggregate_test_100.c3]], aggr=[[min(aggregate_test_100.c1)]]
-04)------TableScan: aggregate_test_100 projection=[c1, c3]
+02)--Aggregate: groupBy=[[aggregate_test_100.c3]], aggr=[[min(aggregate_test_100.c1)]]
+03)----TableScan: aggregate_test_100 projection=[c1, c3]
 physical_plan
 01)CoalescePartitionsExec: fetch=5
-02)--AggregateExec: mode=SinglePartitioned, gby=[c3@0 as c3, min(aggregate_test_100.c1)@1 as min(aggregate_test_100.c1)], aggr=[], lim=[5]
-03)----AggregateExec: mode=FinalPartitioned, gby=[c3@0 as c3], aggr=[min(aggregate_test_100.c1)]
-04)------RepartitionExec: partitioning=Hash([c3@0], 4), input_partitions=4
-05)--------AggregateExec: mode=Partial, gby=[c3@1 as c3], aggr=[min(aggregate_test_100.c1)]
-06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100_with_dates.csv]]}, projection=[c1, c3], file_type=csv, has_header=true
+02)--AggregateExec: mode=FinalPartitioned, gby=[c3@0 as c3], aggr=[min(aggregate_test_100.c1)]
+03)----RepartitionExec: partitioning=Hash([c3@0], 4), input_partitions=4
+04)------AggregateExec: mode=Partial, gby=[c3@1 as c3], aggr=[min(aggregate_test_100.c1)]
+05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100_with_dates.csv]]}, projection=[c1, c3], file_type=csv, has_header=true
 
 
 #
@@ -7348,23 +7346,18 @@ query TT
 EXPLAIN SELECT DISTINCT c3, c2 FROM aggregate_test_100 group by c2, c3 limit 3 offset 10;
 ----
 logical_plan
-01)Limit: skip=10, fetch=3
-02)--Aggregate: groupBy=[[aggregate_test_100.c3, aggregate_test_100.c2]], aggr=[[]]
-03)----Projection: aggregate_test_100.c3, aggregate_test_100.c2
-04)------Aggregate: groupBy=[[aggregate_test_100.c2, aggregate_test_100.c3]], aggr=[[]]
-05)--------TableScan: aggregate_test_100 projection=[c2, c3]
+01)Projection: aggregate_test_100.c3, aggregate_test_100.c2
+02)--Limit: skip=10, fetch=3
+03)----Aggregate: groupBy=[[aggregate_test_100.c2, aggregate_test_100.c3]], aggr=[[]]
+04)------TableScan: aggregate_test_100 projection=[c2, c3]
 physical_plan
-01)GlobalLimitExec: skip=10, fetch=3
-02)--AggregateExec: mode=Final, gby=[c3@0 as c3, c2@1 as c2], aggr=[], lim=[13]
-03)----CoalescePartitionsExec
-04)------AggregateExec: mode=Partial, gby=[c3@0 as c3, c2@1 as c2], aggr=[], lim=[13]
-05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-06)----------ProjectionExec: expr=[c3@1 as c3, c2@0 as c2]
-07)------------AggregateExec: mode=Final, gby=[c2@0 as c2, c3@1 as c3], aggr=[]
-08)--------------CoalescePartitionsExec
-09)----------------AggregateExec: mode=Partial, gby=[c2@0 as c2, c3@1 as c3], aggr=[]
-10)------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
-11)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100_with_dates.csv]]}, projection=[c2, c3], file_type=csv, has_header=true
+01)ProjectionExec: expr=[c3@1 as c3, c2@0 as c2]
+02)--GlobalLimitExec: skip=10, fetch=3
+03)----AggregateExec: mode=Final, gby=[c2@0 as c2, c3@1 as c3], aggr=[], lim=[13]
+04)------CoalescePartitionsExec
+05)--------AggregateExec: mode=Partial, gby=[c2@0 as c2, c3@1 as c3], aggr=[], lim=[13]
+06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100_with_dates.csv]]}, projection=[c2, c3], file_type=csv, has_header=true
 
 query II
 SELECT DISTINCT c3, c2 FROM aggregate_test_100 group by c3, c2 order by c3, c2 limit 3 offset 10;

datafusion/sqllogictest/test_files/distinct_on.slt

@@ -397,7 +397,7 @@ FROM aggregate_test_100 GROUP BY c1 ORDER BY c1, agg2;
 logical_plan
 01)Projection: first_value(aggregate_test_100.c1) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST] AS c1, first_value(agg2) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST] AS agg2
 02)--Sort: aggregate_test_100.c1 ASC NULLS LAST
-03)----Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[first_value(aggregate_test_100.c1) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST], first_value(max(aggregate_test_100.c4) AS agg2) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST]]]
+03)----Projection: aggregate_test_100.c1, aggregate_test_100.c1 AS first_value(aggregate_test_100.c1) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST], max(aggregate_test_100.c4) AS agg2 AS first_value(agg2) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST]
 04)------Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[max(aggregate_test_100.c4)]]
 05)--------TableScan: aggregate_test_100 projection=[c1, c4]