[SPARK-58752][CONNECT] Deliver session environment variables to Python UDF workers

[RamonZhou]

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, 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 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. No second copy is maintained as
  session state, 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 holds the request's snapshot. A cached plan bakes the environment into every
  Python function it contains, so an entry is only reusable by a request carrying the same
  environment. Note the key is therefore as large as the environment, and a configuration is stored
  before anything validates it, so the limits below bound what reaches a worker rather than what
  the session and its cache hold.
• 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, 36 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.
• Plan cache: an entry is keyed on the snapshot the plan was built with; a change to the
  configurations midway through planning does not mis-key the entry, and the plan built under the
  old environment is not reused afterwards; five successive environments produce five entries; an
  oversized environment still does not stop an ordinary query being planned and cached.
• Session lifecycle also covers the Connect cloneSession path through the session manager, not
  only SparkSession.cloneSession.
• 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.

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. The end-to-end suite was
run against a real Connect server and real Python workers: 11 tests, all passing.

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)