This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new a84fcaae64d0 [SPARK-58241][PYTHON] Fix the columnar Arrow Python UDF
runner writing the input schema in the legacy protocol position
a84fcaae64d0 is described below
commit a84fcaae64d031cd68341805eff2d9193df1c79f
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Tue Jul 21 17:55:11 2026 -0700
[SPARK-58241][PYTHON] Fix the columnar Arrow Python UDF runner writing the
input schema in the legacy protocol position
### What changes were proposed in this pull request?
Two fixes to the Python worker init-message protocol:
1. **Protocol fix**: `ColumnarArrowPythonWithNamedArgumentRunner` (the
columnar-input Arrow Python UDF runner) still wrote `schema.json` as a UTF
string at the head of the UDF command section for `SQL_ARROW_BATCHED_UDF`. The
message-based worker protocol moved that schema into `evalConf` under the
`"input_type"` key -- which the row-based `ArrowPythonWithNamedArgumentRunner`
already does -- and the worker's `WorkerInitInfo.from_stream` therefore parses
the command section as a UDF cou [...]
2. **Fail-fast hardening**: the worker blocked *forever* instead of failing
because `SparkSocketMessageReceiver._do_get_init_message` materializes the
entire init message but never marked the resulting `ZeroCopyByteStream`
finished, so an over-read waits on a condition for a chunk that can never
arrive (the worker is single-threaded at that point). The receiver now calls
`finish()` on the fully-materialized stream, so any future writer/reader
protocol mismatch surfaces as an immediate [...]
### Why are the changes needed?
Any Python UDF whose input plan produces columnar Arrow batches deadlocks:
the JVM task thread waits for worker output, the worker waits for init-message
bytes that were never sent, forever. This was latent since the columnar-input
runner was added because no in-tree source produced Arrow-backed
`ColumnarBatch` input for it; the Arrow-cached scan (SPARK-57268) is the first,
so a simple `df.cache()` (with the Arrow cache serializer) followed by any
Arrow-optimized Python UDF now hangs [...]
Diagnosed by instrumenting `WorkerInitInfo.from_stream`: the parse is
byte-exact through the conf sections and desyncs precisely at the UDF command
section, matching the extra `schema.json` write; the worker's `faulthandler`
stack shows it blocked in `ZeroCopyByteStream._try_read_bytes` under
`UDFInfo.from_stream`.
### Does this PR introduce _any_ user-facing change?
No. It fixes a hang; the wire format for the columnar runner now matches
what the worker has expected all along.
### How was this patch tested?
- New e2e suite `pyspark.sql.tests.arrow.test_arrow_python_udf_cached` (own
module because `spark.sql.cache.serializer` is a static conf latched into a
JVM-wide singleton): Arrow-optimized Python UDFs over an Arrow-cached relation
-- the exact deadlock scenario -- including an unselected pass-through column.
Hung forever before the fix; passes in seconds with it.
- New unit tests in `pyspark.tests.test_spark_message_receiver` for the
socket receiver: declared content reads back, and an over-read past the
declared init-message length raises `EOFError` instead of blocking.
- Existing `pyspark.tests.test_zero_copy_byte_stream`,
`pyspark.tests.test_spark_message_receiver`, and
`pyspark.sql.tests.arrow.test_arrow_python_udf` all pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
This pull request and its description were written by Claude Code.
Closes #57403 from viirya/arrow-columnar-udf-init-protocol.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
(cherry picked from commit e6d66abda0ac420f9f8448b6b10e6efbb3c67fbd)
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
dev/sparktestsupport/modules.py | 1 +
.../socket/spark_socket_message_receiver.py | 7 +-
.../tests/arrow/test_arrow_python_udf_cached.py | 97 ++++++++++++++++++++++
.../pyspark/tests/test_spark_message_receiver.py | 30 +++++++
.../python/ColumnarArrowPythonRunner.scala | 16 +++-
5 files changed, 148 insertions(+), 3 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index fdd2f2fa598f..3edebcee6537 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -593,6 +593,7 @@ pyspark_sql = Module(
"pyspark.sql.tests.arrow.test_arrow_cogrouped_map_misc",
"pyspark.sql.tests.arrow.test_arrow_grouped_map",
"pyspark.sql.tests.arrow.test_arrow_python_udf",
+ "pyspark.sql.tests.arrow.test_arrow_python_udf_cached",
"pyspark.sql.tests.arrow.test_arrow_udf",
"pyspark.sql.tests.arrow.test_arrow_udf_grouped_agg",
"pyspark.sql.tests.arrow.test_arrow_udf_scalar",
diff --git a/python/pyspark/messages/socket/spark_socket_message_receiver.py
b/python/pyspark/messages/socket/spark_socket_message_receiver.py
index de3db2fe6611..7099f0207c1f 100644
--- a/python/pyspark/messages/socket/spark_socket_message_receiver.py
+++ b/python/pyspark/messages/socket/spark_socket_message_receiver.py
@@ -43,7 +43,12 @@ class SparkSocketMessageReceiver(SparkMessageReceiver):
message_length = read_int(self._infile)
message_content = self._infile.read(message_length)
- return ZeroCopyByteStream(memoryview(message_content))
+ # The init message is fully materialized, so mark the stream finished:
a parse that
+ # runs past the declared message length (e.g. a writer/reader protocol
mismatch) must
+ # fail fast with EOFError instead of blocking forever on a chunk that
will never come.
+ stream = ZeroCopyByteStream(memoryview(message_content))
+ stream.finish()
+ return stream
def _do_get_data_stream(self) -> BinaryIO:
# For socket communication, we just pass along the underlying socket
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
new file mode 100644
index 000000000000..962302730e69
--- /dev/null
+++ b/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py
@@ -0,0 +1,97 @@
+#
+# 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.
+#
+
+import unittest
+
+from pyspark.sql.functions import col, udf
+from pyspark.sql import Row
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import (
+ assertDataFrameEqual,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+
+
[email protected](
+ not have_pandas or not have_pyarrow,
+ pandas_requirement_message or pyarrow_requirement_message, # type:
ignore[arg-type]
+)
+class ArrowPythonUDFCachedInputTests(ReusedSQLTestCase):
+ """
+ Arrow-optimized Python UDFs reading columnar input from an Arrow-cached
relation.
+
+ The Arrow-cached scan is a columnar Arrow-backed source, so it drives the
columnar Arrow
+ Python runner (ColumnarArrowPythonWithNamedArgumentRunner) end to end --
the path an
+ Arrow-backed DSv2 source would also take. This pins the init-message
protocol between that
+ runner and the worker: a mismatched section there desyncs the worker's
init parse and hangs
+ the task (SPARK-58241).
+
+ This suite runs in its own module because spark.sql.cache.serializer is a
static conf latched
+ into a JVM-wide singleton on first cache materialization.
+ """
+
+ @classmethod
+ def conf(cls):
+ conf = super().conf()
+ # Static conf: must be set before the session is created.
+ conf.set(
+ "spark.sql.cache.serializer",
+
"org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer",
+ )
+ return conf
+
+ @classmethod
+ def setUpClass(cls):
+ super().setUpClass()
+ cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled",
"true")
+
+ @classmethod
+ def tearDownClass(cls):
+ try:
+ cls.spark.conf.unset("spark.sql.execution.pythonUDF.arrow.enabled")
+ finally:
+ super().tearDownClass()
+
+ def test_udf_on_cached_strings(self):
+ df = self.spark.createDataFrame(
+ [("hello",), (None,), ("world",)], schema="v string"
+ ).cache()
+ try:
+ result = df.select(udf(lambda x: x.upper() if x is not None else
None)(col("v")))
+ assertDataFrameEqual(result, [Row("HELLO"), Row(None),
Row("WORLD")])
+ finally:
+ df.unpersist()
+
+ def test_udf_on_cached_numeric_column_alongside_unselected_string(self):
+ # Only the UDF's input columns are serialized to the worker; the
column that is not a
+ # UDF input rides through the pass-through recombination, so select it
alongside the
+ # UDF result to observe that its values line up with the corresponding
rows.
+ df = self.spark.createDataFrame([(1, "a"), (2, "b")], schema="i long,
s string").cache()
+ try:
+ result = df.select(col("s"), udf(lambda x: x + 1,
"long")(col("i")))
+ assertDataFrameEqual(result, [Row("a", 2), Row("b", 3)])
+ finally:
+ df.unpersist()
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/tests/test_spark_message_receiver.py
b/python/pyspark/tests/test_spark_message_receiver.py
index 398bae922e5b..944a5084546b 100644
--- a/python/pyspark/tests/test_spark_message_receiver.py
+++ b/python/pyspark/tests/test_spark_message_receiver.py
@@ -18,8 +18,10 @@ import io
import unittest
from typing import BinaryIO
+from pyspark.messages.socket.spark_socket_message_receiver import
SparkSocketMessageReceiver
from pyspark.messages.spark_message_receiver import SparkMessageReceiver
from pyspark.messages.zero_copy_byte_stream import ZeroCopyByteStream
+from pyspark.serializers import SpecialLengths
class StubMessageReceiver(SparkMessageReceiver):
@@ -73,6 +75,34 @@ class SparkMessageReceiverTests(unittest.TestCase):
getattr(receiver, invalid_call)()
+class SparkSocketMessageReceiverTests(unittest.TestCase):
+ """Tests for the socket-framed receiver."""
+
+ def _framed_init(self, content: bytes) -> io.BytesIO:
+ frame = io.BytesIO()
+ frame.write(SpecialLengths.START_OF_INIT_MESSAGE.to_bytes(4, "big",
signed=True))
+ frame.write(len(content).to_bytes(4, "big", signed=True))
+ frame.write(content)
+ frame.seek(0)
+ return frame
+
+ def test_init_message_reads_declared_content(self):
+ receiver = SparkSocketMessageReceiver(self._framed_init(b"abcdef"))
+ stream = receiver.get_init_message()
+ self.assertEqual(bytes(stream.read(6)), b"abcdef")
+
+ def test_init_message_over_read_raises_eof_instead_of_blocking(self):
+ # The init message is fully materialized, so the stream must be marked
finished:
+ # a parse that runs past the declared length (e.g. a writer/reader
protocol
+ # mismatch, see SPARK-58241) must fail fast with EOFError -- with an
unfinished
+ # stream it would block forever waiting for a chunk that never arrives.
+ receiver = SparkSocketMessageReceiver(self._framed_init(b"abcdef"))
+ stream = receiver.get_init_message()
+ stream.read(4)
+ with self.assertRaises(EOFError):
+ stream.read(4)
+
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
index dcc915d5021a..41a563a7302f 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala
@@ -51,10 +51,22 @@ private[python] class
ColumnarArrowPythonWithNamedArgumentRunner(
override protected def runnerConf: Map[String, String] = super.runnerConf ++
pythonRunnerConf
- override protected def writeUDF(dataOut: DataOutputStream): Unit = {
+ // The input schema travels in evalConf (key "input_type"), exactly like
+ // ArrowPythonWithNamedArgumentRunner. It must NOT be written into the UDF
command section:
+ // the worker's WorkerInitInfo.from_stream reads that section as a UDF count
followed by UDF
+ // entries, so an extra UTF string there desyncs the whole init-message
parse -- the worker
+ // then blocks waiting for bytes past the end of the init message and the
task hangs forever.
+ override protected def evalConf: Map[String, String] = {
if (evalType == PythonEvalType.SQL_ARROW_BATCHED_UDF) {
- PythonWorkerUtils.writeUTF(schema.json, dataOut)
+ super.evalConf ++ Map(
+ "input_type" -> schema.json
+ )
+ } else {
+ super.evalConf
}
+ }
+
+ override protected def writeUDF(dataOut: DataOutputStream): Unit = {
PythonUDFRunner.writeUDFs(dataOut, funcs, argMetas)
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
