andygrove commented on code in PR #4989:
URL: https://github.com/apache/datafusion-comet/pull/4989#discussion_r3625770206
##########
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:
Calling `collectFirst` on the plan directly resolves to
`TreeNode.collectFirst`, which stops at `AdaptiveSparkPlanExec` since that is a
`LeafExecNode`, so it would miss the exchange whenever AQE is on.
I went with the `AdaptiveSparkPlanHelper` version instead, which unwraps AQE
and query stage nodes in `allChildren`. Same single partial function you
suggested, just the curried form:
```scala
spillCount = collectFirst(shuffled.queryExecution.executedPlan) {
case e: CometShuffleExchangeExec => e.metrics("spill_count").value
}.get
```
Kept it as an assignment to the outer `var` so the enclosing helper method
still returns the count. Applied in 44f2f0f19.
##########
native/shuffle/src/partitioners/multi_partition.rs:
##########
@@ -174,6 +178,7 @@ impl<T: PartitionWriter>
MultiPartitionShuffleRepartitioner<T> {
runtime: Arc<RuntimeEnv>,
batch_size: usize,
tracing_enabled: bool,
+ max_buffer_bytes: usize,
Review Comment:
Done
--
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]
