comphead commented on code in PR #4989:
URL: https://github.com/apache/datafusion-comet/pull/4989#discussion_r3625246923
##########
spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala:
##########
@@ -446,6 +446,39 @@ class CometNativeShuffleSuite extends CometTestBase with
AdaptiveSparkPlanHelper
}
}
+ test("native shuffle: maxBufferBytes triggers spilling") {
+ withParquetTable((0 until 20000).map(i => (i, (i + 1).toLong, s"str$i")),
"tbl") {
+ def spillCountWithMaxBufferBytes(maxBufferBytes: String): Long = {
+ var spillCount = 0L
+ withSQLConf(CometConf.COMET_SHUFFLE_MAX_BUFFER_BYTES.key ->
maxBufferBytes) {
+ val shuffled = sql("SELECT * FROM tbl").repartition(10, $"_1")
+ checkShuffleAnswer(shuffled, 1)
+
+ // Materialize the shuffled data. The metrics live on this plan's
exchange, so the
+ // query has to run before they are read; checkShuffleAnswer
executes copies built
+ // from the logical plan and leaves this one's accumulators
untouched.
+ shuffled.collect()
+ spillCount = find(shuffled.queryExecution.executedPlan) {
+ case _: CometShuffleExchangeExec => true
+ case _ => false
+ }.map(_.metrics("spill_count").value).get
+ }
Review Comment:
```suggestion
val spillCount =
shuffled.queryExecution.executedPlan.collectFirst {
case e: CometShuffleExchangeExec => e.metrics("spill_count").value
}.get
```
--
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]
