Copilot commented on code in PR #26356:
URL: https://github.com/apache/datafusion/pull/26356#discussion_r4299039619


##########
datafusion/optimizer/src/eliminate_group_by_constant.rs:
##########
@@ -56,33 +57,86 @@ impl OptimizerRule for EliminateGroupByConstant {
                     })
                     .collect();
 
-                let (redundant, required): (Vec<_>, Vec<_>) = aggregate
+                let is_redundant: Vec<bool> = aggregate
                     .group_expr
                     .iter()
-                    .partition(|expr| is_redundant_group_expr(expr, 
&group_by_columns));
+                    .map(|expr| is_redundant_group_expr(expr, 
&group_by_columns))
+                    .collect();
+                let redundant_count = is_redundant.iter().filter(|r| 
**r).count();
                 // Return now if no simplification can be done. We also bail 
out
                 // if applying the optimization would eliminate all of the
                 // grouping expressions (e.g., GROUP BY on only constant
                 // expressions): this would turn a grouped aggregate into an
                 // ungrouped aggregate, which changes query semantics (grouped
                 // aggregates produce an empty result set on an empty input,
                 // whereas ungrouped aggregates return a single row).
-                if redundant.is_empty() || required.is_empty() {
+                // A grouping set adds output fields that do not map one to one
+                // to `group_expr`, so the rule also leaves it unchanged.
+                if redundant_count == 0
+                    || redundant_count == aggregate.group_expr.len()
+                    || aggregate
+                        .group_expr
+                        .iter()
+                        .any(|expr| matches!(expr, Expr::GroupingSet(_)))
+                {
                     return 
Ok(Transformed::no(LogicalPlan::Aggregate(aggregate)));
                 }
 
-                let simplified_aggregate = 
LogicalPlan::Aggregate(Aggregate::try_new(
-                    aggregate.input,
-                    required.into_iter().cloned().collect(),
-                    aggregate.aggr_expr.clone(),
-                )?);
-
-                let projection_expr =
-                    
aggregate.group_expr.into_iter().chain(aggregate.aggr_expr);
+                let original_schema = Arc::clone(&aggregate.schema);
+                let required = aggregate
+                    .group_expr
+                    .iter()
+                    .zip(&is_redundant)
+                    .filter(|(_, redundant)| !**redundant)
+                    .map(|(expr, _)| expr.clone())
+                    .collect();
+                let simplified_aggregate =
+                    Aggregate::try_new(aggregate.input, required, 
aggregate.aggr_expr)?;
+                let simplified_schema = 
Arc::clone(&simplified_aggregate.schema);
+
+                // Refer to the output columns of the simplified aggregate by
+                // position. Do not re-plan the original group expressions on 
top
+                // of it: an expression can refer to an input column through a
+                // name that the aggregate output does not carry, so a rebuilt
+                // projection does not always resolve. A redundant expression
+                // refers only to bare GROUP BY columns, which the simplified
+                // aggregate outputs unchanged.
+                let mut projection_expr =
+                    Vec::with_capacity(original_schema.fields().len());
+                let mut next_output = 0;
+                for (i, (expr, redundant)) in aggregate
+                    .group_expr
+                    .into_iter()
+                    .zip(is_redundant)
+                    .enumerate()
+                {
+                    let (qualifier, field) = 
original_schema.qualified_field(i);
+                    let expr = if redundant {
+                        expr.unalias()

Review Comment:
   `unalias()` discards metadata attached to a redundant group-key alias, and 
`alias_qualified()` below restores only its name and qualifier. For example, 
grouping by `col("a")` and `lit(1u32).alias_with_metadata("k", Some(metadata))` 
now loses `k`'s field metadata in the output schema. Keep the original 
expression here to preserve the alias and its metadata, and add a regression 
assertion for the output field metadata.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to