diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeShuffleWithLocalRead.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeShuffleWithLocalRead.scala index cf1c7ecedd5b..997588013af0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeShuffleWithLocalRead.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/OptimizeShuffleWithLocalRead.scala @@ -55,12 +55,29 @@ object OptimizeShuffleWithLocalRead extends AQEShuffleReadRule { } } - private def createLocalRead(plan: SparkPlan): AQEShuffleReadExec = { - plan match { + private def createLocalRead(plan: SparkPlan): SparkPlan = { + val (shuffleStage, advisoryParallelism) = plan match { case c @ AQEShuffleReadExec(s: ShuffleQueryStageExec, _) => - AQEShuffleReadExec(s, getPartitionSpecs(s, Some(c.partitionSpecs.length))) + (s, Some(c.partitionSpecs.length)) case s: ShuffleQueryStageExec => - AQEShuffleReadExec(s, getPartitionSpecs(s, None)) + (s, None) + } + val partitionSpecs = getPartitionSpecs(shuffleStage, advisoryParallelism) + val maxReducerPartitionsPerTask = + conf.getConf(SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK) + val exceedsReducerPartitionLimit = advisoryParallelism.isDefined && partitionSpecs.exists { + case spec: PartialMapperPartitionSpec => + spec.endReducerIndex - spec.startReducerIndex > maxReducerPartitionsPerTask + case spec: CoalescedMapperPartitionSpec => + spec.numReducers > maxReducerPartitionsPerTask + case _ => false + } + // Local reads use mapper-oriented specs, whose reducer spans can differ from the bounded + // coalesced specs they replace. Keep the existing read if the rewrite would exceed the limit. + if (exceedsReducerPartitionLimit) { + plan + } else { + AQEShuffleReadExec(shuffleStage, partitionSpecs) } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala index 8cf6fbf921da..58fda8d76327 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala @@ -293,11 +293,41 @@ class AdaptiveQueryExecSuite } } + test("Preserve the reducer partition limit through local shuffle read optimization") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80", + SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "10", + SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK.key -> "2") { + val (_, adaptivePlan) = runAdaptiveAndVerifyResult( + "SELECT * FROM testData join testData2 ON key = a where value = '1'") + assert(findTopLevelBroadcastHashJoin(adaptivePlan).size == 1) + + // The two shuffle sides exercise both mapper-oriented local-read spec types. Neither rewrite + // is safe with this reducer limit, so the bounded coalesced reads must remain in the plan. + val localReads = collect(adaptivePlan) { + case read: AQEShuffleReadExec if read.isLocalRead => read + } + assert(localReads.isEmpty) + + val coalescedReads = collect(adaptivePlan) { + case read: AQEShuffleReadExec if read.hasCoalescedPartition => read + } + assert(coalescedReads.nonEmpty) + coalescedReads.flatMap(_.partitionSpecs).foreach { + case spec: CoalescedPartitionSpec => + assert(spec.endReducerIndex - spec.startReducerIndex <= 2) + case spec => fail(s"Unexpected shuffle partition spec: $spec") + } + } + } + test("Reuse the default parallelism in local shuffle read") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "80", - SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") { + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.COALESCE_PARTITIONS_MAX_REDUCER_PARTITIONS_PER_TASK.key -> "2") { val (plan, adaptivePlan) = runAdaptiveAndVerifyResult( "SELECT * FROM testData join testData2 ON key = a where value = '1'") val smj = findTopLevelSortMergeJoin(plan)