[SPARK-57931][CORE] Restore worker channel blocking mode after pipelined Python UDF execution

[viirya]

What changes were proposed in this pull request?

The pipelined Python UDF path (SPARK-56642) switches a borrowed worker's SocketChannel from non-blocking to blocking mode in createPipelinedDataIn and never restores it. With worker reuse enabled (the default), the worker is returned to the idle pool with its channel still in blocking mode.

This changes PythonWorkerFactory.create() to normalize a reused daemon worker's channel back to non-blocking before calling refresh(), so a pooled worker is always handed to the next task in the same (non-blocking) mode as a freshly created one — restoring the invariant that a daemon worker taken from the pool is non-blocking.

Why are the changes needed?

PythonWorker.refresh() only opens a selector when the channel is non-blocking. A pooled worker left in blocking mode therefore comes back with a null selector / selectionKey. Code on the non-pipelined (single-threaded NIO selector) path dereferences worker.selector / worker.selectionKey, so a worker in this corrupted state would throw a NullPointerException there.

In the current OSS code this does not surface as an end-to-end failure, because the worker-factory cache key (PythonWorkersKey) includes the worker envVars, and the pipelined path adds SPARK_PIPELINED_UDF=1 to envVars before requesting a worker (BasePythonRunner.compute). Pipelined and non-pipelined tasks therefore draw from separate idle pools: a worker left in blocking mode only returns to the pipelined pool, and the next borrower from that pool is again a pipelined task, which unconditionally re-sets the channel to blocking and does not use the selector. So the corrupted state is currently masked by pool isolation.

That masking is fragile. It relies on the two pools staying disjoint via envVars; it does not fix the broken invariant that a pooled daemon worker is non-blocking. Any worker-management layer that pools or reuses workers across that boundary — e.g. reusing a warmed worker regardless of whether the previous task was pipelined — will hand a blocking-mode worker to selector-path code and hit the NullPointerException. Fixing it at the pool boundary (create()) restores the invariant unconditionally and is robust to how workers are pooled.

The fix is applied in create() rather than in the pipelined path's task-completion listener because the worker is released back to the pool from the reader iterator when it reaches END_OF_STREAM, which runs before the task-completion listener; a restore in the listener would therefore run after the worker is already back in the pool. Normalizing at the single pool exit point (create()) is correct regardless of that ordering.

Does this PR introduce any user-facing change?

No. This hardens an invariant in an opt-in feature (spark.python.udf.pipelined.enabled) that has not been released yet; the pool-isolation behavior above means it is not an observable failure in OSS today.

How was this patch tested?

New unit test in PythonWorkerFactorySuite that constructs a PythonWorker over a loopback channel, puts it in the blocking state the pipelined path leaves behind (where refresh() opens no selector), and asserts refreshNonBlocking() — the method create() uses when reusing a pooled worker — normalizes the channel back to non-blocking and re-opens the selector. The test uses a mock channel rather than a real daemon worker so it runs in the core module's test environment (which has no pyspark on PYTHONPATH).

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code

[uros-b]

cc @HyukjinKwon for PythonWorkerFactory and PythonRunner

[viirya]

Merge Summary:

• merged into master 20d8cce
• merged into branch-4.x 72d854a

Posted by merge_spark_pr.py

[viirya]

Thanks @uros-b @HyukjinKwon