RamonZhou opened a new pull request, #58264:
URL: https://github.com/apache/spark/pull/58264
### What changes were proposed in this pull request?
Python UDFs on Spark Connect run in worker processes whose environment is
always empty: the Connect
planner builds every Python function with no environment variables. This
lets a session carry an
environment for its Python workers through session configurations under a
reserved prefix, one
configuration per variable.
spark.conf.set("spark.pythonWorkerEnv.MY_SETTING", "abc")
@udf("string")
def f(_):
import os
return os.environ["MY_SETTING"] # "abc"; previously a KeyError
Changes, all confined to `sql/connect/server` apart from four new error
sub-conditions:
- `PythonWorkerEnvironment` (new): reads the environment from the session's
configurations,
validates it, fingerprints it for the plan cache, and hands out a fresh
mutable copy for a single
Python function.
- `SparkConnectPlanner.transformPythonFunction` populates
`SimplePythonFunction.envVars` from it.
Every Python function family built at that site therefore receives the
environment. Covered by
tests: scalar Python UDFs, Arrow-batched UDFs, scalar pandas UDFs and
their iterator variant,
`mapInPandas` and `mapInArrow`. Reaching the same site by construction but
not separately tested
here: grouped-map, cogrouped-map, stateful pandas functions, streaming
`foreach` /
`foreachBatch` callbacks, and Python listeners.
- The session plan cache is keyed on a fingerprint of the environment in
addition to the relation.
- Three internal, cluster-level configurations bound the environment: at
most 100 variables, names
at most 512 characters, and 128 KiB total measured as the sum of the UTF-8
lengths of every name
and value. Zero accepts no user-provided environment at all; a negative
value is rejected.
Design notes a reviewer may want:
- **The configurations are the authoritative session state.** Nothing is
cached beside them, so the
environment follows the session wherever ordinary session configurations
follow it: it survives
reattach and retry with no code, and `SQLConf.clone()` carries it into a
session created by
`cloneSession`, while `newSession` correctly starts without one. There is
a test for each.
- **One snapshot per request.** A request reads the environment once and
uses that snapshot for the
cache lookup, for building its Python functions, and for the cache
insertion. Re-reading would
race with a concurrent configuration write and could store a plan built
with one environment
under the key of another.
- **The cache key is a fingerprint, not the environment.** A configuration
is stored before
anything validates it, so the environment is unbounded at that point;
putting it in a key would
let a session multiply the memory it holds by the cache size just by
issuing ordinary cacheable
queries. A SHA-256 digest is a fixed size whatever the environment holds,
and lengths are folded
in so that shifting a name/value boundary cannot collide.
- **A fresh mutable copy per function is required.** `BasePythonRunner`
takes the map by reference
and writes its own entries into it before launching a worker, so a shared
map would leak entries
between functions and an immutable one would fail the assignment.
- **Validation happens when a Python function is built**, so that every way
of writing a
configuration is covered by one check: the Connect config RPC, SQL `SET`,
and the
application-level configurations merged into a new session all arrive
there. The cost is that an
invalid environment stays in the session until the user corrects it; the
benefit is that it
cannot reach a worker, and that one invalid entry fails only the queries
that would install it
rather than every query in the session.
- **A value containing NUL is rejected.** A process environment cannot carry
it, and the JDK's own
rejection embeds the offending value in its message, so this has to be
caught before a worker
launch is attempted.
- **Names are preserved case-sensitively.** On a case-sensitive operating
system `FOO` and `foo` are
therefore distinct; Windows process environments are case-insensitive, so
what a worker observes
there is the platform's business.
- The accepted name pattern is deliberately stricter than the OS requires.
POSIX permits any byte
except `=` and NUL, and container platforms accept their own broader sets,
but a name outside
`[A-Za-z_][A-Za-z0-9_]*` cannot be referenced portably from a shell. It is
a portability policy,
not a description of what a process environment can hold.
- The name pattern is checked with a whole-string match rather than a
search. An anchored pattern
that is searched for would accept a name with a trailing newline, since
`$` also matches before a
terminating line break.
- Rejections reuse the existing `INVALID_SPARK_CONFIG` condition rather than
adding a new top-level
one. A message may name a variable but never carries its value, and a name
is truncated and has
its control characters escaped, so a rejection cannot forge log lines.
Note the name itself is
user-chosen, so it is only as safe as what the user put in it.
- An empty value is accepted (`FOO=` in a shell). A null value needs no
handling: `SQLConf`
rejects one on the way in, so a config request with an absent value fails
rather than storing
null.
Python UDTFs (`transformPythonTableFunction`) and Python data sources
(`transformPythonDataSource`)
have their own construction sites and keep receiving an empty environment;
they are follow-ups.
Not addressed here, and worth a reviewer's attention: a user can set a name
that Spark's own worker
protocol uses. Variables the runner sets unconditionally
(`SPARK_AUTH_SOCKET_TIMEOUT`,
`SPARK_BUFFER_SIZE`, `PYTHONPATH`, `PYTHON_WORKER_FACTORY_SECRET`, ...)
always overwrite a user
value, but ones set only under a condition — `SPARK_REUSE_WORKER`,
`SPARK_PIPELINED_UDF`,
`SPARK_HIDE_TRACEBACK` and others — are not removed when that condition is
false, so a user value
survives. Deciding between rejecting Spark-owned names and explicitly
clearing every one of them is
left to a follow-up.
### Why are the changes needed?
Code that reads `os.environ` behaves differently inside a Python UDF than
outside it, and a Spark
Connect client has no way to influence it. A user can set a value, read it
successfully from driver
code, and get a `KeyError` for the same name inside a UDF.
This is also the gap that blocks moving existing workloads onto Spark
Connect: on classic compute an
executor environment can be configured for the application through
`spark.executorEnv.*`, but that
is application-scoped and set before the context starts, so it has no
session-scoped equivalent a
Connect client can use.
### Does this PR introduce _any_ user-facing change?
Yes. Session configurations under `spark.pythonWorkerEnv.` are now read and
installed in the
environment of the Python worker processes that run the session's Python
functions, so `os.environ`
inside a Python UDF can see them. Previously these configurations had no
effect, and the worker
environment was always empty.
Setting a malformed or oversized environment now fails the queries that
would install it, with
`INVALID_SPARK_CONFIG.INVALID_PYTHON_WORKER_ENV_VAR_NAME`,
`INVALID_SPARK_CONFIG.INVALID_PYTHON_WORKER_ENV_VAR_VALUE`,
`INVALID_SPARK_CONFIG.PYTHON_WORKER_ENV_TOO_MANY_VARIABLES`, or
`INVALID_SPARK_CONFIG.PYTHON_WORKER_ENV_TOO_LARGE`.
The three new bounding configurations are internal.
### How was this patch tested?
`PythonWorkerEnvironmentSuite` (new, 35 tests):
- Reading: variables under the prefix, configurations outside the prefix
ignored, an empty value,
case sensitivity, no configurations at all, and that a null value cannot
be installed.
- Validation: a malformed name (including a trailing newline and an empty
name), a name over the
length limit with the name bounded in the message, a value containing NUL,
more variables than
the limit, exactly the limit, a total size over the limit, and a total
size that exceeds the
limit only when counted in UTF-8 bytes rather than characters.
- Message safety: a name carrying newlines, tabs, DEL and an ANSI escape is
escaped, and the
message carries no control characters.
- Limits: each of the three is exercised at a non-default value, zero
accepts nothing, and a
negative value is rejected by the configuration itself.
- Fingerprint: empty for an empty environment, stable and order-independent,
distinct for distinct
environments, fixed length however large the environment, and not confused
by shifting a
name/value boundary.
- Delivery: every scalar family and `mapInPandas` / `mapInArrow` receive the
environment; an empty
one when nothing is set; each function gets an independent mutable copy;
an invalid environment
fails planning of a Python function but not of a plan without one.
- Plan cache: a plan cached under one environment is not reused under
another, and the key holds no
environment values.
`SparkConnectPythonWorkerEnvTests` (new, end-to-end through a real Connect
client and a real Python
worker): a UDF reads the value from `os.environ`; an unset name is not
visible; an update is picked
up; `unset` removes it; an empty value arrives as empty; a platform-owned
variable
(`PYTHONUNBUFFERED`) still wins; an invalid name and a NUL value fail the
query without printing the
value; and `mapInPandas` sees the environment.
`SparkConnectSessionHolderSuite` was updated for the new plan cache key and
its plan cache tests
still pass. `SparkThrowableSuite` passes with the new error sub-conditions.
build/sbt "connect/testOnly *PythonWorkerEnvironmentSuite
*SparkConnectSessionHolderSuite"
build/sbt "core/testOnly *SparkThrowableSuite"
python/run-tests --testnames
pyspark.sql.tests.connect.test_connect_python_worker_env
`connect/scalastyle`, `connect/Test/scalastyle`, and scalafmt (with CI's
`changedOnly=false`) are
clean.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (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]
