[SPARK-58241][PYTHON] Fix the columnar Arrow Python UDF runner writing the input schema in the legacy protocol position

[viirya]

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 count followed by UDF entries. The stray schema string desyncs the entire parse: its length prefix is misread as the UDF count, subsequent bytes are misread as field lengths, and the worker ends up trying to read a garbage-sized value past the end of the init message. This PR aligns the columnar runner with the row-based one: schema via evalConf, writeUDF writes only the UDF list.

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 EOFError naming the short read instead of a silent, undiagnosable hang.

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 with default configurations. Note ArrowEvalPythonExec.doExecute routes through the columnar evaluator whenever child.supportsColumnar, regardless of spark.sql.execution.arrow.pythonUDF.columnarInput.enabled, so the config is not a workaround.

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.

[dongjoon-hyun]

+1, LGTM. Thank you, @viirya .

[uros-b]

Thank you @viirya and @dongjoon-hyun!

[Yicong-Huang]

I think this is a legit fix and thanks for catching that! Left one nit inline.

[Yicong-Huang]

The comment says unselected columns "ride through the pass-through recombination," but the assertion only checks the UDF output ([Row(2), Row(3)]) -- the unselected s column is projected away by the select and never observed. Not vacuous (the pass-through row-count assert would still fire on a size mismatch), but to actually observe the pass-through the comment names, consider also selecting s and asserting its value.

[viirya]

Good point -- the comment claimed more than the assertion observed. The test now selects s alongside the UDF result and asserts [Row("a", 2), Row("b", 3)], so the pass-through values (and their row alignment with the UDF output) are observed directly rather than implied by the row-count check.

[viirya]

Merge Summary:

• merged into master e6d66ab
• merged into branch-4.x a84fcaa

Posted by merge_spark_pr.py

[viirya]

Thanks @dongjoon-hyun @uros-b @Yicong-Huang

[viirya]

I created a backport of this to branch-4.2 #57419.