xiedeyantu commented on code in PR #21058:
URL: https://github.com/apache/datafusion/pull/21058#discussion_r3051594558
##########
datafusion/optimizer/src/analyzer/resolve_grouping_function.rs:
##########
@@ -184,40 +191,43 @@ 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),
Review Comment:
done!
--
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]
