jerrypeng commented on code in PR #59523:
URL: https://github.com/apache/spark/pull/59523#discussion_r4429440922
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/PipelinedShuffleSqlSuite.scala:
##########
@@ -213,7 +218,89 @@ class PipelinedShuffleSqlSuite extends SparkFunSuite
}
}
+ private def distributedSequence(
+ spark: SparkSession,
+ input: DataFrame,
+ cache: Boolean): DataFrame = {
+ classic.Dataset.ofRows(spark.asInstanceOf[classic.SparkSession],
+ AttachDistributedSequence(AttributeReference("index", LongType)(),
+ input.queryExecution.analyzed, cache))
+ }
+
+ private def assertDistributedSequenceResult(df: DataFrame): Unit = {
+ val rows = df.collect()
+ assert(rows.map(_.getLong(0)).sorted === (0L until 4000L).toArray)
+ assert(rows.map(_.getLong(1)).sorted === (0L until 4000L).toArray)
+ }
+
for (aqe <- Seq(false, true)) {
+ test(s"distributed sequence cache construction retains regular shuffles
with AQE=$aqe") {
+ withPipelinedSession("pipelined-distributed-sequence-cache", aqe,
+ channelBatchSize = 8, channelQueueCapacity = 1) { spark =>
+ val cacheKey = "pandas_on_Spark.compute.default_index_cache"
+ for {
+ storage <- Seq(None, Some("MEMORY_ONLY"), Some("LOCAL_CHECKPOINT"),
Some("NONE"))
+ upperShuffle <- Seq(false, true)
+ } {
+ // An absent setting exercises pandas-on-Spark's default
MEMORY_AND_DISK_SER.
+ storage match {
+ case Some(level) => spark.conf.set(cacheKey, "\"" + level + "\"")
+ case None => spark.conf.unset(cacheKey)
+ }
+ val indexed = distributedSequence(spark,
+ spark.range(0, 4000, 1, 2).repartition(4).toDF(), cache = true)
+ val df = if (upperShuffle) indexed.repartition(3) else indexed
Review Comment:
Changed the upper shuffle to `repartition(1)` and added a `SinglePartition`
assertion so AQE continues below it and reaches the cached operator. The lower
four-partition exchange and exchange-count checks remain.
I verified the AQE upper-shuffle case by temporarily disabling only the
cached-operator predicate: both exchanges became pipelined, and the
regular-shuffle assertion failed before `collect()`. The corrected test passes
with the predicate present.
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
