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/20 02:08:30 UTC
(spark) branch branch-4.3 updated: [SPARK-60305][PYTHON] Keep the Python runner's input buffer alive when an Arrow write is interrupted
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new d9878fb26064 [SPARK-60305][PYTHON] Keep the Python runner's input buffer alive when an Arrow write is interrupted
d9878fb26064 is described below
commit d9878fb26064e04dd8abcf9647324c2080a0cbfa
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Mon Oct 19 19:08:01 2026 -0700
[SPARK-60305][PYTHON] Keep the Python runner's input buffer alive when an Arrow write is interrupted
### 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 now freed only on the thread that writes into it: after the last input is sent, as before, or, new in this PR, when an exception escape [...]
#59332 (SPARK-60113) adds the same `CloseShieldOutputStream` file for the pipelined writer, so the two PRs stay independent: whichever merges second has no diff for it. If #59332 merges first, this PR's commit no longer adds the file, so its backports to branch-4.2, 4.1 and 4.0, which #59332 doesn't reach, need to bring it along. 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-o [...]
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 [...]
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 `CloseShieldOutputStreamSuite`:
- `CloseShieldOutputStream` forwards bulk writes as bulk writes and does not close the underlying stream.
- A channel write with the interrupt flag already set closes the channel on the writing thread, and the shielded buffer stays open.
- 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, that the thread stops with `ClosedByInterruptException`, and that the buffer is freed afterwards. A later read must then fail with `Cannot call methods on a closed DirectByteBufferOutputStream`, which checks [...]
Each test fails when the part of the change it covers is reverted. The first two tests port unchanged to all branches; the third only needs the `ReaderInputStream` constructor call adjusted. 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_grou [...]
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.
Closes #59565 from viirya/SPARK-60305-close-shield.
Authored-by: Liang-Chi Hsieh <vi...@gmail.com>
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
(cherry picked from commit 058c570f99e0097cb9d2c08155f8f46549e30a5d)
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
.../spark/api/python/CloseShieldOutputStream.scala | 41 ++++++
.../org/apache/spark/api/python/PythonRunner.scala | 30 ++++-
.../api/python/CloseShieldOutputStreamSuite.scala | 147 +++++++++++++++++++++
3 files changed, 214 insertions(+), 4 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala b/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala
new file mode 100644
index 000000000000..33ccef7ab0b1
--- /dev/null
+++ b/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.python
+
+import java.io.{FilterOutputStream, OutputStream}
+
+/**
+ * An output stream that forwards writes to `underlying` but does not close it on [[close()]].
+ *
+ * Python runners serialize their input into a `DirectByteBufferOutputStream`, whose `close()`
+ * frees its native buffer. Arrow writers write that stream through `Channels.newChannel(...)`,
+ * an interruptible channel: when another thread interrupts the writing thread inside a channel
+ * write (e.g. a task kill with `interruptThread = true`), `Thread.interrupt()` closes the
+ * channel, and so the wrapped stream, on the interrupting thread while the writing thread may
+ * still be copying into the buffer. Handing the writer this stream keeps the buffer alive; the
+ * writing thread still stops with a `ClosedByInterruptException`, and the buffer is freed only
+ * by its owner.
+ */
+private[python] class CloseShieldOutputStream(underlying: OutputStream)
+ extends FilterOutputStream(underlying) {
+
+ // FilterOutputStream writes an array one byte at a time.
+ override def write(b: Array[Byte], off: Int, len: Int): Unit = out.write(b, off, len)
+
+ override def close(): Unit = flush()
+}
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 5026d65cf37c..1926f7f5e47a 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -964,16 +964,22 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
private[this] var writerIfbhThreadLocalValue: Object = null
private[this] val temp = new Array[Byte](1)
private[this] val bufferStream = new DirectByteBufferOutputStream()
+ // The writer gets `bufferStream` only through this stream. Arrow writers write through an
+ // interruptible channel, so a task kill that interrupts this thread closes the stream they
+ // write to on the killing thread; that must not free `bufferStream` under the writer. It is
+ // freed only on this thread, by `writeAdditionalInputToPythonWorker()`, once all input is
+ // sent or when writing fails. See `CloseShieldOutputStream`.
+ private[this] val writerStream = new CloseShieldOutputStream(bufferStream)
/**
* Buffers data to be written to the Python worker until the socket is
* available for write.
* A best-effort attempt is made to not grow the buffer beyond "spark.buffer.size". See
- * `writeAdditionalInputToPythonWorker()` for details.
+ * `writeBufferedInputToPythonWorker()` for details.
*/
private[this] var buffer: ByteBuffer = _
private[this] var hasInput = true
- writer.open(new DataOutputStream(bufferStream))
+ writer.open(new DataOutputStream(writerStream))
buffer = bufferStream.toByteBuffer
override def read(): Int = {
@@ -1133,11 +1139,27 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
}
}
+ private def writeAdditionalInputToPythonWorker(): Unit = {
+ try {
+ writeBufferedInputToPythonWorker()
+ } catch {
+ case t: Throwable =>
+ // Writing failed or the task was killed mid-write, so the close at the end of
+ // `writeBufferedInputToPythonWorker()` is never reached. Free the buffer here, on the
+ // thread that writes into it, instead of leaving it to GC. `buffer` may be null or
+ // still point into that memory, so replace it with an empty one: a later call then
+ // fails fast in `bufferStream` instead of touching the freed memory.
+ buffer = ByteBuffer.allocate(0)
+ bufferStream.close()
+ throw t
+ }
+ }
+
/**
* Reads input data from `writer.inputIterator` into `buffer` and writes the buffer to the
* Python worker if the socket is available for writing.
*/
- private def writeAdditionalInputToPythonWorker(): Unit = {
+ private def writeBufferedInputToPythonWorker(): Unit = {
var acceptsInput = true
while (acceptsInput && (hasInput || buffer.hasRemaining)) {
if (!buffer.hasRemaining && hasInput) {
@@ -1146,7 +1168,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
// Set the `buffer` to null to make it eligible for GC
buffer = null
- val dataOut = new DataOutputStream(bufferStream)
+ val dataOut = new DataOutputStream(writerStream)
// Try not to grow the buffer much beyond `bufferSize`. This is inevitable for large
// input rows.
while (bufferStream.size() < bufferSize && hasInput && !shouldFlush()) {
diff --git a/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala b/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala
new file mode 100644
index 000000000000..e34d34422557
--- /dev/null
+++ b/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala
@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.api.python
+
+import java.io.{DataOutputStream, FilterOutputStream}
+import java.net.{InetAddress, InetSocketAddress}
+import java.nio.ByteBuffer
+import java.nio.channels.{Channels, ClosedByInterruptException, ServerSocketChannel, SocketChannel}
+import java.util.{ArrayList => JArrayList, HashMap => JHashMap}
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+import java.util.concurrent.atomic.AtomicBoolean
+
+import org.apache.spark.{LocalSparkContext, SparkContext, SparkEnv, SparkException, SparkFunSuite}
+import org.apache.spark.util.DirectByteBufferOutputStream
+
+class CloseShieldOutputStreamSuite extends SparkFunSuite with LocalSparkContext {
+
+ test("SPARK-60305: forwards bulk writes and does not close the underlying stream") {
+ var byteWrites = 0
+ val bufferStream = new DirectByteBufferOutputStream() {
+ override def write(b: Int): Unit = {
+ byteWrites += 1
+ super.write(b)
+ }
+ }
+ val shielded = new CloseShieldOutputStream(bufferStream)
+ shielded.write(Array[Byte](1, 2, 3), 1, 2)
+ assert(byteWrites === 0)
+ shielded.write(4)
+ assert(byteWrites === 1)
+ shielded.close()
+ assert(bufferStream.size() === 3)
+ val bytes = new Array[Byte](3)
+ bufferStream.toByteBuffer.get(bytes)
+ assert(bytes.toSeq === Seq[Byte](2, 3, 4))
+ bufferStream.close()
+ }
+
+ test("SPARK-60305: an interrupted channel write does not close the shielded buffer") {
+ val bufferStream = new DirectByteBufferOutputStream()
+ val channel =
+ Channels.newChannel(new DataOutputStream(new CloseShieldOutputStream(bufferStream)))
+ // With the interrupt flag already set, the channel closes itself, and the stream it writes
+ // to, on this thread before it writes.
+ Thread.currentThread().interrupt()
+ try {
+ intercept[ClosedByInterruptException](channel.write(ByteBuffer.wrap(new Array[Byte](16))))
+ } finally {
+ Thread.interrupted()
+ }
+ assert(bufferStream.size() === 16)
+ bufferStream.close()
+ }
+
+ test("SPARK-60305: interrupting the task thread inside a channel write keeps the " +
+ "ReaderInputStream buffer alive") {
+ sc = new SparkContext("local", "test")
+ val func = SimplePythonFunction(Seq.empty[Byte], new JHashMap[String, String](),
+ new JArrayList[String](), "python3", "3", new JArrayList(), null)
+ val runner = new PythonRunner(Seq(ChainedPythonFunctions(Seq(func))), None)
+ val inWrite = new CountDownLatch(1)
+ val resume = new AtomicBoolean(false)
+ @volatile var writeError: Throwable = null
+
+ val server = ServerSocketChannel.open()
+ server.bind(new InetSocketAddress(InetAddress.getLoopbackAddress, 0))
+ val client = SocketChannel.open(server.getLocalAddress)
+ val accepted = server.accept()
+ client.configureBlocking(false)
+ val worker = PythonWorker(client).refresh()
+ try {
+ val writer = new runner.Writer(SparkEnv.get, worker, Iterator.empty, 0, null) {
+ // An interrupt closes the stream the writer writes to; with the shield, this only
+ // flushes.
+ override def open(dataOut: DataOutputStream): Unit = dataOut.close()
+ override protected def writeCommand(dataOut: DataOutputStream): Unit = {}
+ // Writes the way the Arrow inputs do: through `Channels.newChannel(dataOut)`, which the
+ // JVM closes, together with `dataOut`, when the writing thread is interrupted. The
+ // write blocks until the test has interrupted the thread, then forwards to `dataOut`.
+ override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
+ val channel = Channels.newChannel(new FilterOutputStream(dataOut) {
+ override def write(b: Array[Byte], off: Int, len: Int): Unit = {
+ inWrite.countDown()
+ while (!resume.get()) Thread.onSpinWait()
+ try {
+ out.write(b, off, len)
+ } catch {
+ case t: Throwable =>
+ writeError = t
+ throw t
+ }
+ }
+ })
+ channel.write(ByteBuffer.wrap(new Array[Byte](16)))
+ false
+ }
+ }
+ val input = new runner.ReaderInputStream(worker, writer, None, 0, false, null)
+ @volatile var error: Throwable = null
+ val taskThread = new Thread(() => {
+ try {
+ input.read(new Array[Byte](1))
+ } catch {
+ case t: Throwable => error = t
+ }
+ })
+ taskThread.start()
+ assert(inWrite.await(1, TimeUnit.MINUTES), s"task thread failed before the write: $error")
+ // Like Task.kill(interruptThread = true). This closes the channel and `dataOut` on this
+ // thread while the task thread is still inside the write.
+ taskThread.interrupt()
+ resume.set(true)
+ taskThread.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!taskThread.isAlive)
+ // The write into the buffer still succeeds: the buffer was not freed under the writer.
+ assert(writeError === null)
+ assert(error.isInstanceOf[ClosedByInterruptException])
+ // The failed write frees the buffer on the task thread instead of leaving it to GC, and
+ // a later read fails instead of reusing the freed buffer. Closing the peer first makes a
+ // regression return -1 rather than wait for output.
+ accepted.close()
+ val e = intercept[SparkException](input.read(new Array[Byte](1)))
+ assert(e.getMessage.contains("closed DirectByteBufferOutputStream"))
+ } finally {
+ // Never leave the writer spinning, even if the test failed before releasing it.
+ resume.set(true)
+ worker.stop()
+ accepted.close()
+ server.close()
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org