viirya opened a new pull request, #57403: URL: https://github.com/apache/spark/pull/57403
### 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. -- 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]
