This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new a8265e36f96d [SPARK-58004][SQL] Add hard fan-in bound to AQE shuffle
partition coalescing
a8265e36f96d is described below
commit a8265e36f96d59c7f22369d401fc0c0146fdefa0
Author: Chao Sun <[email protected]>
AuthorDate: Tue Jul 7 21:42:47 2026 -0700
[SPARK-58004][SQL] Add hard fan-in bound to AQE shuffle partition coalescing
### Why are the changes needed?
AQE currently coalesces adjacent shuffle partitions using their total byte
size. It does not limit how many original reducer partitions one task may span.
For example, consider a shuffle with 100,000 mostly empty or tiny reducer
partitions and a 64 MiB advisory target. AQE may create a task for a range such
as `[0, 20,000)` because the range is still smaller than 64 MiB. The task must
nevertheless handle the metadata and shuffle fan-in for 20,000 reducer
partitions. On workloads with very high initial partition counts, this overhead
can make byte-only coalescing impractical.
Users need an independent bound on reducer-partition fan-in so they can
keep AQE coalescing enabled without allowing a single task to span an
arbitrarily large reducer range.
This resolves
[SPARK-58004](https://issues.apache.org/jira/browse/SPARK-58004).
### What changes were proposed in this pull request?
This PR adds:
```
spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask
```
The setting is a positive integer with a default of `2147483647`
(`Int.MaxValue`). The default preserves the existing behavior.
When the setting is `128`, for example, AQE may create `[0, 128)` and
`[128, 256)`, but never a `CoalescedPartitionSpec` spanning more than 128
original reducer partitions. The existing byte target still applies and may
split a range earlier.
The bound applies in both skew-aware and ordinary coalescing:
- Empty reducer partitions count toward the span.
- Backward merging of a small partition cannot cross the bound.
- The final small-tail merge cannot cross the bound.
- Multiple shuffle inputs use the same bounded reducer ranges.
- Existing `PartialReducerPartitionSpec` entries remain unchanged.
`ShufflePartitionsUtil.coalescePartitions` accepts the new bound as a
defaulted final argument. Its default comes from the SQL configuration, so
existing callers retain the unbounded behavior.
This limit is independent of the shuffle backend. It bounds the reducer
range represented by a `CoalescedPartitionSpec`, not the number of physical
shuffle blocks fetched by a task. AQE does not have reliable block-count
information at this planning point, so a separate block-count limit is outside
the scope of this PR.
### Does this PR introduce _any_ user-facing change?
Yes. Users can opt in with a setting such as:
```
spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask=128
```
With the default value, query behavior is unchanged.
### How was this patch tested?
Added focused tests for:
- The unbounded default.
- Many tiny and empty reducer partitions.
- A bound of one.
- Backward and final-tail merging.
- An individually oversized reducer partition.
- Skew partition specifications.
- Multiple shuffle inputs.
- Invalid zero and negative values.
- The bound in an executed AQE shuffle-read plan.
Ran:
```
./build/sbt "sql/testOnly
org.apache.spark.sql.execution.ShufflePartitionsUtilSuite"
./build/sbt "sql/testOnly
org.apache.spark.sql.execution.CoalesceShufflePartitionsSuite"
./build/sbt "hive/testOnly
org.apache.spark.sql.configaudit.SparkConfigBindingPolicySuite"
./build/sbt "catalyst/scalastyle" "sql/scalastyle" "sql/Test/scalastyle"
```
Results:
- `ShufflePartitionsUtilSuite`: 19 tests passed.
- `CoalesceShufflePartitionsSuite`: 20 tests passed.
- `SparkConfigBindingPolicySuite`: 3 tests passed.
- All three Scalastyle checks passed with zero findings.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #57087 from sunchao/dev/chao/codex/aqe-coalesce-fanin-oss.
Authored-by: Chao Sun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
docs/sql-performance-tuning.md | 8 ++
.../org/apache/spark/sql/internal/SQLConf.scala | 11 ++
.../adaptive/CoalesceShufflePartitions.scala | 4 +-
.../execution/adaptive/ShufflePartitionsUtil.scala | 52 +++++--
.../execution/CoalesceShufflePartitionsSuite.scala | 40 +++++-
.../sql/execution/ShufflePartitionsUtilSuite.scala | 152 ++++++++++++++++++++-
6 files changed, 252 insertions(+), 15 deletions(-)
diff --git a/docs/sql-performance-tuning.md b/docs/sql-performance-tuning.md
index ebae89cbe5e0..bbb449360287 100644
--- a/docs/sql-performance-tuning.md
+++ b/docs/sql-performance-tuning.md
@@ -292,6 +292,14 @@ This feature coalesces the post shuffle partitions based
on the map output stati
</td>
<td>3.2.0</td>
</tr>
+ <tr>
+
<td><code>spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask</code></td>
+ <td>Int.MaxValue</td>
+ <td>
+ The maximum number of contiguous reducer partitions that may be
coalesced into a single task. This limits the reducer-partition fan-in
independently of the advisory partition size.
+ </td>
+ <td>4.3.0</td>
+ </tr>
<tr>
<td><code>spark.sql.adaptive.coalescePartitions.initialPartitionNum</code></td>
<td>(none)</td>
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index c348ded151fc..d4463e060af4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -1158,6 +1158,17 @@ object SQLConf {
.checkValue(_ > 0, "minPartitionSize must be positive")
.createWithDefaultString("1MB")
+ val COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK =
+
buildConf("spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask")
+ .doc("The maximum number of contiguous reducer partitions that may be
coalesced into a " +
+ "single task. This puts a hard limit on the reducer-partition fan-in
of a coalesced " +
+ "shuffle task, independent of the advisory partition size.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .intConf
+ .checkValue(_ > 0, "The maximum number of reducer partitions per task
must be positive.")
+ .createWithDefault(Int.MaxValue)
+
val COALESCE_PARTITIONS_MIN_PARTITION_NUM =
buildConf("spark.sql.adaptive.coalescePartitions.minPartitionNum")
.internal()
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
index f172942f6ec5..d44667ea2fa3 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/CoalesceShufflePartitions.scala
@@ -110,7 +110,9 @@ case class CoalesceShufflePartitions(session: SparkSession)
extends AQEShuffleRe
advisoryTargetSize = advisoryTargetSize,
minNumPartitions = minNumPartitions,
minPartitionSize = minPartitionSize,
- coalesceGroup.shuffleStages.map(_.shuffleStage.id))
+ shuffleStageIds = coalesceGroup.shuffleStages.map(_.shuffleStage.id),
+ maxReducerPartitionsPerTask =
+
conf.getConf(SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK))
if (newPartitionSpecs.nonEmpty) {
coalesceGroup.shuffleStages.zip(newPartitionSpecs).map { case
(stageInfo, partSpecs) =>
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
index 7b161bd17f42..e0bad017b8f8 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/ShufflePartitionsUtil.scala
@@ -22,6 +22,7 @@ import scala.collection.mutable.ArrayBuffer
import org.apache.spark.{MapOutputStatistics, MapOutputTrackerMaster, SparkEnv}
import org.apache.spark.internal.{Logging, LogKeys}
import org.apache.spark.sql.execution.{CoalescedPartitionSpec,
PartialReducerPartitionSpec, ShufflePartitionSpec}
+import org.apache.spark.sql.internal.SQLConf
object ShufflePartitionsUtil extends Logging {
final val SMALL_PARTITION_FACTOR = 0.2
@@ -47,7 +48,10 @@ object ShufflePartitionsUtil extends Logging {
advisoryTargetSize: Long,
minNumPartitions: Int,
minPartitionSize: Long,
- shuffleStageIds: Seq[Int] = Seq.empty): Seq[Seq[ShufflePartitionSpec]] =
{
+ shuffleStageIds: Seq[Int] = Seq.empty,
+ maxReducerPartitionsPerTask: Int =
+
SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK.defaultValue.get
+ ): Seq[Seq[ShufflePartitionSpec]] = {
assert(mapOutputStatistics.length == inputPartitionSpecs.length)
if (mapOutputStatistics.isEmpty) {
@@ -66,14 +70,21 @@ object ShufflePartitionsUtil extends Logging {
log"${MDC(LogKeys.ADVISORY_TARGET_SIZE, advisoryTargetSize)}, actual
target size " +
log"${MDC(LogKeys.TARGET_SIZE, targetSize)}, minimum partition size: " +
log"${MDC(LogKeys.PARTITION_SIZE, minPartitionSize)}")
+ if (maxReducerPartitionsPerTask < Int.MaxValue) {
+ logInfo(log"For shuffle(${MDC(LogKeys.SHUFFLE_IDS, shuffleIds)}),
maximum reducer " +
+ log"partitions per task: " +
+ log"${MDC(LogKeys.MAX_NUM_PARTITIONS, maxReducerPartitionsPerTask)}")
+ }
// If `inputPartitionSpecs` are all empty, it means skew join optimization
is not applied.
if (inputPartitionSpecs.forall(_.isEmpty)) {
coalescePartitionsWithoutSkew(
- mapOutputStatistics, targetSize, minPartitionSize, shuffleStageIds)
+ mapOutputStatistics, targetSize, minPartitionSize, shuffleStageIds,
+ maxReducerPartitionsPerTask)
} else {
coalescePartitionsWithSkew(
- mapOutputStatistics, inputPartitionSpecs, targetSize,
minPartitionSize, shuffleStageIds)
+ mapOutputStatistics, inputPartitionSpecs, targetSize,
minPartitionSize, shuffleStageIds,
+ maxReducerPartitionsPerTask)
}
}
@@ -81,7 +92,8 @@ object ShufflePartitionsUtil extends Logging {
mapOutputStatistics: Seq[Option[MapOutputStatistics]],
targetSize: Long,
minPartitionSize: Long,
- shuffleStageIds: Seq[Int]): Seq[Seq[ShufflePartitionSpec]] = {
+ shuffleStageIds: Seq[Int],
+ maxReducerPartitionsPerTask: Int): Seq[Seq[ShufflePartitionSpec]] = {
// `ShuffleQueryStageExec#mapStats` returns None when the input RDD has 0
partitions,
// we should skip it when calculating the `partitionStartIndices`.
val validMetrics = mapOutputStatistics.flatten
@@ -104,7 +116,8 @@ object ShufflePartitionsUtil extends Logging {
val numPartitions = validMetrics.head.bytesByPartitionId.length
val newPartitionSpecs = coalescePartitionsAndGetSpecs(
- 0, numPartitions, validMetrics, targetSize, minPartitionSize)
+ 0, numPartitions, validMetrics, targetSize, minPartitionSize,
+ maxReducerPartitionsPerTask)
if (newPartitionSpecs.length < numPartitions) {
attachDataSize(mapOutputStatistics, newPartitionSpecs)
} else {
@@ -117,7 +130,8 @@ object ShufflePartitionsUtil extends Logging {
inputPartitionSpecs: Seq[Option[Seq[ShufflePartitionSpec]]],
targetSize: Long,
minPartitionSize: Long,
- shuffleStageIds: Seq[Int]): Seq[Seq[ShufflePartitionSpec]] = {
+ shuffleStageIds: Seq[Int],
+ maxReducerPartitionsPerTask: Int): Seq[Seq[ShufflePartitionSpec]] = {
// Do not coalesce if any of the map output stats are missing or if not
all shuffles have
// partition specs, which should not happen in practice.
if (!mapOutputStatistics.forall(_.isDefined) ||
!inputPartitionSpecs.forall(_.isDefined)) {
@@ -166,6 +180,7 @@ object ShufflePartitionsUtil extends Logging {
validMetrics,
targetSize,
minPartitionSize,
+ maxReducerPartitionsPerTask,
allowReturnEmpty = true)
newPartitionSpecsSeq.zip(attachDataSize(mapOutputStatistics,
partitionSpecs))
.foreach(spec => spec._1 ++= spec._2)
@@ -197,6 +212,7 @@ object ShufflePartitionsUtil extends Logging {
validMetrics,
targetSize,
minPartitionSize,
+ maxReducerPartitionsPerTask,
allowReturnEmpty = true)
newPartitionSpecsSeq.zip(attachDataSize(mapOutputStatistics,
partitionSpecs))
.foreach(spec => spec._1 ++= spec._2)
@@ -247,6 +263,7 @@ object ShufflePartitionsUtil extends Logging {
mapOutputStatistics: Seq[MapOutputStatistics],
targetSize: Long,
minPartitionSize: Long,
+ maxReducerPartitionsPerTask: Int,
allowReturnEmpty: Boolean = false): Seq[CoalescedPartitionSpec] = {
// `minPartitionSize` is useful for cases like [64MB, 0.5MB, 64MB]: we
can't do coalesce,
// because merging 0.5MB to either the left or right partition will exceed
the target size.
@@ -264,6 +281,10 @@ object ShufflePartitionsUtil extends Logging {
}
}
+ def isWithinMaxReducerPartitions(rangeStart: Int, rangeEnd: Int): Boolean
= {
+ rangeEnd - rangeStart <= maxReducerPartitionsPerTask
+ }
+
while (i < end) {
// We calculate the total size of i-th shuffle partitions from all
shuffles.
var totalSizeOfCurrentPartition = 0L
@@ -273,13 +294,21 @@ object ShufflePartitionsUtil extends Logging {
j += 1
}
- // If including the `totalSizeOfCurrentPartition` would exceed the
target size and the
- // current size has reached the `minPartitionSize`, then start a new
coalesced partition.
- if (i > latestSplitPoint && coalescedSize + totalSizeOfCurrentPartition
> targetSize) {
+ // The reducer-partition limit is a hard bound, so it takes precedence
over the size-based
+ // packing and minimum partition size. Advance the split point even for
an empty range so
+ // empty reducer partitions count toward the bound without creating an
empty task.
+ if (i > latestSplitPoint && i - latestSplitPoint >=
maxReducerPartitionsPerTask) {
+ createPartitionSpec()
+ latestSplitPoint = i
+ latestPartitionSize = coalescedSize
+ coalescedSize = totalSizeOfCurrentPartition
+ } else if (i > latestSplitPoint &&
+ coalescedSize + totalSizeOfCurrentPartition > targetSize) {
if (coalescedSize < minPartitionSize) {
// the current partition size is below `minPartitionSize`.
// pack it with the smaller one between the two adjacent partitions
(before and after).
- if (latestPartitionSize > 0 && latestPartitionSize <
totalSizeOfCurrentPartition) {
+ if (latestPartitionSize > 0 && latestPartitionSize <
totalSizeOfCurrentPartition &&
+
isWithinMaxReducerPartitions(partitionSpecs.last.startReducerIndex, i)) {
// pack with the before partition.
partitionSpecs(partitionSpecs.length - 1) =
CoalescedPartitionSpec(partitionSpecs.last.startReducerIndex, i)
@@ -304,7 +333,8 @@ object ShufflePartitionsUtil extends Logging {
i += 1
}
- if (coalescedSize < minPartitionSize && latestPartitionSize > 0) {
+ if (coalescedSize < minPartitionSize && latestPartitionSize > 0 &&
+ isWithinMaxReducerPartitions(partitionSpecs.last.startReducerIndex,
end)) {
// pack with the last partition.
partitionSpecs(partitionSpecs.length - 1) =
CoalescedPartitionSpec(partitionSpecs.last.startReducerIndex, end)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
index 3ef22ccb77e9..fbc7ba71d6ce 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/CoalesceShufflePartitionsSuite.scala
@@ -62,7 +62,8 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite
with SQLConfHelper
f: SparkSession => Unit,
targetPostShuffleInputSize: Int,
minNumPostShufflePartitions: Option[Int],
- enableIOEncryption: Boolean = false): Unit = {
+ enableIOEncryption: Boolean = false,
+ maxReducerPartitionsPerTask: Int = Int.MaxValue): Unit = {
val sparkConf =
new SparkConf(false)
.setMaster("local[*]")
@@ -83,6 +84,11 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite
with SQLConfHelper
case None =>
sparkConf.set(SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key, "1")
}
+ if (maxReducerPartitionsPerTask < Int.MaxValue) {
+ sparkConf.set(
+ SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK.key,
+ maxReducerPartitionsPerTask.toString)
+ }
val spark = SparkSession.builder()
.config(sparkConf)
@@ -90,6 +96,38 @@ class CoalesceShufflePartitionsSuite extends SparkFunSuite
with SQLConfHelper
try f(spark) finally spark.stop()
}
+ test("limit the number of reducer partitions per coalesced task") {
+ val test: SparkSession => Unit = { spark =>
+ val agg = spark.range(0, 100, 1, numInputPartitions)
+ .selectExpr("id % 5 as key")
+ .groupBy("key")
+ .count()
+
+ QueryTest.checkAnswer(
+ agg,
+ spark.range(0, 5).selectExpr("id as key", "20L as count")
+ .collect().toImmutableArraySeq)
+
+ val finalPlan = stripAQEPlan(agg.queryExecution.executedPlan)
+ val shuffleReads = finalPlan.collect {
+ case r @ CoalescedShuffleRead() => r
+ }
+ assert(shuffleReads.length === 1)
+ val reducerRanges = shuffleReads.head.partitionSpecs.map {
+ case spec: CoalescedPartitionSpec =>
+ (spec.startReducerIndex, spec.endReducerIndex)
+ case spec => fail(s"Unexpected shuffle partition spec: $spec")
+ }
+ assert(reducerRanges === Seq((0, 2), (2, 4), (4, 5)))
+ }
+
+ withSparkSession(
+ test,
+ targetPostShuffleInputSize = 1024 * 1024,
+ minNumPostShufflePartitions = None,
+ maxReducerPartitionsPerTask = 2)
+ }
+
Seq(Some(5), None).foreach { minNumPostShufflePartitions =>
val testNameNote = minNumPostShufflePartitions match {
case Some(numPartitions) => "(minNumPostShufflePartitions: " +
numPartitions + ")"
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/ShufflePartitionsUtilSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/ShufflePartitionsUtilSuite.scala
index c87ba4b7d18e..6146a2f6dc64 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/ShufflePartitionsUtilSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/ShufflePartitionsUtilSuite.scala
@@ -20,6 +20,7 @@ package org.apache.spark.sql.execution
import org.apache.spark.{LocalSparkContext, MapOutputStatistics,
MapOutputTrackerMaster, SparkConf, SparkContext, SparkEnv, SparkFunSuite}
import org.apache.spark.scheduler.MapStatus
import org.apache.spark.sql.execution.adaptive.ShufflePartitionsUtil
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.storage.BlockManagerId
import org.apache.spark.util.ArrayImplicits._
@@ -30,7 +31,8 @@ class ShufflePartitionsUtilSuite extends SparkFunSuite with
LocalSparkContext {
expectedPartitionStartIndices: Seq[Seq[CoalescedPartitionSpec]],
targetSize: Long,
minNumPartitions: Int = 1,
- minPartitionSize: Long = 0): Unit = {
+ minPartitionSize: Long = 0,
+ maxReducerPartitionsPerTask: Int = Int.MaxValue): Unit = {
val mapOutputStatistics = bytesByPartitionIdArray.zipWithIndex.map {
case (bytesByPartitionId, index) =>
Some(new MapOutputStatistics(index, bytesByPartitionId))
@@ -40,7 +42,9 @@ class ShufflePartitionsUtilSuite extends SparkFunSuite with
LocalSparkContext {
Seq.fill(mapOutputStatistics.length)(None),
targetSize,
minNumPartitions,
- minPartitionSize)
+ minPartitionSize,
+ shuffleStageIds = Seq.empty,
+ maxReducerPartitionsPerTask = maxReducerPartitionsPerTask)
assert(estimatedPartitionStartIndices.length ===
expectedPartitionStartIndices.length)
estimatedPartitionStartIndices.zip(expectedPartitionStartIndices).foreach {
case (actual, expect) => assert(actual === expect)
@@ -197,6 +201,150 @@ class ShufflePartitionsUtilSuite extends SparkFunSuite
with LocalSparkContext {
}
}
+ test("max reducer partitions per task is unbounded by default") {
+ val bytesByPartitionId = Array.fill[Long](10)(1)
+ val expected = Seq(CoalescedPartitionSpec(0, 10, 10))
+ checkEstimation(Array(bytesByPartitionId), Seq(expected), targetSize = 100)
+ }
+
+ test("max reducer partitions per task limits many tiny partitions") {
+ val bytesByPartitionId = Array.fill[Long](10)(1)
+ val expected = Seq(
+ CoalescedPartitionSpec(0, 4, 4),
+ CoalescedPartitionSpec(4, 8, 4),
+ CoalescedPartitionSpec(8, 10, 2))
+ checkEstimation(
+ Array(bytesByPartitionId),
+ Seq(expected),
+ targetSize = 100,
+ maxReducerPartitionsPerTask = 4)
+ }
+
+ test("max reducer partitions per task of one does not coalesce") {
+ val bytesByPartitionId = Array.fill[Long](3)(1)
+ checkEstimation(
+ Array(bytesByPartitionId),
+ expectedPartitionStartIndices = Seq.empty,
+ targetSize = 100,
+ maxReducerPartitionsPerTask = 1)
+ }
+
+ test("max reducer partitions per task counts empty partitions") {
+ val bytesByPartitionId = Array[Long](1, 0, 0, 0, 0, 0, 0, 0, 1)
+ val expected = Seq(
+ CoalescedPartitionSpec(0, 4, 1),
+ CoalescedPartitionSpec(8, 9, 1))
+ checkEstimation(
+ Array(bytesByPartitionId),
+ Seq(expected),
+ targetSize = 100,
+ maxReducerPartitionsPerTask = 4)
+ }
+
+ test("max reducer partitions per task prevents an oversized backward merge")
{
+ val bytesByPartitionId = Array[Long](30, 30, 1, 1, 100)
+ val expected = Seq(
+ CoalescedPartitionSpec(0, 2, 60),
+ CoalescedPartitionSpec(2, 5, 102))
+ checkEstimation(
+ Array(bytesByPartitionId),
+ Seq(expected),
+ targetSize = 60,
+ minPartitionSize = 10,
+ maxReducerPartitionsPerTask = 3)
+ }
+
+ test("max reducer partitions per task prevents an oversized final tail
merge") {
+ val bytesByPartitionId = Array[Long](60, 1, 1, 1)
+ val expected = Seq(
+ CoalescedPartitionSpec(0, 1, 60),
+ CoalescedPartitionSpec(1, 4, 3))
+ checkEstimation(
+ Array(bytesByPartitionId),
+ Seq(expected),
+ targetSize = 60,
+ minPartitionSize = 10,
+ maxReducerPartitionsPerTask = 3)
+ }
+
+ test("max reducer partitions per task keeps an oversized reducer partition
intact") {
+ val bytesByPartitionId = Array[Long](200, 1, 1, 1)
+ val expected = Seq(
+ CoalescedPartitionSpec(0, 1, 200),
+ CoalescedPartitionSpec(1, 3, 2),
+ CoalescedPartitionSpec(3, 4, 1))
+ checkEstimation(
+ Array(bytesByPartitionId),
+ Seq(expected),
+ targetSize = 100,
+ maxReducerPartitionsPerTask = 2)
+ }
+
+ test("max reducer partitions per task applies to multiple shuffles") {
+ val bytesByPartitionId1 = Array.fill[Long](7)(1)
+ val bytesByPartitionId2 = Array.tabulate[Long](7)(i => if (i % 2 == 0) 2
else 0)
+ val expected1 = Seq(
+ CoalescedPartitionSpec(0, 3, 3),
+ CoalescedPartitionSpec(3, 6, 3),
+ CoalescedPartitionSpec(6, 7, 1))
+ val expected2 = Seq(
+ CoalescedPartitionSpec(0, 3, 4),
+ CoalescedPartitionSpec(3, 6, 2),
+ CoalescedPartitionSpec(6, 7, 2))
+ checkEstimation(
+ Array(bytesByPartitionId1, bytesByPartitionId2),
+ Seq(expected1, expected2),
+ targetSize = 100,
+ maxReducerPartitionsPerTask = 3)
+ }
+
+ test("max reducer partitions per task preserves skew partition specs") {
+ val bytesByPartitionId1 = Array.fill[Long](8)(1)
+ val bytesByPartitionId2 = Array.fill[Long](8)(2)
+ val skewSpecs = Seq(
+ PartialReducerPartitionSpec(2, 0, 1, 1),
+ PartialReducerPartitionSpec(2, 1, 2, 1))
+ val specs1 = Seq.tabulate(2)(i => CoalescedPartitionSpec(i, i + 1, 1)) ++
+ skewSpecs ++ Seq.tabulate(5)(i => CoalescedPartitionSpec(i + 3, i + 4,
1))
+ val otherSideSkewSpecs = Seq.fill(2)(CoalescedPartitionSpec(2, 3, 2))
+ val specs2 = Seq.tabulate(2)(i => CoalescedPartitionSpec(i, i + 1, 2)) ++
+ otherSideSkewSpecs ++ Seq.tabulate(5)(i => CoalescedPartitionSpec(i + 3,
i + 4, 2))
+
+ val expected1 = Seq(CoalescedPartitionSpec(0, 2, 2)) ++ skewSpecs ++ Seq(
+ CoalescedPartitionSpec(3, 5, 2),
+ CoalescedPartitionSpec(5, 7, 2),
+ CoalescedPartitionSpec(7, 8, 1))
+ val expected2 = Seq(CoalescedPartitionSpec(0, 2, 4)) ++ otherSideSkewSpecs
++ Seq(
+ CoalescedPartitionSpec(3, 5, 4),
+ CoalescedPartitionSpec(5, 7, 4),
+ CoalescedPartitionSpec(7, 8, 2))
+
+ val coalesced = ShufflePartitionsUtil.coalescePartitions(
+ Seq(
+ Some(new MapOutputStatistics(0, bytesByPartitionId1)),
+ Some(new MapOutputStatistics(1, bytesByPartitionId2))),
+ Seq(Some(specs1), Some(specs2)),
+ advisoryTargetSize = 100,
+ minNumPartitions = 1,
+ minPartitionSize = 0,
+ shuffleStageIds = Seq.empty,
+ maxReducerPartitionsPerTask = 2)
+ assert(coalesced === Seq(expected1, expected2))
+ }
+
+ test("max reducer partitions per task must be positive") {
+ val conf = new SQLConf
+
assert(conf.getConf(SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK)
===
+ Int.MaxValue)
+ Seq(0, -1).foreach { value =>
+ intercept[IllegalArgumentException] {
+ conf.setConfString(
+ SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK.key,
+ value.toString)
+ }
+ }
+ }
+
test("enforce minimal number of coalesced partitions") {
val targetSize = 100
val minNumPartitions = 2
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
