kosiew commented on code in PR #24859:
URL: https://github.com/apache/datafusion/pull/24859#discussion_r3955646518
##########
datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:
##########
@@ -333,8 +333,13 @@ async fn simple_aggregate() -> Result<()> {
#[tokio::test]
async fn aggregate_distinct_with_having() -> Result<()> {
- roundtrip("SELECT a, count(distinct b) FROM data GROUP BY a, c HAVING
count(b) > 100")
- .await
+ let ctx = create_context_without_single_distinct_to_group_by().await?;
Review Comment:
I think using the custom context here makes sense. Substrait does not
represent the aggregate/group-expression aliases introduced by this rewrite, so
this looks like an existing plan-identity limitation rather than a semantic
regression.
Would it be worth adding a separate test using the default context that
exercises the rewritten plan through Substrait and checks its schema and query
results instead of exact plan identity? This test could then stay focused on
identity roundtripping with `single_distinct_aggregation_to_group_by` disabled.
##########
datafusion/optimizer/src/single_distinct_to_groupby.rs:
##########
@@ -96,7 +170,58 @@ fn is_single_distinct_agg(aggr_expr: &[Expr]) ->
Result<bool> {
return Ok(false);
}
}
- Ok(aggregate_count == aggr_expr.len() && fields_set.len() == 1)
+ if aggregate_count != aggr_expr.len() || fields_set.len() != 1 {
+ return Ok(false);
+ }
+ if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs,
input_schema)? {
+ return Ok(false);
+ }
+ Ok(true)
+}
+
+/// Whether the rewrite is worth extending to a plan that only qualifies
because
+/// of the non-distinct `count`.
+///
+/// The rewrite is not free: every other aggregate moves down to the inner
group
+/// by, which has a row per `(group, distinct value)` pair rather than per
group,
+/// and each one keeps its state at that finer grain. What pays for it is
taking
+/// the distinct aggregate off `GroupsAccumulatorAdapter`, whose one boxed
+/// accumulator per group is the expensive shape. A distinct aggregate that
+/// already has a specialized `GroupsAccumulator` never went near the adapter,
so
+/// there is nothing to buy and only the inner group by to pay for: ClickBench
+/// Q22, whose `count(DISTINCT "UserID")` is over an `Int64`, measured a 132%
+/// increase in peak memory pool reservation when the rewrite applied to it.
+///
+/// `Some(false)` is the only answer that buys anything. `Some(true)` says the
+/// call never reaches the adapter. `None` says the function does not answer
the
+/// question from the argument types, which is the default and so the answer
for
+/// almost every function. Silence is not evidence, and reading it as
`Some(false)`
+/// would open this path to every such function: `sum(DISTINCT int_col)`
beside a
+/// `count(*)` measured 3.15x the peak memory once rewritten, over 4,000,000
rows
+/// in 2,000 groups.
+///
+/// The predicate is a proxy, not the true discriminator. What decides the
+/// outcome is the cost per distinct value on each side, which this rule cannot
+/// see. The proxy is deliberately conservative in the direction that leaves a
+/// plan alone.
+///
+/// The existing tolerance of a non-distinct `sum`, `min` or `max` predates
this
+/// and is left alone: narrowing it would change plans that have always been
+/// rewritten, which no measurement here calls for.
+fn rewrite_pays_for_count(
+ distinct_aggs: &[(&Arc<AggregateUDF>, &Vec<Expr>)],
Review Comment:
Small cleanup: I don't think this needs to require `Vec<Expr>`. Could this
take slices instead, for example `&[(&Arc<AggregateUDF>, &[Expr])]`, and use
`args.as_slice()` when collecting? That seems a little more idiomatic and
matches the surrounding slice-based APIs. No behavior change intended.
--
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]
