diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs index 2775d62144c56..41d09db7c2bbe 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -826,6 +826,9 @@ fn extract_expressions(expr: &Expr, result: &mut Vec) { let col = Column::new(qualifier, field_name); result.push(Expr::Column(col)) } + result.push(Expr::Column(Column::from_name( + Aggregate::INTERNAL_GROUPING_ID, + ))); } else { let (qualifier, field_name) = expr.qualified_name(); let col = Column::new(qualifier, field_name); @@ -1106,6 +1109,27 @@ mod test { ) } + #[test] + fn common_aggregate_grouping_set_preserves_internal_id() -> Result<()> { + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .aggregate( + vec![grouping_set(vec![vec![col("a")]])], + vec![avg(col("b")).alias("first"), avg(col("b")).alias("second")], + )? + .filter(col(Aggregate::INTERNAL_GROUPING_ID).eq(lit(0_u8)))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @ r" + Filter: __grouping_id = UInt8(0) + Projection: test.a, __grouping_id, __common_expr_1 AS first, __common_expr_1 AS second + Aggregate: groupBy=[[GROUPING SETS ((test.a))]], aggr=[[avg(test.b) AS __common_expr_1]] + TableScan: test + " + ) + } + #[test] fn subexpr_in_same_order() -> Result<()> { let table_scan = test_table_scan()?; @@ -1288,20 +1312,31 @@ mod test { #[test] fn test_extract_expressions_from_grouping_set() -> Result<()> { - let mut result = Vec::with_capacity(3); + let mut result = Vec::with_capacity(4); let grouping = grouping_set(vec![vec![col("a"), col("b")], vec![col("c")]]); extract_expressions(&grouping, &mut result); - assert!(result.len() == 3); + assert_eq!( + result, + vec![ + col("a"), + col("b"), + col("c"), + col(Aggregate::INTERNAL_GROUPING_ID), + ] + ); Ok(()) } #[test] fn test_extract_expressions_from_grouping_set_with_identical_expr() -> Result<()> { - let mut result = Vec::with_capacity(2); + let mut result = Vec::with_capacity(3); let grouping = grouping_set(vec![vec![col("a"), col("b")], vec![col("a")]]); extract_expressions(&grouping, &mut result); - assert!(result.len() == 2); + assert_eq!( + result, + vec![col("a"), col("b"), col(Aggregate::INTERNAL_GROUPING_ID),] + ); Ok(()) }