Vivek1106-04 opened a new pull request, #58102: URL: https://github.com/apache/spark/pull/58102
### 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](https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala) 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*: ```scala 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: ```scala 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. ```sql 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) -- 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]
