RamonZhou opened a new pull request, #58510:
URL: https://github.com/apache/spark/pull/58510
### What changes were proposed in this pull request?
Follow-up to SPARK-58752, which let a session set environment variables for
its Python UDF
workers via `spark.pythonWorkerEnv.*`, but installed them for the regular
scalar Python UDF
only. This extends the environment to every Python worker a session launches
for a Python
function that session supplied.
`PythonWorkerEnvironment.appliesTo(evalType)` is removed. While only three
evaluation types
were covered, an allowlist was a useful guard; now that every runner
installs the environment
it would be a second source of truth, and a new evaluation type routed
through an existing
runner would silently receive nothing. The scope is now the set of runners
that install it, so
such a type is covered by construction.
Runners that gain the merge:
| Runner | Serves |
| --- | --- |
| `ArrowPythonUDTFRunner` | Arrow UDTFs |
| `CoGroupedArrowPythonRunner` | cogrouped map, pandas and Arrow |
| `TransformWithStateInPySparkPythonRunner` | the `transformWithState`
evaluation types |
| `ApplyInPandasWithStatePythonRunner` | `applyInPandasWithState` |
| `PythonPlannerRunner` | UDTF `analyze`, data-source planning,
streaming-sink commit |
| `PythonStreamingSourceRunner` | Python streaming data sources |
| `StreamingPythonRunner` | `foreachBatch` |
`BasePythonUDFRunner` and `BaseArrowPythonRunner` already merged, gated on
`appliesTo`; with the
gate gone they cover the pandas, Arrow, iterator, grouped-map,
grouped-aggregate, window,
`mapInPandas` and `mapInArrow` families that already reached them.
`StreamingPythonRunner` lives in `core`, which cannot see `sql/core`, and
`foreachBatch` uses it
directly, so it takes an optional `sessionEnvVars` parameter that callers
with a session supply.
Rewriting the caller's `PythonFunction` instead was rejected: it would
silently no-op for any
implementation other than `SimplePythonFunction`. While there, `envVars`
becomes a copy rather
than an alias of `func.envVars`, which the runner mutates -- a function
handed to two runners
previously accumulated the other's variables.
### Why are the changes needed?
A session's environment should not depend on which serialization mode or
Python function family
a user picked. Before this change a plain `udf` received the environment
while the same logic
written as a `pandas_udf`, a UDTF or `applyInPandas` did not, which is a
difference users cannot
predict from anything in the API.
### Does this PR introduce _any_ user-facing change?
Yes. `spark.pythonWorkerEnv.*` now reaches the families listed above, where
it previously had no
effect. `docs/configuration.md` is updated accordingly.
Two consequences worth calling out:
- A worker that outlives the query that launched it keeps the values it
started with. The
streaming paths launch one worker per query rather than one per batch, so
a change made while a
streaming query is running reaches that query's worker only on restart.
Documented.
- Validation is part of installing the environment, so a malformed
environment now fails the
newly covered paths, which previously ignored it. Only reachable by
writing straight to
`SQLConf` and bypassing `RuntimeConfig.set`, which validates on write.
Pinned by a test.
A Spark Connect listener added through `addListener` runs its callbacks in
the client process
rather than in a worker Spark launched, so no session environment applies to
it; the docs now say
so. The server-side `PythonStreamingQueryListener` does launch a worker and
does receive the
environment, but it is reached only through the `add_listener` command,
which PySpark no longer
uses, so it has no PySpark-level test.
### How was this patch tested?
`PythonWorkerEnvironmentSuite`, 46 tests. The two tests that asserted the
old narrow scope are
rewritten to assert the new contract, including one pinning that an invalid
environment now fails
runners that previously ignored it. New helpers cover the Arrow UDTF and
cogrouped runners
directly, since each declares its own `envVars` and would not be caught by
testing a sibling.
`test_python_worker_env`, 28 tests, run on both classic and Connect through
the shared mixin:
`mapInPandas`, `mapInArrow`, scalar and iterator pandas UDFs,
`applyInPandas`, cogrouped
`applyInPandas`, grouped-aggregate and window pandas UDFs, and row and Arrow
UDTFs. The
data-source test reads the variable in `schema()`, which runs in the
planning worker, so it fails
if only the execution-side workers received the environment.
`test_parity_python_worker_env` adds a `foreachBatch` test that records what
the worker observed
in a file, using a file source with `availableNow` so the query is bounded
rather than timed.
Run locally: `PythonWorkerEnvironmentSuite` 46/46, classic
`test_python_worker_env` 30/30,
Connect parity 30/30, `dev/lint-scala` and both `dev/lint-python` ruff
passes clean.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5
--
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]
