neilconway commented on code in PR #25803:
URL: https://github.com/apache/datafusion/pull/25803#discussion_r4116040038
##########
datafusion/sql/tests/cases/plan_to_sql.rs:
##########
@@ -4611,6 +4623,99 @@ fn
roundtrip_approx_percentile_cont_within_group_with_centroids()
Ok(())
}
+/// Ordering of an ordered aggregate that uses the argument-list syntax
+/// (`array_agg(x ORDER BY y)`) must survive unparsing, not just the
+/// `WITHIN GROUP` spelling. See #25796.
Review Comment:
"Not just the `WITHIN GROUP` spelling" seems confusing to future readers
because it implies they know about the bug we're fixing here. The future state
is just that there are two syntax variants and they both need tests.
##########
datafusion/sql/src/unparser/expr.rs:
##########
@@ -408,18 +408,30 @@ impl Unparser<'_> {
..
} = &agg.params;
- // if this is a WITHIN GROUP aggregate, skip the prepended arg
- let (args_to_use, within_group) =
- if agg.func.supports_within_group_clause() &&
!order_by.is_empty() {
- let args_to_use =
self.function_args_to_sql(&args[1..])?;
- let within_group = order_by
- .iter()
- .map(|sort_expr| self.sort_to_sql(sort_expr))
- .collect::<Result<Vec<ast::OrderByExpr>>>()?;
- (args_to_use, within_group)
- } else {
- (self.function_args_to_sql(args)?, Vec::new())
- };
+ // An aggregate's `order_by` is spelled one of two ways in SQL:
+ // `WITHIN GROUP (ORDER BY ..)` for ordered-set aggregates,
and an
+ // `ORDER BY` clause inside the argument list for every other
+ // aggregate. Both map onto `params.order_by`, so pick the
spelling
+ // the function asks for rather than dropping the ordering.
Review Comment:
"rather than dropping the ordering" is a bit of a strange comment to make
(why would we do that?), I would omit it.
##########
datafusion/sql/src/unparser/expr.rs:
##########
@@ -408,18 +408,30 @@ impl Unparser<'_> {
..
} = &agg.params;
- // if this is a WITHIN GROUP aggregate, skip the prepended arg
- let (args_to_use, within_group) =
- if agg.func.supports_within_group_clause() &&
!order_by.is_empty() {
- let args_to_use =
self.function_args_to_sql(&args[1..])?;
- let within_group = order_by
- .iter()
- .map(|sort_expr| self.sort_to_sql(sort_expr))
- .collect::<Result<Vec<ast::OrderByExpr>>>()?;
- (args_to_use, within_group)
- } else {
- (self.function_args_to_sql(args)?, Vec::new())
- };
+ // An aggregate's `order_by` is spelled one of two ways in SQL:
+ // `WITHIN GROUP (ORDER BY ..)` for ordered-set aggregates,
and an
+ // `ORDER BY` clause inside the argument list for every other
+ // aggregate. Both map onto `params.order_by`, so pick the
spelling
+ // the function asks for rather than dropping the ordering.
+ let use_within_group =
+ agg.func.supports_within_group_clause() &&
!order_by.is_empty();
+ let order_by_sql = order_by
+ .iter()
+ .map(|sort_expr| self.sort_to_sql(sort_expr))
+ .collect::<Result<Vec<ast::OrderByExpr>>>()?;
+
+ let (args_to_use, within_group, clauses) = if use_within_group
{
+ // if this is a WITHIN GROUP aggregate, skip the prepended
arg
+ (self.function_args_to_sql(&args[1..])?, order_by_sql,
vec![])
+ } else if order_by.is_empty() {
+ (self.function_args_to_sql(args)?, vec![], vec![])
+ } else {
+ (
+ self.function_args_to_sql(args)?,
+ vec![],
+
vec![ast::FunctionArgumentClause::OrderBy(order_by_sql)],
+ )
+ };
Review Comment:
Wouldn't this be a bit cleaner if we branched on `order_by.is_empty()` up
top? And then handled the two variants of when `order_by` is non-empty nested
beneath that.
##########
datafusion/sql/tests/cases/plan_to_sql.rs:
##########
@@ -4611,6 +4623,99 @@ fn
roundtrip_approx_percentile_cont_within_group_with_centroids()
Ok(())
}
+/// Ordering of an ordered aggregate that uses the argument-list syntax
+/// (`array_agg(x ORDER BY y)`) must survive unparsing, not just the
+/// `WITHIN GROUP` spelling. See #25796.
+#[test]
+fn roundtrip_ordered_aggregate_order_by() -> Result<(), DataFusionError> {
+ roundtrip_statement_with_dialect_helper!(
+ sql: "SELECT first_name, last_value(age ORDER BY salary) FROM person
GROUP BY first_name",
+ parser_dialect: GenericDialect {},
+ unparser_dialect: UnparserDefaultDialect {},
+ expected: @"SELECT person.first_name, last_value(person.age ORDER BY
person.salary ASC NULLS LAST) FROM person GROUP BY person.first_name",
+ );
+ roundtrip_statement_with_dialect_helper!(
+ sql: "SELECT first_name, first_value(age ORDER BY salary DESC) FROM
person GROUP BY first_name",
+ parser_dialect: GenericDialect {},
+ unparser_dialect: UnparserDefaultDialect {},
+ expected: @"SELECT person.first_name, first_value(person.age ORDER BY
person.salary DESC NULLS FIRST) FROM person GROUP BY person.first_name",
+ );
+ roundtrip_statement_with_dialect_helper!(
+ sql: "SELECT first_name, array_agg(age ORDER BY salary) FROM person
GROUP BY first_name",
+ parser_dialect: GenericDialect {},
+ unparser_dialect: UnparserDefaultDialect {},
+ expected: @"SELECT person.first_name, array_agg(person.age ORDER BY
person.salary ASC NULLS LAST) FROM person GROUP BY person.first_name",
+ );
+ roundtrip_statement_with_dialect_helper!(
+ sql: "SELECT first_name, string_agg(CAST(age AS VARCHAR), ',' ORDER BY
salary) FROM person GROUP BY first_name",
+ parser_dialect: GenericDialect {},
+ unparser_dialect: UnparserDefaultDialect {},
+ expected: @"SELECT person.first_name, string_agg(CAST(person.age AS
VARCHAR), ',' ORDER BY person.salary ASC NULLS LAST) FROM person GROUP BY
person.first_name",
+ );
+ roundtrip_statement_with_dialect_helper!(
+ sql: "SELECT first_name, array_agg(DISTINCT age ORDER BY salary, id)
FROM person GROUP BY first_name",
+ parser_dialect: GenericDialect {},
+ unparser_dialect: UnparserDefaultDialect {},
+ expected: @"SELECT person.first_name, array_agg(DISTINCT person.age
ORDER BY person.salary ASC NULLS LAST, person.id ASC NULLS LAST) FROM person
GROUP BY person.first_name",
+ );
+ Ok(())
+}
+
+/// The emitted SQL must not merely *look* right: re-planning it has to produce
+/// the same plan, so a frozen-but-broken snapshot cannot hide an ordering that
+/// was dropped or moved to a clause the parser reads differently. Covers both
+/// spellings, since the unparser picks between them. See #25796.
+#[test]
+fn ordered_aggregate_order_by_survives_replanning() -> Result<(),
DataFusionError> {
+ let state = MockSessionState::default()
+ .with_aggregate_function(
+ datafusion_functions_aggregate::array_agg::array_agg_udaf(),
+ )
+ .with_aggregate_function(
+ datafusion_functions_aggregate::first_last::first_value_udaf(),
+ )
+ .with_aggregate_function(
+ datafusion_functions_aggregate::first_last::last_value_udaf(),
+ )
+ .with_aggregate_function(
+ datafusion_functions_aggregate::string_agg::string_agg_udaf(),
+ )
+ .with_aggregate_function(
+
datafusion_functions_aggregate::percentile_cont::percentile_cont_udaf(),
+ )
+ .with_expr_planner(Arc::new(CoreFunctionPlanner::default()));
+ let context = MockContextProvider { state };
+ let sql_to_rel = SqlToRel::new(&context);
+ let unparser = Unparser::default();
+
+ for sql in [
+ "SELECT first_name, last_value(age ORDER BY salary) FROM person GROUP
BY first_name",
+ "SELECT first_name, first_value(age ORDER BY salary DESC) FROM person
GROUP BY first_name",
+ "SELECT first_name, array_agg(age ORDER BY salary) FROM person GROUP
BY first_name",
+ "SELECT first_name, array_agg(DISTINCT age ORDER BY salary, id) FROM
person GROUP BY first_name",
+ "SELECT first_name, string_agg(CAST(age AS VARCHAR), ',' ORDER BY
salary) FROM person GROUP BY first_name",
+ "SELECT first_name, percentile_cont(0.5) WITHIN GROUP (ORDER BY age)
FROM person GROUP BY first_name",
+ ] {
+ let plan = sql_to_rel.sql_statement_to_plan(
+ Parser::new(&GenericDialect {})
+ .try_with_sql(sql)?
+ .parse_statement()?,
+ )?;
+ let unparsed = unparser.plan_to_sql(&plan)?.to_string();
+ let replanned = sql_to_rel.sql_statement_to_plan(
+ Parser::new(&GenericDialect {})
+ .try_with_sql(&unparsed)?
+ .parse_statement()?,
+ )?;
+ assert_eq!(
+ plan.display_indent().to_string(),
+ replanned.display_indent().to_string(),
+ "unparsing changed the plan for `{sql}`, emitting `{unparsed}`"
+ );
+ }
+ Ok(())
+}
Review Comment:
I'd delete this test and add the SQL in question to the existing
`roundtrip_statement` test.
--
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]
