Addressed comments

Author: xiedeyantu

datafusion/optimizer/src/analyzer/resolve_grouping_function.rs

@@ -99,10 +99,17 @@ fn replace_grouping_exprs(
     {
         match expr {
             Expr::AggregateFunction(ref function) if is_grouping_function(&expr) => {
+                let grouping_id_type = is_grouping_set
+                    .then(|| {
+                        schema
+                            .field_with_name(None, Aggregate::INTERNAL_GROUPING_ID)
+                            .map(|f| f.data_type().clone())
+                    })
+                    .transpose()?;
                 let grouping_expr = grouping_function_on_id(
                     function,
                     &group_expr_to_bitmap_index,
-                    is_grouping_set,
+                    grouping_id_type,
                 )?;
                 projection_exprs.push(Expr::Alias(Alias::new(
                     grouping_expr,
@@ -184,41 +191,40 @@ fn validate_args(
 fn grouping_function_on_id(
     function: &AggregateFunction,
     group_by_expr: &HashMap<&Expr, usize>,
-    is_grouping_set: bool,
+    // None means not a grouping set (result is always 0).
+    grouping_id_type: Option<DataType>,
 ) -> Result<Expr> {
     validate_args(function, group_by_expr)?;
     let args = &function.params.args;
 
     // Postgres allows grouping function for group by without grouping sets, the result is then
     // always 0
-    if !is_grouping_set {
+    let Some(grouping_id_type) = grouping_id_type else {
         return Ok(Expr::Literal(ScalarValue::from(0i32), None));
-    }
+    };
 
-    let group_by_expr_count = group_by_expr.len();
-    let literal = |value: usize| {
-        if group_by_expr_count < 8 {
-            Expr::Literal(ScalarValue::from(value as u8), None)
-        } else if group_by_expr_count < 16 {
-            Expr::Literal(ScalarValue::from(value as u16), None)
-        } else if group_by_expr_count < 32 {
-            Expr::Literal(ScalarValue::from(value as u32), None)
-        } else {
-            Expr::Literal(ScalarValue::from(value as u64), None)
-        }
+    // Use the actual __grouping_id column type to size literals correctly. This
+    // accounts for duplicate-ordinal bits that `Aggregate::grouping_id_type`
+    // packs into the high bits of the column, which a simple count of grouping
+    // expressions would miss.
+    let literal = |value: usize| match &grouping_id_type {
+        DataType::UInt8 => Expr::Literal(ScalarValue::from(value as u8), None),
+        DataType::UInt16 => Expr::Literal(ScalarValue::from(value as u16), None),
+        DataType::UInt32 => Expr::Literal(ScalarValue::from(value as u32), None),
+        _ => Expr::Literal(ScalarValue::from(value as u64), None),
     };
     let grouping_id_column = Expr::Column(Column::from(Aggregate::INTERNAL_GROUPING_ID));
-    if args.len() == group_by_expr_count
+    if args.len() == group_by_expr.len()
         && args
             .iter()
             .rev()
             .enumerate()
             .all(|(idx, expr)| group_by_expr.get(expr) == Some(&idx))
     {
-        let n = group_by_expr_count;
+        let n = group_by_expr.len();
         // Mask the ordinal bits above position `n` so only the semantic bitmask is visible.
-        // (1 << n) - 1 masks the low n bits.
-        let semantic_mask: u64 = if n >= 64 { u64::MAX } else { (1u64 << n) - 1 };
+        // checked_shl returns None when n >= 64 (all bits are semantic), mapping to u64::MAX.
+        let semantic_mask: u64 = 1u64.checked_shl(n as u32).map_or(u64::MAX, |m| m - 1);
         let masked_id =
             bitwise_and(grouping_id_column.clone(), literal(semantic_mask as usize));
         return Ok(cast(masked_id, DataType::Int32));

datafusion/physical-plan/src/aggregates/mod.rs

@@ -2027,7 +2027,7 @@ pub fn evaluate_group_by(
     batch: &RecordBatch,
 ) -> Result<Vec<Vec<ArrayRef>>> {
     let max_ordinal = max_duplicate_ordinal(&group_by.groups);
-    let mut ordinal_per_pattern: HashMap<&Vec<bool>, usize> = HashMap::new();
+    let mut ordinal_per_pattern: HashMap<&[bool], usize> = HashMap::new();
     let exprs = evaluate_expressions_to_arrays(
         group_by.expr.iter().map(|(expr, _)| expr),
         batch,