HyukjinKwon commented on code in PR #59282:
URL: https://github.com/apache/spark/pull/59282#discussion_r4246824969
##########
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:
**Nit (P3):** This assertion runs right after `queue1.spill(Long.MaxValue,
queue2)` has moved every in-memory buffer except the last to disk, so queue1
now holds [disk, disk, lastInMemory]. `spill` never moves disk buffers or the
last buffer, so this call returns 0 even if the `trigger eq this` short-circuit
is removed. No other test spills a queue for itself while it still has
spillable pages: the older tests pass `trigger = null`.
Could we check the self-triggered spill first, while queue1 still holds
several in-memory buffers? For example:
```scala
val used = queue1.getUsed
assert(queue1.spill(Long.MaxValue, queue1) === 0)
assert(queue1.getUsed === used)
assert(queue1.spill(Long.MaxValue, queue2) === used - smallPageSize)
```
**Verification:**
- **Regression:** With the self-trigger short-circuit removed from
HybridQueue.spill, the self-triggered spill check fails because in-memory
buffers move to disk and memory is released.
- **Behavior:** With the short-circuit in place, a self-triggered spill on a
queue with several in-memory buffers releases nothing and keeps its buffers,
and a spill triggered by a distinct equal-argument queue releases all but the
last page.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala:
##########
@@ -181,8 +181,13 @@ private[python] case class DiskRowQueue(
*
* HybridRowQueue could be safely appended in one thread, and pulled in
another thread in the same
* time.
+ *
+ * This is intentionally not a case class. As a memory consumer it must use
identity equality:
+ * otherwise two queues created with the same arguments in one task are equal,
so only one of
+ * them is tracked (and offered for spilling) by the [[TaskMemoryManager]],
and each refuses to
+ * spill when the other one triggers it.
*/
-case class HybridRowQueue(
+class HybridRowQueue(
Review Comment:
Thanks, verified in the current head. `close()` now swaps `queues` under the
lock and closes the detached buffers after releasing it, so it no longer holds
the queue monitor while calling `TaskMemoryManager.freePage`. A spill that runs
after the swap sees an empty list. The new test forces the two lock orders to
overlap and asserts the handshake and the spill result, so this concern is
resolved.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4228283943","thread_id":"inline:4228283943","verdict_sha256":"7782a87aad767a49068a837bf008a6a354048f59a8e28f0e2a3a92aae2c57855"}
-->
--
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]
