Cherry-Pick: Preserve grouping ID during aggregate CSE (#24144)

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -826,6 +826,9 @@ fn extract_expressions(expr: &Expr, result: &mut Vec<Expr>) {
             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(())
     }