You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by vi...@apache.org on 2026/10/12 22:16:13 UTC
(spark) branch master updated: [SPARK-60071][SQL] Use identity equality for HybridRowQueue
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new ef4f7e0ece1e [SPARK-60071][SQL] Use identity equality for HybridRowQueue
ef4f7e0ece1e is described below
commit ef4f7e0ece1e02a4f07437deea824e029e3ad9a2
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Mon Oct 12 15:16:00 2026 -0700
[SPARK-60071][SQL] Use identity equality for HybridRowQueue
### What changes were proposed in this pull request?
`HybridQueue`, the base class of `HybridRowQueue`, now defines final identity `equals`/`hashCode`, and `HybridQueue.spill` checks `trigger eq this`. Every current and future subclass, including a case class, is therefore an identity-equal memory consumer. `HybridRowQueue` is changed from a case class to a regular class, because copying or destructuring a memory consumer by its constructor arguments is not meaningful. It gets a `toString` that shows its width, its `lockFree` flag and i [...]
`HybridQueue.close()` now detaches its queue list while holding the queue's lock and closes the detached queues, which frees their pages, after releasing it.
### Why are the changes needed?
`HybridRowQueue` is a `MemoryConsumer` (via `HybridQueue`), but as a case class its `equals`/`hashCode` compare the constructor arguments. Two queues created in the same task with the same `TaskMemoryManager`, temp dir, number of fields, `SerializerManager` and `lockFree` flag are therefore equal, which breaks memory management:
1. `TaskMemoryManager` tracks consumers in a `HashSet`. A second queue that equals an already registered one is never added, so it is never offered for spilling when another consumer needs memory, and its memory is missing from the per-consumer breakdown reported on OOM.
2. `HybridQueue.spill(size, trigger)` returns 0 when `trigger == this`, which calls `equals`. A queue therefore refuses to spill when an equal but distinct queue triggers the spill.
This happens in real queries. For example, with `coalesce`, `CoalescedRDD` evaluates several parent partitions in one task, and each one creates its own queue in `EvalPythonEvaluatorFactory` with the same width and local dir. These queues are only closed at task completion. Both sides of a sort-merge join with Python UDF eval nodes of the same input width are another example.
Making equal-argument queues spill candidates for each other also exposes them to an existing lock-order inversion. `TaskMemoryManager.acquireExecutionMemory` holds the `TaskMemoryManager` monitor while it calls `spill` on another consumer, and `HybridQueue.spill` then takes the queue's lock. `HybridQueue.close()` took the locks in the opposite order: it freed pages through `TaskMemoryManager.freePage` while holding the queue's lock. With pipelined Python UDF execution, `add()` runs o [...]
This was found during the review of #58978.
### Does this PR introduce _any_ user-facing change?
No, apart from the fix: the extra queues are now spilled under memory pressure and show up in the memory usage breakdown, where each queue is shown as `HybridRowQueue(numFields=..., lockFree=...)<identity hash>`.
### How was this patch tested?
Added three tests to `RowQueueSuite`, all of which fail without the fix:
- `hybrid queues with the same fields are distinct memory consumers`: memory pressure from another consumer must be able to spill the second of two equal-argument queues (without the fix, the requesting consumer gets 0 bytes).
- `hybrid queue spills for another queue with the same fields`: `spill` triggered by a distinct queue with the same arguments must release memory (without the fix, it returns 0).
- `closing a hybrid queue does not deadlock with a concurrent spill of it`: latches and an overridden `TaskMemoryManager.freePage` force `close()` and a spill of the same queue, run with the `TaskMemoryManager` monitor held, to overlap (without the fix, the two threads deadlock and the test times out).
```
build/sbt 'sql/testOnly *RowQueueSuite *PythonForeachWriterSuite'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
Closes #59282 from viirya/hybrid-row-queue-identity-equality.
Authored-by: Liang-Chi Hsieh <vi...@gmail.com>
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
.../spark/sql/execution/python/HybridQueue.scala | 27 ++++-
.../spark/sql/execution/python/RowQueue.scala | 22 +++-
.../spark/sql/execution/python/RowQueueSuite.scala | 127 ++++++++++++++++++++-
3 files changed, 168 insertions(+), 8 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/HybridQueue.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/HybridQueue.scala
index 4eb0a1425b91..c7f9f8817d42 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/HybridQueue.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/HybridQueue.scala
@@ -62,6 +62,16 @@ abstract class HybridQueue[T, Q <: Queue[T]](
// exposed for testing
private[python] def numQueues(): Int = queues.size()
+ // A memory consumer must use identity equality. TaskMemoryManager tracks consumers in a
+ // HashSet, so two queues that compare equal would be tracked (and offered for spilling) as one.
+ // These are final so that no subclass, including a case class, can switch to value equality.
+ final override def equals(other: Any): Boolean = other match {
+ case ref: AnyRef => this eq ref
+ case _ => false
+ }
+
+ final override def hashCode(): Int = System.identityHashCode(this)
+
protected def createDiskQueue(): Q
protected def createInMemoryQueue(page: MemoryBlock): Q
protected def getRequiredSize(item: T): Long
@@ -70,7 +80,7 @@ abstract class HybridQueue[T, Q <: Queue[T]](
protected def isReadingFromDiskQueue: Boolean = !isInMemoryQueue(reading)
def spill(size: Long, trigger: MemoryConsumer): Long = {
- if (trigger == this) {
+ if (trigger eq this) {
// When it's triggered by itself, it should write upcoming elements into disk instead of
// copying the elements already in the queue.
return 0L
@@ -171,10 +181,17 @@ abstract class HybridQueue[T, Q <: Queue[T]](
reading.close()
reading = null.asInstanceOf[Q]
}
- synchronized {
- while (!queues.isEmpty) {
- queues.remove().close()
- }
+ // Detach the queues under the lock but close them (which frees their pages through the
+ // TaskMemoryManager) after releasing it. `spill` is called with the TaskMemoryManager's
+ // monitor held and then takes this queue's lock, so freeing pages while holding this lock
+ // could deadlock with a concurrent spill of this queue triggered from another thread.
+ val detached = synchronized {
+ val current = queues
+ queues = new java.util.LinkedList[Q]()
+ current
+ }
+ while (!detached.isEmpty) {
+ detached.remove().close()
}
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala
index c0531fcf303c..f83518ed6a94 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala
@@ -181,8 +181,11 @@ private[python] case class DiskRowQueue(
*
* HybridRowQueue could be safely appended in one thread, and pulled in another thread in the same
* time.
+ *
+ * This is not a case class: it is a memory consumer with identity equality (see [[HybridQueue]]),
+ * so copying or destructuring it by its constructor arguments would be meaningless.
*/
-case class HybridRowQueue(
+class HybridRowQueue(
memManager: TaskMemoryManager,
tempDir: File,
numFields: Int,
@@ -209,9 +212,26 @@ case class HybridRowQueue(
override protected def isInMemoryQueue(queue: RowQueue): Boolean =
queue.isInstanceOf[InMemoryRowQueue]
+
+ // Shown in the per-consumer memory usage breakdown and TaskMemoryManager logs. The identity
+ // hash tells apart queues created with the same arguments.
+ override def toString: String =
+ s"HybridRowQueue(numFields=$numFields, lockFree=$lockFree)@" +
+ Integer.toHexString(System.identityHashCode(this))
}
object HybridRowQueue {
+ // Same parameter names as the constructor, like the apply of the former case class, so that
+ // calls with named arguments keep compiling.
+ def apply(
+ memManager: TaskMemoryManager,
+ tempDir: File,
+ numFields: Int,
+ serMgr: SerializerManager,
+ lockFree: Boolean = false): HybridRowQueue = {
+ new HybridRowQueue(memManager, tempDir, numFields, serMgr, lockFree)
+ }
+
def apply(taskMemoryMgr: TaskMemoryManager, file: File, fields: Int): HybridRowQueue = {
HybridRowQueue(taskMemoryMgr, file, fields, SparkEnv.get.serializerManager)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/RowQueueSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/RowQueueSuite.scala
index 3610b3581640..15ba1788c816 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/RowQueueSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/RowQueueSuite.scala
@@ -18,15 +18,18 @@
package org.apache.spark.sql.execution.python
import java.io.File
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
+import scala.concurrent.duration._
import org.apache.spark.{SparkConf, SparkFunSuite}
import org.apache.spark.internal.config._
-import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryConsumer, TestMemoryManager}
import org.apache.spark.security.{CryptoStreamUtils, EncryptionFunSuite}
import org.apache.spark.serializer.{JavaSerializer, SerializerManager}
import org.apache.spark.sql.catalyst.expressions.UnsafeRow
import org.apache.spark.unsafe.memory.{MemoryAllocator, MemoryBlock}
-import org.apache.spark.util.Utils
+import org.apache.spark.util.{ThreadUtils, Utils}
class RowQueueSuite extends SparkFunSuite with EncryptionFunSuite {
@@ -170,4 +173,124 @@ 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
+ // Check the self-triggered spill while queue1 still has spillable in-memory pages.
+ assert(queue1.spill(Long.MaxValue, queue1) === 0)
+ assert(queue1.getUsed === used)
+ assert(queue1.spill(Long.MaxValue, queue2) === used - smallPageSize)
+ drainAndClose(queue1, n)
+ queue2.close()
+ }
+
+ test("closing a hybrid queue does not deadlock with a concurrent spill of it") {
+ val closeFreeingPage = new CountDownLatch(1)
+ val spillHoldingTaskMemoryManager = new CountDownLatch(1)
+ @volatile var closingThread: Thread = null
+ val taskM = new TaskMemoryManager(new TestMemoryManager(smallPageConf), 0) {
+ override def freePage(page: MemoryBlock, consumer: MemoryConsumer): Unit = {
+ if (Thread.currentThread() eq closingThread) {
+ // Pause close() at its first page free until the spilling thread holds this
+ // TaskMemoryManager's monitor, so the two lock orders are forced to overlap.
+ closeFreeingPage.countDown()
+ assert(spillHoldingTaskMemoryManager.await(10, TimeUnit.SECONDS))
+ }
+ super.freePage(page, consumer)
+ }
+ }
+ val (queue1, queue2) = queuesWithSameArguments(taskM)
+ addRows(queue2, 150)
+ assert(queue2.numQueues() > 2)
+
+ val pool = ThreadUtils.newDaemonFixedThreadPool(2, "row-queue-close-spill")
+ try {
+ val closing = pool.submit[Unit](() => {
+ closingThread = Thread.currentThread()
+ queue2.close()
+ })
+ // Mimic TaskMemoryManager.acquireExecutionMemory, which holds its monitor while it asks
+ // another consumer to spill.
+ val spilling = pool.submit[Long](() => {
+ assert(closeFreeingPage.await(10, TimeUnit.SECONDS))
+ taskM.synchronized {
+ spillHoldingTaskMemoryManager.countDown()
+ queue2.spill(Long.MaxValue, queue1)
+ }
+ })
+ ThreadUtils.awaitResult(closing, 10.seconds)
+ // close() has already detached all of queue2's pages, so the spill has nothing to release.
+ assert(ThreadUtils.awaitResult(spilling, 10.seconds) === 0L)
+ } finally {
+ pool.shutdownNow()
+ }
+ assert(queue2.getUsed === 0)
+ queue1.close()
+ assert(taskM.cleanUpAllAllocatedMemory() === 0)
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org