viirya commented on code in PR #59282:
URL: https://github.com/apache/spark/pull/59282#discussion_r4250951514
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/RowQueueSuite.scala:
##########
@@ -170,4 +173,122 @@ class RowQueueSuite extends SparkFunSuite with
EncryptionFunSuite {
queue.close()
}
}
+
+ // Small pages so that a few rows span several in-memory queues.
+ private val smallPageSize = 1024L
+ private val smallPageConf = new SparkConf(false).set(BUFFER_PAGESIZE,
smallPageSize)
+
+ /** Creates two distinct queues with the same constructor arguments. */
+ private def queuesWithSameArguments(
+ taskM: TaskMemoryManager): (HybridRowQueue, HybridRowQueue) = {
+ val serManager = createSerializerManager(smallPageConf)
+ val tempDir = Utils.createTempDir().getCanonicalFile
+ // The second queue uses named arguments to check that they match the
former case class.
+ (HybridRowQueue(taskM, tempDir, 1, serManager),
+ HybridRowQueue(memManager = taskM, tempDir = tempDir, numFields = 1,
serMgr = serManager))
+ }
+
+ private def addRows(queue: HybridRowQueue, n: Int): Unit = {
+ val row = new UnsafeRow(1)
+ row.pointTo(new Array[Byte](16), 16)
+ (0 until n).foreach { i =>
+ row.setLong(0, i)
+ assert(queue.add(row) === QueueMode.IN_MEMORY)
+ }
+ }
+
+ private def drainAndClose(queue: HybridRowQueue, n: Int): Unit = {
+ (0 until n).foreach { i =>
+ assert(queue.remove().getLong(0) === i)
+ }
+ queue.close()
+ }
+
+ test("hybrid queues with the same fields are distinct memory consumers") {
+ val mem = new TestMemoryManager(smallPageConf)
+ val taskM = new TaskMemoryManager(mem, 0)
+ val (queue1, queue2) = queuesWithSameArguments(taskM)
+ assert(queue1 != queue2)
+ assert(queue1.toString.startsWith("HybridRowQueue(numFields=1,
lockFree=false)@"))
+ assert(queue1.toString != queue2.toString)
+
+ // queue1 is registered first and holds a single page, which it never
spills because the
+ // last queue is kept for writing. queue2 holds several pages and can
release all but one.
+ addRows(queue1, 1)
+ val n = 150
+ addRows(queue2, n)
+ assert(queue2.numQueues() > 2)
+ val queue2Used = queue2.getUsed
+
+ // Memory pressure from another consumer must be able to spill queue2.
+ mem.limit(0)
+ val consumer = new TestMemoryConsumer(taskM)
+ consumer.use(smallPageSize)
+ assert(consumer.getUsed === smallPageSize)
+ assert(queue2.getUsed === smallPageSize)
+ assert(queue2.getUsed < queue2Used)
+ consumer.free(smallPageSize)
+
+ drainAndClose(queue1, 1)
+ drainAndClose(queue2, n)
+ assert(taskM.cleanUpAllAllocatedMemory() === 0)
+ }
+
+ test("hybrid queue spills for another queue with the same fields") {
+ val taskM = new TaskMemoryManager(new TestMemoryManager(smallPageConf), 0)
+ val (queue1, queue2) = queuesWithSameArguments(taskM)
+ val n = 150
+ addRows(queue1, n)
+ assert(queue1.numQueues() > 2)
+ val used = queue1.getUsed
+ assert(queue1.spill(Long.MaxValue, queue2) === used - smallPageSize)
+ assert(queue1.spill(Long.MaxValue, queue1) === 0)
Review Comment:
Good catch, thanks. Done in b3f9259. The self-triggered spill is now checked
first, while queue1 still holds several in-memory buffers, and the test also
asserts that its memory usage is unchanged. I checked that the test now fails
when the `trigger eq this` short-circuit is removed.
--
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]
