cloud-fan commented on a change in pull request #31994:
URL: https://github.com/apache/spark/pull/31994#discussion_r603753304



##########
File path: 
sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
##########
@@ -1543,4 +1543,27 @@ class AdaptiveQueryExecSuite
     assert(materializeLogs(0).startsWith("Materialize query stage 
BroadcastQueryStageExec"))
     assert(materializeLogs(1).startsWith("Materialize query stage 
ShuffleQueryStageExec"))
   }
+
+  test("SPARK-34899: Use origin plan if we can not coalesce shuffle 
partition") {
+    def check(ds: Dataset[Row], origin: ShuffleOrigin): Unit = {
+      ds.collect()
+      val plan = 
ds.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan
+      assert(collect(plan) {
+        case c: CustomShuffleReaderExec => c
+      }.isEmpty)
+      assert(collect(plan) {
+        case s: ShuffleExchangeExec if s.shuffleOrigin == origin && 
s.numPartitions == 3 => s
+      }.size == 1)
+      checkAnswer(ds, testData)
+    }
+
+    withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+      SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "10",
+      SQLConf.COALESCE_PARTITIONS_MIN_PARTITION_NUM.key -> "1",
+      SQLConf.SHUFFLE_PARTITIONS.key -> "3") {
+      check(testData.repartition(), REPARTITION)

Review comment:
       Can we create a new dataframe for testing here? It's unclear how many 
partitions `testData` has and what's the size for each partition.




-- 
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.

For queries about this service, please contact Infrastructure at:
[email protected]



---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to