[SPARK-58004][SQL][FOLLOWUP] Preserve AQE fan-in bound through local shuffle reads

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)
     }
   }
 

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)