viirya opened a new pull request, #59565: URL: https://github.com/apache/spark/pull/59565
### What changes were proposed in this pull request? `BasePythonRunner.ReaderInputStream` (the default, non-pipelined Python runner path) serializes the task input into a `DirectByteBufferOutputStream`. This PR hands the writer that buffer only through a new `CloseShieldOutputStream`, a `FilterOutputStream` that forwards `write(b, off, len)` and `write(b)` and whose `close()` only flushes. The buffer is still freed exactly once, by `ReaderInputStream`'s own `bufferStream.close()` after the last input is sent. `CloseShieldOutputStream` is the same class that #59332 (SPARK-60113) adds for the pipelined writer, so the two PRs stay independent; whichever merges second drops its copy on rebase. This PR does not touch `PipelinedWriterRunnable`. ### Why are the changes needed? The Arrow inputs (`PythonArrowInput`, `ColumnarArrowPythonInput`) write `dataOut` through `Channels.newChannel(dataOut)`, via `ArrowStreamWriter` and `WriteChannel`. That channel is an `AbstractInterruptibleChannel`: each chunk is written between `begin()` and `end()`, and closing the channel closes `dataOut`. When another thread interrupts the task thread inside such a write, e.g. `Task.kill(interruptThread = true)` for a losing speculative attempt or a job cancelled with interrupt-on-cancel, `Thread.interrupt()` runs the channel's interruptor on the interrupting thread. That closes `DataOutputStream` -> `DirectByteBufferOutputStream`, which frees the direct buffer while the task thread may still be copying into it (or copying the old buffer in `grow()`). The result can be a JVM crash, silent memory corruption, or an NPE / internal error instead of a clean task kill. With the shield, the interrupt still closes the channel and the task thread stops with `ClosedByInterruptException` right after the current chunk, but the buffer is not freed under it. This affects Arrow-based Python execution (pandas/Arrow UDFs, Arrow-optimized Python UDFs, `mapInPandas`/`mapInArrow`, `applyInPandas`, ...) since SPARK-44705 made the runner single-threaded (4.0.0). The pickled `BatchEvalPython` path writes `dataOut` directly and is not affected. Other streams in core and sql/core that are written through `Channels.newChannel` (`ByteArrayOutputStream`, `ChunkedByteBufferOutputStream`, `ByteBufferOutputStream`) do not free native memory on close, and `PythonPlannerRunner`'s `DirectByteBufferOutputStream` is not written through a channel. Thanks to @dongjoon-hyun, who found this in the review of #59332 (https://github.com/apache/spark/pull/59332#discussion_r4264575990). ### Does this PR introduce _any_ user-facing change? Yes, as a bug fix: killing a task that runs an Arrow-based Python UDF with thread interruption no longer risks freeing the input buffer while it is being written. There is no API or configuration change. ### How was this patch tested? New tests in `BasePythonRunnerSuite`: - `CloseShieldOutputStream` forwards writes and does not close the underlying stream. - A test through the real `ReaderInputStream`, over a local socket, with a writer that writes through `Channels.newChannel(dataOut)` like the Arrow inputs. It interrupts the writing thread inside the channel write, then checks that the blocked write into the buffer still succeeds and that the thread stops with `ClosedByInterruptException`. Without the fix, the write fails with `Cannot call methods on a closed DirectByteBufferOutputStream`, because the interrupting thread already freed the buffer. Both new tests fail when the fix is reverted. Also ran `BasePythonRunnerSuite`, `DirectByteBufferOutputStreamSuite`, `PythonRDDSuite`, and the PySpark tests `pyspark.sql.tests.arrow.test_arrow_python_udf`, `pyspark.sql.tests.pandas.test_pandas_udf_scalar`, `pyspark.sql.tests.arrow.test_arrow_map` and `pyspark.sql.tests.pandas.test_pandas_grouped_map`, on JDK 17. Backport plan: branch-4.3, branch-4.2, branch-4.1 and branch-4.0 are affected and will get backports after this is merged. ### 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]
