Skip to content

fix: apply single_distinct_to_groupby to aliased distinct aggregates#23486

Open
gmhelmold wants to merge 1 commit into
apache:mainfrom
gmhelmold:fix/single-distinct-aliased-agg
Open

fix: apply single_distinct_to_groupby to aliased distinct aggregates#23486
gmhelmold wants to merge 1 commit into
apache:mainfrom
gmhelmold:fix/single-distinct-aliased-agg

Conversation

@gmhelmold

Copy link
Copy Markdown

Which issue does this PR close?

Rationale for this change

SingleDistinctToGroupBy rewrites a grouped AGG(DISTINCT col) into a cheaper nested double aggregate. Today the rule only recognizes a bare Expr::AggregateFunction in Aggregate.aggr_expr, so it fires for plans from the SQL planner but silently skips the logically-identical plan built through the DataFrame API — that path runs materially slower for the same query.

The divergence is purely in where each front-end places the output alias:

  • SQL (count(DISTINCT v) AS n): the planner hoists AS n into an outer Projection, leaving a bare Expr::AggregateFunction in aggr_expr.
  • DataFrame API (count(col("v")).distinct().alias("n")): .alias(..) wraps the aggregate directly, so aggr_expr holds Expr::Alias(AggregateFunction).

is_single_distinct_agg matched only the bare shape and returned false for Alias(AggregateFunction), so the rule never fired for DataFrame-API distinct aggregates.

What changes are included in this PR?

  • Unwrap a single Expr::Alias layer in the matcher (is_single_distinct_agg) and in the rewrite, mirroring the one-layer unalias() idiom already used in the optimizer (e.g. decorrelate.rs), so both front-end shapes are recognized identically.
  • The user-facing alias is preserved without extra work: the rewrite already rebuilds the final output Projection from the original Aggregate schema (schema.qualified_field(idx)), so AS n survives.
  • Bare-aggregate behavior is unchanged (byte-identical): every pre-existing snapshot passes untouched.

Are these changes tested?

Yes — new snapshot tests in single_distinct_to_groupby.rs:

  • single_distinct_aliased — the DataFrame-API Alias(AggregateFunction) shape is now rewritten, and the output column stays named n.
  • single_distinct_aliased_and_groupby — same, alongside a real GROUP BY column.
  • non_distinct_aliased and two_distinct_aliased_and_groupby — negative cases: an aliased non-distinct aggregate, and two aliased distinct aggregates on different fields, are still (correctly) left untouched.

Reverting the production change makes the two positive tests fail (the rule no longer fires) while every other test stays green, confirming the tests are wired to the fix.

Are there any user-facing changes?

No API changes. DataFrame-API queries using count(distinct …) (and other single-distinct aggregates) now receive the same optimized plan as their SQL equivalents — a performance improvement, with no change to results or output column names.

@github-actions github-actions Bot added the optimizer Optimizer rules label Jul 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

optimizer Optimizer rules

Projects

None yet

Development

Successfully merging this pull request may close these issues.

single_distinct_to_groupby not applied to DataFrame API count(DISTINCT)

1 participant