[SPARK-58387][SQL] Do not pre-aggregate under an Expand when a duplicate-sensitive aggregate is present

[Vivek1106-04]

What changes were proposed in this pull request?

OptimizeExpand (added by SPARK-56315, gated by the internal conf spark.sql.optimizer.optimizeExpandRatio, default -1 = disabled) inserts a de-duplicating Aggregate beneath the Expand produced by RewriteDistinctAggregates. That de-duplication is sound only for pure distinct aggregates, and the conf's own doc states the precondition: "Only applies to pure distinct aggregates without non-distinct aggregates or FILTER clauses."

The guard implementing that precondition tested attributes rather than aggregate functions:

val innerGroupByAttrs = AttributeSet(innerAgg.groupingExpressions.flatMap(_.references))
if (!expand.producedAttributes.subsetOf(innerGroupByAttrs)) return false

The reasoning is that a non-distinct aggregate forces an Expand output column outside the inner GROUP BY. That only holds for aggregates to which RewriteDistinctAggregates assigns a dedicated Expand slot. An aggregate that reaches the inner Aggregate as count(1) — from COUNT(1) or COUNT(*) — references no attribute at all, so it always passes the guard.

This PR rejects duplicate-sensitive aggregates directly, by checking the inner aggregate's aggregate expressions with EliminateDistinct.isDuplicateAgnostic — the same soundness check RemoveRedundantAggregates uses for the same question:

val hasDuplicateSensitiveAgg = innerAgg.aggregateExpressions.exists(_.exists {
  case ae: AggregateExpression =>
    !ae.isDistinct && !EliminateDistinct.isDuplicateAgnostic(ae.aggregateFunction)
  case _ => false
})
if (hasDuplicateSensitiveAgg) return false

Why are the changes needed?

It is a correctness bug: the query returns the count of distinct rows where it must return the count of base rows.

CREATE TABLE oe USING parquet AS SELECT * FROM VALUES (1,5,7),(1,5,7),(1,5,7),(1,6,8),(2,9,9) AS t(k,a,b);

SET spark.sql.optimizer.optimizeExpandRatio=2;

SELECT k, COUNT(DISTINCT a), COUNT(DISTINCT b), COUNT(1) FROM oe GROUP BY k ORDER BY k;
-- before: [1,2,2,2], [2,1,1,1]   WRONG
-- after:  [1,2,2,4], [2,1,1,1]   k=1 has four rows

COUNT(*) behaves identically. The optimized plan before the fix shows the inserted pre-aggregate feeding the Expand whose downstream count(1) was supposed to count base rows:

Aggregate [k], [k, count(a) FILTER (gid=1), count(b) FILTER (gid=2),
                coalesce(first(count(1)) FILTER (gid=0), 0)]
+- Aggregate [k, a, b, gid], [k, a, b, gid, count(1)]
   +- Expand [[k,null,null,0], [k,a,null,1], [k,null,b,2]], [k, a, b, gid]
      +- Aggregate [k, a, b], [k, a, b]          <- inserted by OptimizeExpand
         +- Relation oe[k,a,b] parquet

Note on the JIRA: it also lists COUNT(a) on a non-nullable column as affected, on the grounds that the rewrite normalizes it to count(1). That does not reproduce on master — no such normalization happens, RewriteDistinctAggregates gives it its own Expand slot, and the existing attribute check already rejects it. A test for that case was written and passed before the fix, so it was dropped rather than added as a non-regression test.

Does this PR introduce any user-facing change?

Yes, it fixes wrong results for the queries above. Only when spark.sql.optimizer.optimizeExpandRatio is explicitly set, since the rule is disabled by default.

How was this patch tested?

New tests, verified failing before the fix and passing after:

• OptimizeExpandSuite — "SPARK-58387: skips when a non-distinct count(1) is present": asserts no pre-aggregate is inserted.
• OptimizeExpandQuerySuite — "SPARK-58387: correctness: count distinct with a non-distinct count(1)": the JIRA repro, checked against the expected answer.
• OptimizeExpandQuerySuite — "SPARK-58387: correctness: count distinct with a non-distinct count(*)": checked against the rule-disabled result.

Full suites pass: OptimizeExpandSuite (9 tests) and OptimizeExpandQuerySuite (10 tests).

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5)