qingfureal opened a new pull request, #58877:
URL: https://github.com/apache/spark/pull/58877
### What changes were proposed in this pull request?
Adds `closeOutputStream()` to the `Queue` trait, defaulting to a no-op,
overrides it in
`DiskRowQueue` to close and clear the writer, and calls it from
`HybridQueue.spill()` once a
queue has been drained.
`DiskRowQueue.remove()` needed a related change: it used `out != null` as
its "this is the
first read" signal and opened the input stream inside that branch. Now that
the output stream
can already be closed before the first read, that branch no longer fires, so
the check is
keyed off `in == null` instead. Opening the input stream is also deferred
until there is
actually something to read, so an exhausted queue no longer opens a file
just to return `null`.
Follow-up on SPARK-53481, which extracted the generic `HybridQueue` /
`Queue` classes.
### Why are the changes needed?
`HybridQueue.spill()` converts every in-memory queue but the last into a
disk-backed queue and
drains the rows into it:
```scala
val diskQueue = createDiskQueue()
var item = queue.remove()
while (item != null) {
diskQueue.add(item)
item = queue.remove()
}
released += getPageSize(queue)
queue.close()
diskQueue
```
That queue is complete — only the last queue is ever written to — but
nothing closes its
output stream. `DiskRowQueue` opens the writer eagerly in its constructor:
```scala
private var out = new DataOutputStream(serMgr.wrapForEncryption(
new BufferedOutputStream(new FileOutputStream(file.toString))))
```
and closes it only on the queue's *first read*, inside `remove()`. So after
N spills a task
holds N open `DataOutputStream` -> `BufferedOutputStream` ->
`FileOutputStream` chains, plus
the encryption wrapper when I/O encryption is enabled, each with its own 8
KB buffer, until
those queues are read. That is a file descriptor and a buffer per spilled
queue retained at
exactly the moment memory is scarce — spilling happens *because* the task is
under memory
pressure.
It is directly observable as unflushed data. In the added test, after a
spill the on-disk
bytes grow from 8184 to 12264 on the first read: 4080 bytes of rows were
still sitting in a
retained buffer rather than on disk.
### Does this PR introduce _any_ user-facing change?
No. This is an internal resource-lifetime fix in the Python UDF row queue;
queue semantics and
the rows returned are unchanged.
### How was this patch tested?
Two tests added to `RowQueueSuite`, both run with I/O encryption on and off
via
`encryptionTest`:
- `closeOutputStream flushes the queue and keeps it readable` — pins the new
contract on
`DiskRowQueue`: the rows reach disk, the call is idempotent, the queue
then rejects further
rows, and it is still fully readable afterwards.
- `spilling does not hold an output stream per spilled queue` — spills a
`HybridRowQueue`, then
asserts the on-disk size does not change on the first read. If a spilled
queue were still
holding its output stream, that first read is what would close and flush
it, and the file
would grow.
Each of the two behavioural hunks was confirmed to be load-bearing by
reverting it
individually:
| Reverted hunk | Result |
|---|---|
| `diskQueue.closeOutputStream()` in the spill loop | Both encryption
variants fail: `on-disk bytes grew from 8184 to 12264 on the first read` |
| `remove()` keyed off `in` instead of `out` | `NullPointerException: Cannot
invoke "java.io.DataInputStream.readInt()" because ... DiskRowQueue.in() is
null` — and this also breaks the **pre-existing** `hybrid queue` tests, 8
failures in total |
The second row is the reason the `remove()` change is part of this PR rather
than a separate
cleanup: closing the output stream early without it is a hard regression.
Also ran `RowQueueSuite`, `BatchEvalPythonExecSuite` and
`EvaluatePythonSuite` (35 tests) and
`scalastyle` on `sql/core`: all pass. `RowQueue.scala` holds the only
implementations of
`Queue` / `HybridQueue`, so the new default method changes no other queue.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)
--
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]
