andygrove commented on code in PR #7059:
URL: https://github.com/apache/datafusion-comet/pull/7059#discussion_r4308016940
##########
spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometDiskBlockWriterSuite.scala:
##########
@@ -216,6 +223,125 @@ class CometDiskBlockWriterSuite extends AnyFunSuite {
}
}
+ test("bypass writers with more partitions than full pages fit write without
spilling") {
+ // Every partition writer holds a page as soon as it has a row. With pages
of Spark's page
+ // size, a task fits only a few of them, so with more partitions than that
each new page
+ // spilled a sibling writer holding about one row (#7058).
+ val conf = new SparkConf()
+ .setMaster("local[1]")
+ .setAppName("CometDiskBlockWriterSuite")
+ .set("spark.memory.offHeap.enabled", "true")
+ .set("spark.memory.offHeap.size", "1g")
+ .set("spark.buffer.pageSize", "1m")
+ val sc = new SparkContext(conf)
+ val memoryManager = new TestMemoryManager(conf)
+ // Room for eight pages of Spark's size, for 32 partitions whose rows need
far less.
+ memoryManager.limit(8 * 1024 * 1024)
+ val tmm = new TaskMemoryManager(memoryManager, 0L)
+ try {
+ val taskContext = newTaskContext(tmm, 0L)
+ val partitions = 32
+ val rowsPerPartition = 100
+ val partitioner = new Partitioner {
+ override def numPartitions: Int = partitions
+ override def getPartition(key: Any): Int = key.asInstanceOf[Int] %
partitions
+ }
+ val dep = new CometShuffleDependency[Int, UnsafeRow, UnsafeRow](
+ _rdd = sc.parallelize(Seq.empty[(Int, UnsafeRow)], 1),
+ partitioner = partitioner,
+ serializer = new UnsafeRowSerializer(1),
+ schema = Some(schema),
+ decodeTime = new SQLMetric("nsTiming"))
+ val shuffleExecutorComponents = new
LocalDiskShuffleExecutorComponents(conf)
+ shuffleExecutorComponents.initializeExecutor(
+ sc.applicationId,
+ "driver",
+ JMap.of[String, String]())
+ val writer = new CometBypassMergeSortShuffleWriter[Int, UnsafeRow](
+ SparkEnv.get.blockManager,
+ tmm,
+ taskContext,
+ new CometBypassMergeSortShuffleHandle[Int, UnsafeRow](dep.shuffleId,
dep),
+ 0L,
+ conf,
+ taskContext.taskMetrics.shuffleWriteMetrics,
+ shuffleExecutorComponents,
+ null)
+
+ val toUnsafe = UnsafeProjection.create(schema)
+ // Round-robin, so that each partition gets a row before any gets a
second one.
+ val rows = (0 until partitions * rowsPerPartition).iterator.map { i =>
+ (i % partitions, toUnsafe(InternalRow(new Array[Byte](100))))
+ }
+ writer.write(rows)
+ assert(writer.stop(true).isDefined)
+
+ // No writer had to make room for another, so each partition was written
once, on close.
+ assert(taskContext.taskMetrics.memoryBytesSpilled == 0)
Review Comment:
Added "growing pages under real memory pressure spill but still write full
batches" in 175fb30d3d. It runs 32 writers round-robin against a 3 MiB limit,
which is well under the roughly 8 MiB their batches fill, with pages growing
from 64 KiB to 1 MiB. It asserts every row is written, memory and disk bytes
spilled are above zero, `getMemoryConsumptionForThisTask` and the allocator are
back to zero, and the average rows per written batch is at least 100. I checked
it catches the regression: with the initial page set to the maximum it fails at
exactly 1 row per batch (63997 rows in 63997 batches).
--
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]
