You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "LuciferYang (via GitHub)" <gi...@apache.org> on 2026/04/01 05:20:21 UTC

[PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

LuciferYang opened a new pull request, #55130:
URL: https://github.com/apache/spark/pull/55130

   ### What changes were proposed in this pull request?
   Add an optimizer rule `OptimizeExpand` that inserts a de-duplication aggregate before the Expand operator produced by `RewriteDistinctAggregates`.
   
   Queries with multiple `COUNT(DISTINCT)` on different columns are rewritten by `RewriteDistinctAggregates` into a plan that duplicates each input row N times via an Expand operator (where N = number of distinct groups). This data amplification becomes a bottleneck as N grows.
   
   This rule transforms the plan from:
   
   ```
   Aggregate(outer) -> Aggregate(inner) -> Expand(Nx) -> child
   ```
   
   to:
   
   ```
   Aggregate(outer) -> Aggregate(inner) -> Expand(Nx) -> Aggregate(keys + all distinct cols) -> child
   ```
   
   The inserted pre-aggregation eliminates duplicate rows before the Expand applies its N-fold amplification, significantly reducing shuffled data volume.
   
   The optimization is controlled by `spark.sql.optimizer.optimizeExpandRatio` (default -1 = disabled, minimum 2). It only applies when:
   - The Expand projection count >= the configured threshold
   - All aggregates are pure distinct (no non-distinct aggregates like `SUM(value)`)
   - No FILTER clauses on distinct aggregates
   
   ### Why are the changes needed?
   For queries like `SELECT key, COUNT(DISTINCT a), COUNT(DISTINCT b), COUNT(DISTINCT c) FROM t GROUP BY key`, the Expand operator duplicates each row 3x. With 6 distinct aggregates, it's 6x. The pre-aggregation step eliminates redundant rows before this amplification, converting O(N * input_rows) to O(N * distinct_groups).
   
   ### Benchmark results
   
   Pure distinct aggregates (6 COUNT(DISTINCT), JDK 17):
   
   | Scenario | Baseline | Optimized | Speedup |
   |---|---|---|---|
   | 1K groups, moderate cardinality | 8276ms | 847ms | **9.8x** |
   | 100K groups, moderate cardinality | 13842ms | 3624ms | **3.8x** |
   | 1K groups, low cardinality (5 vals) | 7902ms | 770ms | **10.7x** |
   | No grouping key | 5560ms | 370ms | **22.4x** |
   
   Scaling with number of distinct aggregates (JDK 17):
   
   | Distinct count | Baseline | Optimized | Speedup |
   |---|---|---|---|
   | 2 | 2181ms | 660ms | **3.3x** |
   | 4 | 4872ms | 794ms | **6.1x** |
   | 6 | 8259ms | 836ms | **9.9x** |
   | 8 | 12642ms | 1603ms | **7.9x** |
   
   Mixed aggregates (with SUM): no regression — baseline and optimized are within noise.
   
   Results are consistent across JDK 17, 21, and 25.
   
   
   ### Does this PR introduce _any_ user-facing change?
   No. The optimization is disabled by default (`spark.sql.optimizer.optimizeExpandRatio = -1`). When enabled, query results are unchanged — only execution performance is affected.
   
   
   ### How was this patch tested?
   Add new tests:
   - Unit tests in `OptimizeExpandSuite` (plan-level rule tests)
   - End-to-end tests in `OptimizeExpandQuerySuite` (result correctness)
   - Benchmark in `ExpandBenchmark` with results for JDK 17, 21, and 25
   
   
   ### Was this patch authored or co-authored using generative AI tooling?
   Generated-by:  Claude Code
   


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55130:
URL: https://github.com/apache/spark/pull/55130#issuecomment-4223893998

   > BTW, just a question. Have you test this with other benchmarks like TPCDS (or TPC-H)?
   
   It's a bit of a pity that there are no scenarios in TPC-DS & TPC-H that would trigger the optimization.


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55130:
URL: https://github.com/apache/spark/pull/55130#issuecomment-4276326170

   Thank you @dongjoon-hyun and @xiaoxuandev 


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55130:
URL: https://github.com/apache/spark/pull/55130#issuecomment-4231932386

   rebased. If all the tests pass, I'll merge this pr tomorrow.


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on code in PR #55130:
URL: https://github.com/apache/spark/pull/55130#discussion_r3025719521


##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeExpandSuite.scala:
##########
@@ -0,0 +1,164 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules._
+import org.apache.spark.sql.internal.SQLConf
+
+class OptimizeExpandSuite extends PlanTest {

Review Comment:
   Thanks for the suggestion! Added tests in the latest commit:
   
   **Plan-level test** (`OptimizeExpandSuite`): verifies that for `COUNT(DISTINCT col1 + col2)`, the pre-aggregate groups by leaf attributes `(key, col1, col2, col3)` rather than `(key, col1+col2, col3)`. The comment documents the trade-off — dedup is less effective (more groups than strictly necessary) but correctness is preserved.
   
   **Correctness test** (`OptimizeExpandQuerySuite`): uses coprime moduli (`col1 = id % 7`, `col2 = id % 11`) to generate data where 77 unique `(col1, col2)` pairs map to only 17 unique `col1 + col2` values, ensuring the "less effective dedup" scenario is actually exercised. Verifies optimized results match the non-optimized baseline.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55130:
URL: https://github.com/apache/spark/pull/55130#issuecomment-4270006056

   Merged to master for Apache Spark 4.2.0. Thank you, @LuciferYang and all.


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun closed pull request #55130: [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)`
URL: https://github.com/apache/spark/pull/55130


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55130:
URL: https://github.com/apache/spark/pull/55130#issuecomment-4174983744

   [5818382](https://github.com/apache/spark/pull/55130/commits/58183828c929decbf116ad977a2ae0fa67a8101c) add new benchmark scenarios


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "xiaoxuandev (via GitHub)" <gi...@apache.org>.
xiaoxuandev commented on code in PR #55130:
URL: https://github.com/apache/spark/pull/55130#discussion_r3023448345


##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/OptimizeExpandSuite.scala:
##########
@@ -0,0 +1,164 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules._
+import org.apache.spark.sql.internal.SQLConf
+
+class OptimizeExpandSuite extends PlanTest {

Review Comment:
   Consider a test for `COUNT(DISTINCT col1 + col2)`. The pre-aggregate groups by leaf attributes (col1, col2) not the expression, so dedup is less effective but still correct. Worth documenting.



-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55130:
URL: https://github.com/apache/spark/pull/55130#issuecomment-4223761400

   BTW, just a question. Have you test this with other benchmarks like TPCDS (or TPC-H)?


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org

Re: [PR] [SPARK-56315][SQL] Pre-aggregate before `Expand` to reduce data amplification for multiple `COUNT(DISTINCT)` [spark]

Posted by "LuciferYang (via GitHub)" <gi...@apache.org>.
LuciferYang commented on PR #55130:
URL: https://github.com/apache/spark/pull/55130#issuecomment-4237164968

   cc @cloud-fan 


-- 
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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org