viirya opened a new pull request, #59332:
URL: https://github.com/apache/spark/pull/59332
### What changes were proposed in this pull request?
In pipelined Python UDF mode (`spark.python.udf.pipelined.enabled`), the
task completion listener that stops the writer thread now waits until the
writer has actually exited.
The writer is now wrapped in `BasePythonRunner.PipelinedWriterTask`, which
counts down a latch in a `finally` around the whole writer. The listener
interrupts the writer with `cancel(true)` as before, and then waits on that
latch instead of calling `Future.get()`. If the listener runs before the pool
has started the writer, an `AtomicBoolean` claim makes sure the writer is
skipped, so the listener does not wait on a latch that would never be counted
down.
The wait has no time bound, the same as the `WriterThread.join()` that the
writer thread used before SPARK-44705 (added for SPARK-33277). While it waits,
it logs a warning every 10 seconds. If the listener thread is interrupted (for
example, by a task kill), it keeps waiting and restores the interrupt flag
afterwards.
### Why are the changes needed?
The existing listener does:
```scala
writerFuture.cancel(true)
try { writerFuture.get() } catch { case _: CancellationException | ... => }
```
Its comment says this waits for the writer to exit before the listeners that
run after it (registered earlier, LIFO) free the memory backing the input rows.
But once `cancel(true)` succeeds, `FutureTask.get()` throws
`CancellationException` right away and does not wait for the interrupted
runnable to return. So the writer can still be serializing a row, or adding it
to a `HybridRowQueue` page, when the queue's close listener frees that page.
That is a use-after-free, and with off-heap memory it can crash the executor. A
`HybridQueue.createNewQueue` racing with `close()` can also add a page that
nobody closes until `cleanUpAllAllocatedMemory`.
A task completes while the writer is still running when the output is not
fully consumed, for example with a `LIMIT` or when the task fails.
The wait cannot be bounded safely: if the listener gave up, the listeners
that run next would still free the memory, and throwing from the listener does
not stop them either. In practice the writer exits quickly after the interrupt:
`channel.write` fails with `ClosedByInterruptException`, and the writer loop
checks the interrupt flag between rows.
The pipelined mode was added in 4.3.0 by SPARK-56642, so branch-4.3 and
branch-4.x have the same issue.
### Does this PR introduce _any_ user-facing change?
No. This fixes a bug in an opt-in mode that is not released yet (4.3.0).
### How was this patch tested?
New tests in `BasePythonRunnerSuite`:
- Task completion stays blocked until a running writer that ignores
interrupts exits, and a listener registered earlier (which runs later) sees
that the writer has exited.
- The same holds when the thread running the listeners is interrupted, and
the interrupt flag is preserved.
- A writer stopped before it starts never runs.
Before the fix, a test with the old listener code showed `markTaskCompleted`
returning while the writer was still running. With the old cancel-then-`get()`
behavior put back, the first two new tests fail.
Also ran `pyspark.sql.tests.pandas.test_pipelined_udf`.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
This pull request and its description were written by Isaac.
--
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]
