comphead commented on code in PR #5445:
URL: https://github.com/apache/datafusion-comet/pull/5445#discussion_r3844657348
##########
spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala:
##########
@@ -160,6 +190,59 @@ class CometTaskMetricsSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("native shuffle task metrics include existing child sort spill metrics
once") {
+ val expectedRecords = 20000L
+ val compressibleValue = "native-child-sort-spill-metrics-" * 8
+ withParquetTable(
+ (0 until expectedRecords.toInt).map(index => (index, compressibleValue)),
+ "tbl") {
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
+ CometConf.COMET_SHUFFLE_COMPRESSION_CODEC.key -> "zstd",
+ CometConf.COMET_SHUFFLE_NATIVE_MAX_BUFFER_BYTES.key -> "32k",
+ CometConf.COMET_BATCH_SIZE.key -> "1024",
+ CometConf.COMET_OFFHEAP_MEMORY_POOL_FRACTION.key -> "0.002",
+ CometConf.COMET_RESPECT_DATAFUSION_CONFIGS.key -> "true",
+ "spark.comet.datafusion.execution.spill_compression" -> "zstd",
+ "spark.comet.datafusion.execution.sort_spill_reservation_bytes" ->
"65536",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "4") {
+ val shuffled = sql("SELECT * FROM tbl")
+ .sortWithinPartitions($"_1".desc)
+ .repartition(4, $"_1")
+ val store = spark.sparkContext.statusStore
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+ val stagesBefore = store.stageList(null).map(_.stageId).toSet
+
+ assert(shuffled.collect().length == expectedRecords)
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+
+ val exchange = collectFirst(shuffled.queryExecution.executedPlan) {
+ case native: CometShuffleExchangeExec if native.shuffleType ==
CometNativeShuffle =>
+ native
+ }.getOrElse(fail("Expected a native shuffle exchange"))
+ val childSorts = collect(exchange.child) { case sort: CometSortExec =>
sort }
+ assert(childSorts.nonEmpty, s"Expected a native child
sort:\n${exchange.treeString}")
+
+ val writerDiskSpilled = exchange.metrics("spilled_bytes").value
+ val writerMemorySpilled =
exchange.metrics("memory_spilled_bytes").value
+ val childDiskSpilled =
childSorts.map(_.metrics("spilled_bytes").value).sum
+ assert(childDiskSpilled > 0L, "Native child sort did not spill")
+ assert(childSorts.forall(!_.metrics.contains("memory_spilled_bytes")))
+
+ val shuffleWriteStages = store
+ .stageList(null)
+ .filterNot(stage => stagesBefore.contains(stage.stageId))
+ .filter(_.shuffleWriteRecords > 0L)
Review Comment:
a nit: maybe we can combine it into a single filter
--
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]
