[SPARK-32083][SQL] Apply CoalesceShufflePartitions when input RDD has 0 partitions with AQE

[manuzhang]

What changes were proposed in this pull request?

As suggested by @cloud-fan in #28916 (comment), apply CoalesceShufflePartitions with partitionSpecs of Nil when ShuffleQueryStageExec#mapStats is None.

Why are the changes needed?

For SQL like

SELECT b, COUNT(t1.a) as cnt
FROM t1
INNER JOIN t2
ON t1.id = t2.id
WHERE t1.id > 10
GROUP BY b

when all ids of t1 are smaller than 10. Many unnecessary tasks are launched for the final shuffle stage because CoalesceShufflePartitions is skipped when input RDD has 0 partitions.

Before

After

Does this PR introduce any user-facing change?

No

How was this patch tested?

Updated tests.

[cloud-fan]

if a query stage has multiple leaf shuffles, and only one of them has 0-partition input RDD. What shall we do?

[manuzhang]

I think it's like coalescing one less shuffles and handled by the nonEmpty codes.

[cloud-fan]

cc @maryannxue @JkSelf @viirya

[viirya]

Can you add a comment for the case of 0-partition?

[SparkQA]

Test build #124648 has finished for PR 28954 at commit 5bf6de5.

• This patch fails Spark unit tests.
• This patch merges cleanly.
• This patch adds no public classes.

[cloud-fan]

After more thoughts, maybe a better way is to add a new rule in AdaptiveSparkPlanExec.optimizer, which converts LogicalQueryStage to empty LocalRelation if the size is 0.

This is not really "coalesce partitions" and we'd better not do it in CoalesceShufflePartitions.

[JkSelf]

@manuzhang I run the test("Empty stage coalesced to 0-partition RDD") in AdaptiveQueryExecSuite. It seems there is no unnecessary tasks for the empty partitions. The Spark UI is as following. And can you give a simple examples to reproduce this issue? Thanks.

The related code is :

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
val df1 = spark.range(10).withColumn("a", 'id)
val df2 = spark.range(10).withColumn("b", 'id)
val result = df1.where('a > 10).join(df2.where('b > 10), "id").groupBy('a).count()
result.show()

[manuzhang]

@JkSelf Try changing result.show() to result.collect()

[JkSelf]

@manuzhang It seems still only one stage and no unnecessary task for empty partitions.

related code:

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1")
val df1 = spark.range(10).withColumn("a", 'id)
val df2 = spark.range(10).withColumn("b", 'id)
val result = df1.where('a > 10).join(df2.where('b > 10), "id").groupBy('a).count()
result.collect()

[manuzhang]

@JkSelf Could you check the jobs tab ?

[JkSelf]

@manuzhang Here is the jobs tab.

[viirya]

When input RDD has 0 partitions, even CoalesceShufflePartitions is skipped, why we will launch unnecessary tasks?

[manuzhang]

@JkSelf @viirya
Here is the partial SQL UI of running the same example with default number of shuffle partitions. (the binary is built from master branch till June 29). You can see SortMergeJoin is followed by an Exchange of 200 partitions.

[manuzhang]

@JkSelf
I get the same result as you when I set the numPartitions of source to 1 (By default, it's 16 on my Mac), i.e.

val df1 = spark.range(0, 10, 1, 1).withColumn("a", 'id)
val df2 = spark.range(0, 10, 1, 1).withColumn("b", 'id)
val testDf = df1.where('a > 10).join(df2.where('b > 10), "id").groupBy('a).count()
testDf.collect()

Compare the execution plan with above.

[manuzhang]

@cloud-fan I find an issue with updating metrics if I convert a LogicalQueryStage to LocalRelation when its child's mapStats is empty.

java.util.NoSuchElementException: key not found: 514
	at scala.collection.immutable.Map$Map1.apply(Map.scala:114)
	at org.apache.spark.sql.execution.ui.SQLAppStatusListener.$anonfun$aggregateMetrics$11(SQLAppStatusListener.scala:257)

The missing key is from the metrics of a RangeExec sent to SQLAppStatusListener earlier