xumingming opened a new pull request, #59402:
URL: https://github.com/apache/spark/pull/59402
### What changes were proposed in this pull request?
`RewriteDistinctAggregates` canonicalizes conditional distinct counts into
FILTER form:
```sql
COUNT(DISTINCT IF(cond, base, NULL)) -> COUNT(DISTINCT base)
FILTER (WHERE cond)
COUNT(DISTINCT CASE WHEN cond THEN base END) -> COUNT(DISTINCT base)
FILTER (WHERE cond)
```
so that multiple conditional counts over the same base column collapse from N
distinct groups into one, reducing the Expand fan-out from N x to 1 x. The
union of the member filters then guards the group's shared distinct key
(filters deduplicated by canonicalized form):
```sql
SELECT key,
COUNT(DISTINCT IF(cond1, base, NULL)),
COUNT(DISTINCT IF(cond2, base, NULL))
-- is collapsed into one distinct group whose Expand key is
IF(Or(cond1, cond2), base, NULL)
```
Rows matching no filter project NULL as the key and merge into a single NULL
key per group value, keeping the first (map-side partial) aggregate's output
cardinality bounded by the pre-rewrite plan:
```
distinct(base | Or(conds)) + 1 <= sum_i distinct(base | cond_i) + N
```
The guard applies only when it is safe and never harmful:
- every aggregate in the group is filtered — an unfiltered one needs the raw
base over all rows;
- all filters are deterministic — the guard evaluates each condition a second
time, beside its flag column;
- the base expression must be safe to evaluate unconditionally (plain
columns,
comparisons, Coalesce, struct/map accessors; expressions that can throw,
e.g. Cast under ANSI mode, are excluded), since the guard evaluates the
base
on rows matched by *any* member filter.
A single filtered distinct count (user-written `FILTER` syntax) gets the
guard
as well, bounding its key cardinality the same way. The rewrite is gated by
the
internal config `spark.sql.optimizer.rewriteCountDistinctConditional.enabled`
(default true).
### Why are the changes needed?
N conditional counts over the same column otherwise create N distinct groups,
and the Expand amplifies every input row N times into the partial aggregates.
Collapsing them into one group removes the amplification, and the
union-of-filters guard ensures the collapsed plan is never worse than the
uncollapsed one, regardless of predicate selectivity.
Benchmark (ExpandBenchmark, local), covering both predicate regimes:
- overlapping day windows (9 conditional COUNT(DISTINCT) + 9 conditional SUM,
nested windows covering all rows, 5M rows):
conf off (9 distinct groups): 50,561 ms/op
conf on (1 collapsed group): 10,118 ms/op (5.0x faster)
- disjoint ~1%-selective predicates (10 conditional COUNT(DISTINCT) over a
unique base column, 10M rows):
conf off (10 distinct groups): 15,670 ms/op
conf on (1 guarded group): 1,531 ms/op (10.2x faster)
### Does this PR introduce _any_ user-facing change?
No. Query results are unchanged; this is an internal optimizer rewrite.
`EXPLAIN` plans for conditional distinct counts show the new shape.
### How was this patch tested?
- `RewriteDistinctAggregatesSuite`: guard-shape assertions (Expand key is
`IF(Or(conds), base, NULL)` with the Or leaves set-equal to the member
filters) plus negative controls — an unfiltered member keeps the raw key, a
nondeterministic filter disables the guard.
- `RewriteDistinctAggregatesConditionalQuerySuite` (conf on and off):
high-cardinality base with selective disjoint predicates; overlapping
windows mixed with disjoint flags; NULL conditions; mixed
filtered/unfiltered counts.
- The SPARK-58501 codegen test updated for the new plan shape: a shared
condition is evaluated twice per Expand branch without subexpression
elimination, still once per row with it.
- `ExpandBenchmark` gained two workload sections, one per predicate regime:
overlapping day windows (covering all rows) and disjoint ~1%-selective
predicates over a unique base column, each comparing the conf off and on
plans; the benchmark will be run via the GitHub Actions benchmark workflow
for a consistent environment.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Pi
--
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]
