viirya commented on code in PR #7059:
URL: https://github.com/apache/datafusion-comet/pull/7059#discussion_r4304232593
##########
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:
This covers the case from #7058 well. I'm wondering about the regime where
memory is genuinely short, where the partitions' batches together don't fit.
That is the one path where the growing pages meet `ArrowIPCWriter.spill()`. A
writer whose next, doubled page is refused spills only itself, since its own
usage is always at least `required`, and restarts the batch at 64 KiB. A writer
with no page yet spills the largest sibling instead.
None of the tests run that with growth on. This test asserts nothing
spilled. The page doubling test runs without memory pressure. "memory pressure
spills only writers of the requesting task" passes `pageSize, pageSize` (line
97), which turns the doubling off.
Could you add a case with, say, 32 partitions and a limit well below 32 full
batches? It could assert that every row is written, that the writers did spill,
that `tmm.getMemoryConsumptionForThisTask` returns to zero, and that the
average rows per written batch stays well above one. That last check would
guard the behavior this PR fixes in the regime where spilling is legitimate.
--
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]
