dongjoon-hyun commented on code in PR #58978: URL: https://github.com/apache/spark/pull/58978#discussion_r4109210156
########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonRuntime.scala: ########## @@ -0,0 +1,294 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.python + +import java.nio.ByteBuffer +import java.util.concurrent.{Callable, ExecutionException, TimeoutException, TimeUnit} + +import scala.jdk.CollectionConverters._ + +import jep.{JepException, MainInterpreter, PyConfig, SharedInterpreter} + +import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.api.python.PythonException +import org.apache.spark.internal.Logging +import org.apache.spark.util.{ThreadUtils, Utils} + +/** Owns one interpreter generation per executor plugin lifecycle. */ +private[python] object InProcessPythonRuntime extends Logging { + val SITE_PACKAGES_CONFIG = "spark.inprocess.python.sitePackages" + private val TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" + private var active: InterpreterSession = _ + private var configured = false + + private[python] class LifecycleException(message: String) extends IllegalStateException(message) + + private def configureInterpreter(): Unit = { + if (!configured) { + // Like Python workers, use a stable default hash seed on every executor. This must + // happen before JEP creates its process-wide main interpreter, including on restarts. + MainInterpreter.setInitParams(new PyConfig().setHashSeed(0).setUseHashSeed(true)) + configured = true + } + } + + private[python] def bootstrapScript(script: String): String = { + "try:\n" + script.linesIterator.map(" " + _).mkString("\n") + + "\nexcept BaseException as _bootstrap_error:\n" + + " raise RuntimeError('In-process Python bootstrap failed: ' + " + + "repr(_bootstrap_error)) from None\n" + } + + def initialize(sitePackages: Seq[String] = Seq.empty): Unit = synchronized { + if (active != null && !active.isTerminated) { + active.requireCompatible(sitePackages) + } else { + configureInterpreter() + val candidate = new InterpreterSession(sitePackages) + try { + candidate.initialize() + active = candidate + } catch { + case t: Throwable => Utils.tryWithSafeFinally { throw t } { candidate.shutdown() } + } + } + } + + def currentSession: InterpreterSession = synchronized { + checkState(active != null && active.isRunning) + active + } + + def shutdown(): Unit = { + val session = synchronized { active } + if (session != null) session.shutdown() + } + + private def checkState(running: Boolean): Unit = { + checkState(running, "In-process Python is not running; initialize the executor plugin first") + } + + private def checkState(running: Boolean, message: String): Unit = { + if (!running) throw new IllegalStateException(message) + } + + /** + * Tasks retain this generation, so stale tasks cannot enter a later SparkContext's interpreter. + * Lifecycle operations only hold the monitor while enqueueing work, never while running Python. + */ + private[python] class InterpreterSession(val sitePackages: Seq[String] = Seq.empty) { + private val executor = ThreadUtils.newDaemonSingleThreadExecutor("inprocess-python") + @volatile private var running = true + // Accessed only on the owning thread. + private var interp: SharedInterpreter = _ + + def isRunning: Boolean = running + def isTerminated: Boolean = executor.isTerminated + + def requireCompatible(paths: Seq[String]): Unit = { + if (!isRunning) { + throw new LifecycleException("In-process Python is still stopping. Wait for outstanding " + + "native work to finish or replace the executor process before starting a new context.") + } + if (sitePackages != paths) { + throw new LifecycleException("In-process Python is already running with different " + + "sitePackages. Stop the existing context before changing interpreter configuration.") + } + } + + private[python] def onInterpreterThread[T](body: => T): T = { + val context = Option(TaskContext.get()) + context.foreach(_.killTaskIfInterrupted()) + val gate = new Object + var started = false + var cancelled = false + val future = synchronized { + checkState(running) + executor.submit(new Callable[T] { + override def call(): T = { + gate.synchronized { + if (cancelled) throw new TaskKilledException("Cancelled before Python invocation") + started = true + } + body + } + }) + } + var interrupted = false + try { + while (true) { + val taskCancelled = context.exists(_.isInterrupted()) + if (interrupted || taskCancelled) { + val cancelledBeforeStart = gate.synchronized { + if (started) false else { + cancelled = true + future.cancel(false) + true + } + } + if (cancelledBeforeStart) { + context.foreach(_.killTaskIfInterrupted()) + throw new InterruptedException("Cancelled before Python invocation") + } + } + try { + val result = future.get(100, TimeUnit.MILLISECONDS) + context.foreach(_.killTaskIfInterrupted()) + return result + } catch { + case _: TimeoutException => + case _: InterruptedException => interrupted = true + case e: ExecutionException => throw e.getCause + } + } + throw new IllegalStateException("Unreachable") + } finally { + // Once native work starts, wait for it even after cancellation: the caller still owns + // CDI structs that Python may use. Pending work, however, is safe to cancel immediately. + if (interrupted) Thread.currentThread().interrupt() + } + } + + def initialize(): Unit = onInterpreterThread { + val candidate = new SharedInterpreter() Review Comment: **Deployment: `import jep` runs before `sitePackages` is applied.** In JEP 4.3.2, the first `SharedInterpreter` construction calls `configureInterpreter`, which runs `eval("import jep")` and `from jep import java_import_hook`. Here, `site.addsitedir(sitePackages)` runs only after `new SharedInterpreter()` returns. The YARN and Kubernetes Option B examples in `docs/sql-pyspark-inprocess-udf.md` install `jep` only into the archived venv. They set only `spark.inprocess.python.sitePackages`, with no `PYTHONPATH`. Unless the system Python also has `jep`: - Construction fails with `ModuleNotFoundError: No module named 'jep'`. - The plugin rethrows, and every executor fails to start. The tests and the local example don't hit this because they prepend the jep directory to `PYTHONPATH` first. Could we pass the paths before the first construction, e.g. `SharedInterpreter.setConfig(new JepConfig().addIncludePaths(...))`? Alternatively, the cluster examples could document the `PYTHONPATH` requirement. ########## python/pyspark/sql/pandas/types.py: ########## @@ -253,18 +253,25 @@ def to_arrow_type( ) elif isinstance(dt, VariantType): fields = [ - pa.field("value", pa.binary(), nullable=False), + pa.field( + "value", pa.large_binary() if prefers_large_types else pa.binary(), nullable=False Review Comment: **Behavior change: this also changes existing worker-based Arrow paths.** Aligning Python with JVM `ArrowUtils` makes sense. However, `to_arrow_type` is shared, so with `spark.sql.execution.arrow.useLargeVarTypes=true` it also changes behavior outside in-process UDFs: - `applyInArrow` (grouped, cogrouped, and iterator variants) builds the return schema via `to_arrow_schema(..., prefers_large_types=...)` and checks it with `enforce_schema(..., arrow_cast=False)`. - A function that returns a variant/geometry/geography struct built with `pa.binary()` children used to pass and now fails with `RESULT_COLUMN_TYPES_MISMATCH`. - In the other direction, passing the JVM input through unchanged now passes, where it used to fail. - Classic `DataFrame.toArrow()` now returns `large_binary` children for these types. This contradicts "Existing Python UDF execution remains unchanged" in the PR description, and there are no tests for these paths. Could we split this into a separate JIRA/PR with regression tests for `applyInArrow` and `toArrow`? ########## python/pyspark/sql/tests/connect/test_connect_plan.py: ########## @@ -76,6 +76,16 @@ class SparkConnectPlanTests(PlanOnlyTestFixture): """These test cases exercise the interface to the proto plan generation but do not call Spark.""" + def test_inprocess_udf_registration_is_rejected(self): + from pyspark.errors import PySparkTypeError + from pyspark.inprocess import inprocess_udf Review Comment: **CI: this fails in the Connect-only (`pyspark-client`) job.** `pyspark.inprocess` is added only to `python/packaging/classic/setup.py`, not to `connect_packages` in `python/packaging/client/setup.py`. The scheduled `build_python_connect.yml` job: - installs the `pyspark-client` sdist, - moves `python/pyspark` away, - runs `--modules pyspark-connect`, which includes this module. `SparkConnectPlanTests` is guarded only by `should_test_connect`, so this test errors with `ModuleNotFoundError`. PR CI won't catch it because it runs with the full source tree. Packaging it wouldn't help either. `pyspark/inprocess/udf.py` does `from pyspark import Accumulator, Broadcast, cloudpickle`, and `Broadcast` is exported only when `not is_remote_only()`. Could you add `@unittest.skipIf(is_remote_only(), ...)`, like other Connect tests that need classic PySpark (e.g. `test_connect_function.py`)? Moving this test to a classic-only suite would also work. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonRuntime.scala: ########## @@ -0,0 +1,294 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.python + +import java.nio.ByteBuffer +import java.util.concurrent.{Callable, ExecutionException, TimeoutException, TimeUnit} + +import scala.jdk.CollectionConverters._ + +import jep.{JepException, MainInterpreter, PyConfig, SharedInterpreter} + +import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.api.python.PythonException +import org.apache.spark.internal.Logging +import org.apache.spark.util.{ThreadUtils, Utils} + +/** Owns one interpreter generation per executor plugin lifecycle. */ +private[python] object InProcessPythonRuntime extends Logging { + val SITE_PACKAGES_CONFIG = "spark.inprocess.python.sitePackages" + private val TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" + private var active: InterpreterSession = _ + private var configured = false + + private[python] class LifecycleException(message: String) extends IllegalStateException(message) + + private def configureInterpreter(): Unit = { + if (!configured) { + // Like Python workers, use a stable default hash seed on every executor. This must + // happen before JEP creates its process-wide main interpreter, including on restarts. + MainInterpreter.setInitParams(new PyConfig().setHashSeed(0).setUseHashSeed(true)) + configured = true + } + } + + private[python] def bootstrapScript(script: String): String = { + "try:\n" + script.linesIterator.map(" " + _).mkString("\n") + + "\nexcept BaseException as _bootstrap_error:\n" + + " raise RuntimeError('In-process Python bootstrap failed: ' + " + + "repr(_bootstrap_error)) from None\n" + } + + def initialize(sitePackages: Seq[String] = Seq.empty): Unit = synchronized { + if (active != null && !active.isTerminated) { + active.requireCompatible(sitePackages) + } else { + configureInterpreter() + val candidate = new InterpreterSession(sitePackages) + try { + candidate.initialize() + active = candidate + } catch { + case t: Throwable => Utils.tryWithSafeFinally { throw t } { candidate.shutdown() } + } + } + } + + def currentSession: InterpreterSession = synchronized { + checkState(active != null && active.isRunning) + active + } + + def shutdown(): Unit = { + val session = synchronized { active } + if (session != null) session.shutdown() + } + + private def checkState(running: Boolean): Unit = { + checkState(running, "In-process Python is not running; initialize the executor plugin first") + } + + private def checkState(running: Boolean, message: String): Unit = { + if (!running) throw new IllegalStateException(message) + } + + /** + * Tasks retain this generation, so stale tasks cannot enter a later SparkContext's interpreter. + * Lifecycle operations only hold the monitor while enqueueing work, never while running Python. + */ + private[python] class InterpreterSession(val sitePackages: Seq[String] = Seq.empty) { + private val executor = ThreadUtils.newDaemonSingleThreadExecutor("inprocess-python") Review Comment: **Stability: the interpreter thread runs with the default JVM thread stack.** `ThreadUtils.newDaemonSingleThreadExecutor` doesn't set a stack size, so this thread gets the JVM default (1MB on Linux x64 unless `-Xss` is set). CPython 3.12/3.13 release builds use `C_RECURSION_LIMIT` / `Py_C_RECURSION_LIMIT = 10000`, which assumes a typical ~8MB main-thread stack. So deep native recursion can overflow the 1MB stack before `RecursionError` is raised. Examples: - `json.loads` on a deeply nested document - `repr` or pickle of deep structures - unpickling a deep closure graph during registration With a Python worker, that is a task failure. Here it would be a SIGSEGV that kills the whole executor JVM, and it would repeat on every retry. I haven't reproduced the crash, but these constants make it plausible. How about creating this thread with an explicit larger stack size, e.g. `new Thread(null, runnable, name, 8L << 20)` or more? ########## python/pyspark/inprocess/runtime.py: ########## @@ -0,0 +1,261 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +"""Arrow CDI entry points called on the executor's dedicated JEP interpreter thread. + +Functions are registered once per task and released when that task finishes. Calls +pass only a handle and CDI addresses, so large closures are not copied per batch. +""" + +import sys +import traceback as _traceback +from typing import Any, Callable, Iterable, Optional, Sequence + +import pyarrow as pa +import pyarrow.compute as pc + +from pyspark import cloudpickle +from pyspark.errors import PySparkRuntimeError +from pyspark.sql.pandas.types import to_arrow_type +from pyspark.sql.types import _parse_datatype_json_string +from pyspark.util import try_simplify_traceback + +_UDF_TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" +NullChecker = Callable[[pa.Array], None] +_udfs: dict[str, tuple[Callable[..., pa.Array], pa.DataType, NullChecker, bool, bool]] = {} + + +def _format_exception(hide: bool, simplified: bool) -> str: + kind, error, tb = sys.exc_info() + if hide: + return "".join(_traceback.format_exception_only(kind, error)) + if simplified and tb is not None: + simple_tb = try_simplify_traceback(tb) + if simple_tb is not None: + tb = simple_tb + if error is not None: + error.__cause__ = None + return "".join(_traceback.format_exception(kind, error, tb)) + + +def _inprocess_register( + handle: str, + serialized_udf: Any, + return_type_json: str, + timezone: str, + python_version: str, + large_var_types: bool = False, + hide_traceback: bool = False, + simplified_traceback: bool = False, +) -> None: + try: + embedded_version = "%d.%d" % sys.version_info[:2] + if python_version != embedded_version: + raise PySparkRuntimeError( + errorClass="PYTHON_VERSION_MISMATCH", + messageParameters={ + "worker_version": embedded_version, + "driver_version": python_version, + }, + ) + # JEP exposes direct ByteBuffers through the buffer protocol. Unpickle a separate + # function per task without iterating over a PyJArray one JNI call per byte. + func = cloudpickle.loads(memoryview(serialized_udf)) + expected_type = to_arrow_type( + _parse_datatype_json_string(return_type_json), + timezone=timezone, + prefers_large_types=large_var_types, + error_on_duplicated_field_names_in_struct=True, + ) + checker = _null_checker(expected_type) or (lambda array: None) + _udfs[handle] = (func, expected_type, checker, hide_traceback, simplified_traceback) + except BaseException: + # In JEP, an uncaught SystemExit can terminate the entire executor JVM. + raise RuntimeError( + _UDF_TRACEBACK_SENTINEL + _format_exception(hide_traceback, simplified_traceback) + ) from None + + +def _inprocess_release(handles: Iterable[str]) -> None: + for handle in handles: + _udfs.pop(handle, None) + + +def _nullable_type(data_type: pa.DataType) -> pa.DataType: + def nullable_field(field: pa.Field) -> pa.Field: + return pa.field(field.name, _nullable_type(field.type), nullable=True) + + if pa.types.is_struct(data_type): + return pa.struct([nullable_field(field) for field in data_type]) + if pa.types.is_list(data_type): + return pa.list_(nullable_field(data_type.value_field)) + if pa.types.is_large_list(data_type): + return pa.large_list(nullable_field(data_type.value_field)) + if pa.types.is_map(data_type): + return pa.map_( + _nullable_type(data_type.key_type), + nullable_field(data_type.item_field), + keys_sorted=data_type.keys_sorted, + ) + return data_type + + +def _null_checker(expected_type: pa.DataType) -> Optional[NullChecker]: + """Compile checks only for required fields and their ancestors, once per registration.""" + + def field_checker(field: pa.Field) -> Optional[NullChecker]: + nested = _null_checker(field.type) + if field.nullable: + return nested + + def check(values: pa.Array) -> None: + if values.null_count: + raise ValueError( + f"In-process UDF returned nulls in non-nullable field {field.name}" + ) + if nested is not None: + nested(values) + + return check + + if pa.types.is_struct(expected_type): + fields = [(i, field_checker(f)) for i, f in enumerate(expected_type)] + checks = [(i, check) for i, check in fields if check is not None] + if not checks: + return None + + def check_struct(array: pa.Array) -> None: + # Only children of valid parents are logically visible. + visible = pc.filter(array, pc.is_valid(array)) if array.null_count else array + for i, check in checks: + check(visible.field(i)) + + return check_struct + if pa.types.is_list(expected_type) or pa.types.is_large_list(expected_type): + check = field_checker(expected_type.value_field) + if check is not None: + return lambda array: check(pc.list_flatten(array)) + if pa.types.is_map(expected_type): + key_check = field_checker(expected_type.key_field) + item_check = field_checker(expected_type.item_field) + + def check_map(array: pa.Array) -> None: + visible = pc.filter(array, pc.is_valid(array)) if array.null_count else array + start = visible.offsets[0].as_py() + length = visible.offsets[-1].as_py() - start + if key_check is not None: + key_check(visible.keys.slice(start, length)) Review Comment: **Correctness: `check_map` ignores the offset of the map's entries struct.** `MapArray.keys` / `.items` return the raw children of the entries struct without applying the entries struct's own offset, while `offsets` index into the offset-adjusted entries. So when the entries child is sliced and the map has no top-level nulls, this checks the wrong window. (The earlier thread about sliced entries was fixed via `array.values`, but the compiled `check_map` went back to `keys`/`items`.) Repro against `_validate_result` with PyArrow 25, declared type `MapType(StringType(), LongType(), valueContainsNull=False)`: ```python fields = [pa.field("key", pa.string(), nullable=False), pa.field("value", pa.int64())] entries = pa.StructArray.from_arrays( [pa.array(["h", "a", "b", "c"]), pa.array([0, 1, 2, None])], fields=fields) offsets = pa.array([0, 1, 3], pa.int32()).buffers()[1] m = pa.Array.from_buffers( pa.map_(pa.string(), pa.int64()), 2, [None, offsets], children=[entries.slice(1)]) m.to_pylist() # [[('a', 1)], [('b', 2), ('c', None)]] ``` - The check sees `items[0:3] == [0, 1, 2]` and accepts. The null then reaches the JVM as a non-nullable map value. - Conversely, items `[None, 1, 2, 3]` (no logical null) are rejected with `returned nulls in non-nullable field value`. Two possible fixes: - Use `visible.values.field(0)` / `.field(1)`, which honor the entries offset, sliced by `start`/`length`. - Run the checker after the `concat_arrays` normalization. ########## python/pyspark/inprocess/runtime.py: ########## @@ -0,0 +1,261 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +"""Arrow CDI entry points called on the executor's dedicated JEP interpreter thread. + +Functions are registered once per task and released when that task finishes. Calls +pass only a handle and CDI addresses, so large closures are not copied per batch. +""" + +import sys +import traceback as _traceback +from typing import Any, Callable, Iterable, Optional, Sequence + +import pyarrow as pa +import pyarrow.compute as pc + +from pyspark import cloudpickle +from pyspark.errors import PySparkRuntimeError +from pyspark.sql.pandas.types import to_arrow_type +from pyspark.sql.types import _parse_datatype_json_string +from pyspark.util import try_simplify_traceback + +_UDF_TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" +NullChecker = Callable[[pa.Array], None] +_udfs: dict[str, tuple[Callable[..., pa.Array], pa.DataType, NullChecker, bool, bool]] = {} + + +def _format_exception(hide: bool, simplified: bool) -> str: Review Comment: **Config: `spark.sql.execution.pyspark.udf.tracebackWithLocals.enabled` is ignored.** The worker path honors this conf: `ArrowPythonRunner.tracebackWithLocals` -> `SPARK_TRACEBACK_WITH_LOCALS` -> `capture_locals=True` in `pyspark.util.handle_worker_exception`. Here, `ArrowEvalPythonExec` passes only `hideTraceback` and `simplifiedTraceback`, and this formatter never captures locals. So the same failing UDF shows locals as an `arrow_udf` but not as an `inprocess_udf`. This function largely duplicates the nested `format_exception` in `handle_worker_exception`. Extracting a shared helper would keep both paths in sync. ########## python/pyspark/inprocess/runtime.py: ########## @@ -0,0 +1,261 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +"""Arrow CDI entry points called on the executor's dedicated JEP interpreter thread. + +Functions are registered once per task and released when that task finishes. Calls +pass only a handle and CDI addresses, so large closures are not copied per batch. +""" + +import sys +import traceback as _traceback +from typing import Any, Callable, Iterable, Optional, Sequence + +import pyarrow as pa +import pyarrow.compute as pc + +from pyspark import cloudpickle +from pyspark.errors import PySparkRuntimeError +from pyspark.sql.pandas.types import to_arrow_type +from pyspark.sql.types import _parse_datatype_json_string +from pyspark.util import try_simplify_traceback + +_UDF_TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" +NullChecker = Callable[[pa.Array], None] +_udfs: dict[str, tuple[Callable[..., pa.Array], pa.DataType, NullChecker, bool, bool]] = {} + + +def _format_exception(hide: bool, simplified: bool) -> str: + kind, error, tb = sys.exc_info() + if hide: + return "".join(_traceback.format_exception_only(kind, error)) + if simplified and tb is not None: + simple_tb = try_simplify_traceback(tb) + if simple_tb is not None: + tb = simple_tb + if error is not None: + error.__cause__ = None + return "".join(_traceback.format_exception(kind, error, tb)) + + +def _inprocess_register( + handle: str, + serialized_udf: Any, + return_type_json: str, + timezone: str, + python_version: str, + large_var_types: bool = False, + hide_traceback: bool = False, + simplified_traceback: bool = False, +) -> None: + try: + embedded_version = "%d.%d" % sys.version_info[:2] + if python_version != embedded_version: + raise PySparkRuntimeError( + errorClass="PYTHON_VERSION_MISMATCH", + messageParameters={ + "worker_version": embedded_version, + "driver_version": python_version, + }, + ) + # JEP exposes direct ByteBuffers through the buffer protocol. Unpickle a separate + # function per task without iterating over a PyJArray one JNI call per byte. + func = cloudpickle.loads(memoryview(serialized_udf)) + expected_type = to_arrow_type( + _parse_datatype_json_string(return_type_json), Review Comment: **Correctness: Python UDT return types must be importable on executors.** The return type reaches the executor as JSON (`udfs.map(_.dataType.json)`) and is rebuilt here via `_parse_datatype_json_string`. `UserDefinedType.fromJson` then imports the `pyClass` by module name. Python workers instead receive the return type through cloudpickle (`(func, returnType)`), so UDTs defined in the driver's `__main__` work there. So `inprocess_udf(return_type=MyUDT())` fails every task with `[UNSUPPORTED_OPERATION] unpickling user defined types` when `MyUDT` is defined in a notebook or driver script, or in any module not installed on executors. `to_arrow_type` and `ArrowUtils` otherwise support UDTs. Could we pickle the return type together with the function, or document this limitation? ########## python/pyspark/inprocess/udf.py: ########## @@ -0,0 +1,221 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +Python API for in-process UDF registration. + +Usage:: + + import pyarrow.compute as pc + from pyspark.inprocess import inprocess_udf + from pyspark.sql.types import LongType + + @inprocess_udf(return_type=LongType()) + def double(x): + # x is a pa.Array; return a pa.Array + return pc.multiply(x, 2) + + df.select(double(df.value)).show() +""" + +import io +import sys +from functools import update_wrapper +from inspect import getfullargspec +from typing import Any, Callable, Optional, Union + +from pyspark import Accumulator, Broadcast, cloudpickle +from pyspark.errors import PySparkTypeError, PySparkValueError +from pyspark.sql.column import Column +from pyspark.sql.types import DataType, _parse_datatype_string +from pyspark.util import PythonEvalType + + +class _InProcessPickler(cloudpickle.CloudPickler): + def reducer_override(self, obj: Any) -> Any: + if isinstance(obj, (Broadcast, Accumulator)): + raise TypeError("In-process UDFs do not support Spark broadcasts or accumulators") + return super().reducer_override(obj) + + +def _serialize_udf(func: Callable) -> bytes: + buffer = io.BytesIO() + _InProcessPickler(buffer).dump(func) + return buffer.getvalue() + + +class InProcessUDFWrapper: + """ + Wraps a Python function as an in-process UDF. + + Returned by ``@inprocess_udf``. Calling an instance with Spark ``Column`` + arguments creates a ``Column`` expression backed by ``PythonUDF`` + on the JVM side. + """ + + def __init__( + self, func: Callable, return_type: Union[DataType, str], deterministic: bool = True + ) -> None: + if not isinstance(return_type, (DataType, str)): + raise PySparkTypeError( + errorClass="NOT_EXPECTED_TYPE", + messageParameters={ + "expected_type": "DataType or str", + "arg_name": "return_type", + "arg_type": type(return_type).__name__, + }, + ) + self._return_type = return_type + self._parsed_return_type: Optional[DataType] = None + self.evalType = PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF + self._deterministic: bool = deterministic + self._name: str = getattr(func, "__name__", "inprocess_udf") + + argspec = getfullargspec(func) + if not argspec.args and argspec.varargs is None and not argspec.kwonlyargs: + raise PySparkValueError( + errorClass="INVALID_PANDAS_UDF", + messageParameters={"detail": "0-arg inprocess_udfs are not supported."}, + ) + self._func = func + self._serialized: Optional[bytes] = None + update_wrapper(self, func, updated=()) + + @property + def func(self) -> Callable: + return self._func + + @property + def returnType(self) -> DataType: + if self._parsed_return_type is None: + self._parsed_return_type = ( Review Comment: **Validation: Arrow-unsupported return types are not rejected on the driver.** `UserDefinedFunction.returnType` calls `_check_return_type`, so `arrow_udf` rejects types that `to_arrow_type` doesn't support (e.g. `YearMonthIntervalType`, `CalendarIntervalType`) right away on the driver with `NOT_IMPLEMENTED`. Here the type is only parsed, and JVM `ArrowUtils` accepts these types. As a result, `@inprocess_udf("interval year to month")`: - passes analysis and planning, - then fails in `_inprocess_register` in every non-empty task with `UNSUPPORTED_DATA_TYPE_FOR_ARROW_CONVERSION`. Because the failure is deterministic, the job only aborts after `spark.task.maxFailures` retries. `CharType`/`VarcharType` behave the same way. Could we call `to_arrow_type` (or reuse `_check_return_type`) when building the column? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonRuntime.scala: ########## @@ -0,0 +1,294 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.python + +import java.nio.ByteBuffer +import java.util.concurrent.{Callable, ExecutionException, TimeoutException, TimeUnit} + +import scala.jdk.CollectionConverters._ + +import jep.{JepException, MainInterpreter, PyConfig, SharedInterpreter} + +import org.apache.spark.{TaskContext, TaskKilledException} +import org.apache.spark.api.python.PythonException +import org.apache.spark.internal.Logging +import org.apache.spark.util.{ThreadUtils, Utils} + +/** Owns one interpreter generation per executor plugin lifecycle. */ +private[python] object InProcessPythonRuntime extends Logging { + val SITE_PACKAGES_CONFIG = "spark.inprocess.python.sitePackages" + private val TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" + private var active: InterpreterSession = _ + private var configured = false + + private[python] class LifecycleException(message: String) extends IllegalStateException(message) + + private def configureInterpreter(): Unit = { + if (!configured) { + // Like Python workers, use a stable default hash seed on every executor. This must + // happen before JEP creates its process-wide main interpreter, including on restarts. + MainInterpreter.setInitParams(new PyConfig().setHashSeed(0).setUseHashSeed(true)) + configured = true + } + } + + private[python] def bootstrapScript(script: String): String = { + "try:\n" + script.linesIterator.map(" " + _).mkString("\n") + + "\nexcept BaseException as _bootstrap_error:\n" + + " raise RuntimeError('In-process Python bootstrap failed: ' + " + + "repr(_bootstrap_error)) from None\n" + } + + def initialize(sitePackages: Seq[String] = Seq.empty): Unit = synchronized { + if (active != null && !active.isTerminated) { + active.requireCompatible(sitePackages) + } else { + configureInterpreter() + val candidate = new InterpreterSession(sitePackages) + try { + candidate.initialize() + active = candidate + } catch { + case t: Throwable => Utils.tryWithSafeFinally { throw t } { candidate.shutdown() } + } + } + } + + def currentSession: InterpreterSession = synchronized { + checkState(active != null && active.isRunning) + active + } + + def shutdown(): Unit = { + val session = synchronized { active } + if (session != null) session.shutdown() + } + + private def checkState(running: Boolean): Unit = { + checkState(running, "In-process Python is not running; initialize the executor plugin first") + } + + private def checkState(running: Boolean, message: String): Unit = { + if (!running) throw new IllegalStateException(message) + } + + /** + * Tasks retain this generation, so stale tasks cannot enter a later SparkContext's interpreter. + * Lifecycle operations only hold the monitor while enqueueing work, never while running Python. + */ + private[python] class InterpreterSession(val sitePackages: Seq[String] = Seq.empty) { + private val executor = ThreadUtils.newDaemonSingleThreadExecutor("inprocess-python") + @volatile private var running = true + // Accessed only on the owning thread. + private var interp: SharedInterpreter = _ + + def isRunning: Boolean = running + def isTerminated: Boolean = executor.isTerminated + + def requireCompatible(paths: Seq[String]): Unit = { + if (!isRunning) { + throw new LifecycleException("In-process Python is still stopping. Wait for outstanding " + + "native work to finish or replace the executor process before starting a new context.") + } + if (sitePackages != paths) { + throw new LifecycleException("In-process Python is already running with different " + + "sitePackages. Stop the existing context before changing interpreter configuration.") + } + } + + private[python] def onInterpreterThread[T](body: => T): T = { + val context = Option(TaskContext.get()) + context.foreach(_.killTaskIfInterrupted()) + val gate = new Object + var started = false + var cancelled = false + val future = synchronized { + checkState(running) + executor.submit(new Callable[T] { + override def call(): T = { + gate.synchronized { + if (cancelled) throw new TaskKilledException("Cancelled before Python invocation") + started = true + } + body + } + }) + } + var interrupted = false + try { + while (true) { + val taskCancelled = context.exists(_.isInterrupted()) + if (interrupted || taskCancelled) { + val cancelledBeforeStart = gate.synchronized { + if (started) false else { + cancelled = true + future.cancel(false) + true + } + } + if (cancelledBeforeStart) { + context.foreach(_.killTaskIfInterrupted()) + throw new InterruptedException("Cancelled before Python invocation") + } + } + try { + val result = future.get(100, TimeUnit.MILLISECONDS) + context.foreach(_.killTaskIfInterrupted()) + return result + } catch { + case _: TimeoutException => + case _: InterruptedException => interrupted = true + case e: ExecutionException => throw e.getCause + } + } + throw new IllegalStateException("Unreachable") + } finally { + // Once native work starts, wait for it even after cancellation: the caller still owns + // CDI structs that Python may use. Pending work, however, is safe to cancel immediately. + if (interrupted) Thread.currentThread().interrupt() + } + } + + def initialize(): Unit = onInterpreterThread { + val candidate = new SharedInterpreter() + try { + candidate.set("_site_packages", sitePackages.asJava) + candidate.exec(bootstrapScript( + """import os, site, sys + |_configured = [os.path.abspath(p) for p in _site_packages] + |_before = set(sys.path) + |for _path in _configured: + | site.addsitedir(_path) + |_added = [p for p in sys.path if p not in _before and p not in _configured] + |_preferred = list(dict.fromkeys(_configured + _added)) + |sys.path[:] = _preferred + [p for p in sys.path if p not in _preferred] + |del _site_packages, _configured, _before, _added, _preferred + |""".stripMargin)) + candidate.exec(bootstrapScript("from pyspark.inprocess.runtime import " + Review Comment: **Deployment: the embedded `sys.path` doesn't include Spark's own `pyspark.zip`.** Python workers always prepend `PythonUtils.sparkPythonPath` (`$SPARK_HOME/python/lib/pyspark.zip` and py4j) in `PythonWorkerFactory`. The bootstrap here adds only `sitePackages` and then imports `pyspark.inprocess.runtime`. On Kubernetes, nothing sets the executor `PYTHONPATH`, and the official `apache/spark` image doesn't pip-install pyspark. So the Option A image recommended in the docs (`FROM apache/spark:latest` + `pip install jep pyarrow ...`) fails at plugin init with `No module named 'pyspark'`. That recipe has two more problems: - `jep` 4.3.2 ships only as an sdist, and the image has no compiler or Python headers, so `pip install jep` fails. - The image's `python3` is 3.10, below the documented 3.11+ requirement. Also, configured `sitePackages` are placed ahead of existing paths. On YARN, a pip-installed `pyspark` in the venv could therefore shadow the version-matched `pyspark.zip`, even though the JVM/Python contract here is version-coupled. Prepending `PythonUtils.sparkPythonPath`, as the worker does, would address most of this. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonUDFBuilder.scala: ########## @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.python + +import java.util.{Collections, List => JList} + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction} +import org.apache.spark.sql.Column +import org.apache.spark.sql.catalyst.expressions.PythonUDF +import org.apache.spark.sql.classic.{ColumnNodeExpression, ExpressionUtils} +import org.apache.spark.sql.types.DataType + +/** + * JVM-side builder for in-process [[PythonUDF]] expressions, called from the Python API + * via py4j's JVM reflection bridge (``sc._jvm.org.apache.spark...InProcessPythonUDFBuilder``). + * + * Accepts Java-typed arguments as passed by PySpark's ``sc._jvm`` proxy and returns a + * [[Column]] backed by a [[PythonUDF]] with the in-process evaluation type. + */ +object InProcessPythonUDFBuilder { + + /** + * Build a [[Column]] backed by an in-process [[PythonUDF]] expression. + * + * @param name display name (Python function ``__name__``) + * @param serializedFunc cloudpickle bytes of the Python UDF + * @param returnTypeJson JSON string of the Spark SQL return type + * @param jColumns Java List of JVM [[Column]] objects (the UDF inputs) + * @param deterministic whether the UDF always returns the same output for the same input; + * set to false for UDFs that use randomness or external state + * @param pythonVersion driver's Python major.minor version + * @return [[Column]] backed by an in-process [[PythonUDF]] expression + */ + def build( + name: String, + serializedFunc: Array[Byte], + returnTypeJson: String, + jColumns: JList[Column], + deterministic: Boolean, + pythonVersion: String): Column = { + val returnType = DataType.fromJson(returnTypeJson) + val inputExprs = jColumns.asScala.map(col => ColumnNodeExpression(col.node)).toSeq + val function = new SimplePythonFunction( + serializedFunc, + Collections.emptyMap[String, String](), + Collections.emptyList[String](), + "", + pythonVersion, + Collections.emptyList(), + null) + ExpressionUtils.column(PythonUDF( Review Comment: **Validation: named-argument checks are skipped.** For eval types that support named arguments, `UserDefinedPythonFunction.builder` runs `NamedParametersSupport.splitAndCheckNamedArguments(e, name, SQLConf.get.resolver)`. This builder creates `PythonUDF` directly and skips that check. With the default case-insensitive resolver, `pandas_udf`/`arrow_udf` reject `f(x=df.a, X=df.b)` with `DUPLICATE_ROUTINE_PARAMETER_ASSIGNMENT.DOUBLE_NAMED_ARGUMENT_REFERENCE`. Here it reaches Python as `{"x": a, "X": b}`: - `def f(x, **kw)` binds it silently. - `def f(x, y)` fails at task runtime with `TypeError`. Could we apply the same check here? ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/InProcessPythonUDFBuilder.scala: ########## @@ -0,0 +1,72 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.execution.python + +import java.util.{Collections, List => JList} + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction} +import org.apache.spark.sql.Column +import org.apache.spark.sql.catalyst.expressions.PythonUDF +import org.apache.spark.sql.classic.{ColumnNodeExpression, ExpressionUtils} +import org.apache.spark.sql.types.DataType + +/** + * JVM-side builder for in-process [[PythonUDF]] expressions, called from the Python API + * via py4j's JVM reflection bridge (``sc._jvm.org.apache.spark...InProcessPythonUDFBuilder``). + * + * Accepts Java-typed arguments as passed by PySpark's ``sc._jvm`` proxy and returns a + * [[Column]] backed by a [[PythonUDF]] with the in-process evaluation type. + */ +object InProcessPythonUDFBuilder { + + /** + * Build a [[Column]] backed by an in-process [[PythonUDF]] expression. + * + * @param name display name (Python function ``__name__``) + * @param serializedFunc cloudpickle bytes of the Python UDF + * @param returnTypeJson JSON string of the Spark SQL return type + * @param jColumns Java List of JVM [[Column]] objects (the UDF inputs) + * @param deterministic whether the UDF always returns the same output for the same input; + * set to false for UDFs that use randomness or external state + * @param pythonVersion driver's Python major.minor version + * @return [[Column]] backed by an in-process [[PythonUDF]] expression + */ + def build( + name: String, + serializedFunc: Array[Byte], + returnTypeJson: String, + jColumns: JList[Column], + deterministic: Boolean, + pythonVersion: String): Column = { + val returnType = DataType.fromJson(returnTypeJson) + val inputExprs = jColumns.asScala.map(col => ColumnNodeExpression(col.node)).toSeq + val function = new SimplePythonFunction( + serializedFunc, + Collections.emptyMap[String, String](), Review Comment: **Contract: `spark.pythonWorkerEnv.*` is not applied.** `docs/configuration.md` says `spark.pythonWorkerEnv.*` applies to scalar UDFs in each of their forms, and the worker runners apply it via `PythonWorkerEnvironment.mergeValidated`. Here, the function is created with an empty env map and nothing else applies it. So after `spark.conf.set("spark.pythonWorkerEnv.FOO", "bar")`: - inside an in-process UDF, `os.environ["FOO"]` raises `KeyError` (or returns the executor process's value), - while the same code as a `pandas_udf` sees `"bar"`. Since the interpreter is process-wide and shared across sessions, supporting this may not be feasible. In that case, it should be documented or rejected explicitly. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala: ########## @@ -102,11 +103,12 @@ case class ArrowEvalPythonExec( // The Arrow FieldVectors are extracted directly from ArrowColumnVector and // serialized to IPC, bypassing the row-based ArrowWriter conversion. override def supportsColumnar: Boolean = - child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled + evalType != PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF && Review Comment: **Design: the in-process eval type is special-cased inside a worker-oriented node.** Eval type 258 is now handled in four places here: `supportsColumnar`, the `doExecute` guard, the early `return` in `evaluatorFactory`, and `supportedPythonEvalTypes`. The node also still carries worker-only state: `jobArtifactUUID`, `sessionUUID`, and `columnarEvaluatorFactory`, which would send 258 to a Python worker if reached. Missing any one of these guards would route 258 to a worker. The `doExecute` guard is actually required, because it branches on `child.supportsColumnar`. For example, `InMemoryTableScanExec` supports both row and columnar output, so no `ColumnarToRow` is inserted. How about mapping `ArrowEvalPython(..., SQL_SCALAR_ARROW_INPROCESS_UDF)` in `SparkStrategies.PythonEvals` to a dedicated `InProcessArrowEvalPythonExec extends EvalPythonExec`? - It would still reuse `EvalPythonEvaluatorFactory`. - Logical rules on `ArrowEvalPython` would be unaffected. - `InsertSortForLimitAndOffset` matches `EvalPythonExec`, so it would keep working. That would remove all of these special cases. ########## python/pyspark/sql/udf.py: ########## @@ -856,6 +856,14 @@ def register( [Row(sum_udf(v1)=1), Row(sum_udf(v1)=5)] """ + # Avoid importing the optional PyArrow-backed module for ordinary UDF registration. Review Comment: **Nit: this guard is now redundant.** `InProcessUDFWrapper` now has `asNondeterministic` and `evalType = 258`, so it enters the existing `if hasattr(f, "asNondeterministic"):` branch below. 258 is not in that eval-type allow-list, so the wrapper is already rejected with the same `INVALID_UDF_EVAL_TYPE`. Both new tests only check the error class, so they would pass without this block. The same applies to `python/pyspark/sql/connect/udf.py`. The comment is also slightly inaccurate: `pyspark/inprocess/udf.py` doesn't import PyArrow; only `runtime.py` does. Could we remove this block, including the `sys.modules` lookup, from both files? ########## python/pyspark/inprocess/runtime.py: ########## @@ -0,0 +1,261 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +"""Arrow CDI entry points called on the executor's dedicated JEP interpreter thread. + +Functions are registered once per task and released when that task finishes. Calls +pass only a handle and CDI addresses, so large closures are not copied per batch. +""" + +import sys +import traceback as _traceback +from typing import Any, Callable, Iterable, Optional, Sequence + +import pyarrow as pa +import pyarrow.compute as pc + +from pyspark import cloudpickle +from pyspark.errors import PySparkRuntimeError +from pyspark.sql.pandas.types import to_arrow_type +from pyspark.sql.types import _parse_datatype_json_string +from pyspark.util import try_simplify_traceback + +_UDF_TRACEBACK_SENTINEL = "__INPROCESS_UDF_TRACEBACK__:" +NullChecker = Callable[[pa.Array], None] +_udfs: dict[str, tuple[Callable[..., pa.Array], pa.DataType, NullChecker, bool, bool]] = {} + + +def _format_exception(hide: bool, simplified: bool) -> str: + kind, error, tb = sys.exc_info() + if hide: + return "".join(_traceback.format_exception_only(kind, error)) + if simplified and tb is not None: + simple_tb = try_simplify_traceback(tb) + if simple_tb is not None: + tb = simple_tb + if error is not None: + error.__cause__ = None + return "".join(_traceback.format_exception(kind, error, tb)) + + +def _inprocess_register( + handle: str, + serialized_udf: Any, + return_type_json: str, + timezone: str, + python_version: str, + large_var_types: bool = False, + hide_traceback: bool = False, + simplified_traceback: bool = False, +) -> None: + try: + embedded_version = "%d.%d" % sys.version_info[:2] + if python_version != embedded_version: + raise PySparkRuntimeError( + errorClass="PYTHON_VERSION_MISMATCH", + messageParameters={ + "worker_version": embedded_version, + "driver_version": python_version, + }, + ) + # JEP exposes direct ByteBuffers through the buffer protocol. Unpickle a separate + # function per task without iterating over a PyJArray one JNI call per byte. + func = cloudpickle.loads(memoryview(serialized_udf)) + expected_type = to_arrow_type( + _parse_datatype_json_string(return_type_json), + timezone=timezone, + prefers_large_types=large_var_types, + error_on_duplicated_field_names_in_struct=True, + ) + checker = _null_checker(expected_type) or (lambda array: None) + _udfs[handle] = (func, expected_type, checker, hide_traceback, simplified_traceback) + except BaseException: + # In JEP, an uncaught SystemExit can terminate the entire executor JVM. + raise RuntimeError( + _UDF_TRACEBACK_SENTINEL + _format_exception(hide_traceback, simplified_traceback) + ) from None + + +def _inprocess_release(handles: Iterable[str]) -> None: + for handle in handles: + _udfs.pop(handle, None) + + +def _nullable_type(data_type: pa.DataType) -> pa.DataType: + def nullable_field(field: pa.Field) -> pa.Field: + return pa.field(field.name, _nullable_type(field.type), nullable=True) + + if pa.types.is_struct(data_type): + return pa.struct([nullable_field(field) for field in data_type]) + if pa.types.is_list(data_type): + return pa.list_(nullable_field(data_type.value_field)) + if pa.types.is_large_list(data_type): + return pa.large_list(nullable_field(data_type.value_field)) + if pa.types.is_map(data_type): + return pa.map_( + _nullable_type(data_type.key_type), + nullable_field(data_type.item_field), + keys_sorted=data_type.keys_sorted, + ) + return data_type + + +def _null_checker(expected_type: pa.DataType) -> Optional[NullChecker]: + """Compile checks only for required fields and their ancestors, once per registration.""" + + def field_checker(field: pa.Field) -> Optional[NullChecker]: + nested = _null_checker(field.type) + if field.nullable: + return nested + + def check(values: pa.Array) -> None: + if values.null_count: + raise ValueError( + f"In-process UDF returned nulls in non-nullable field {field.name}" + ) + if nested is not None: + nested(values) + + return check + + if pa.types.is_struct(expected_type): + fields = [(i, field_checker(f)) for i, f in enumerate(expected_type)] + checks = [(i, check) for i, check in fields if check is not None] + if not checks: + return None + + def check_struct(array: pa.Array) -> None: + # Only children of valid parents are logically visible. + visible = pc.filter(array, pc.is_valid(array)) if array.null_count else array Review Comment: **Performance: a single null parent still copies the whole result.** This improved since the previous round. However, when a struct/map result has any top-level null, `pc.filter` materializes the whole result, including sibling children that aren't checked. The checker is always active for: - every `MapType` result, because map keys are always non-nullable, - Variant/Geometry/Geography results, because their children are non-nullable. Measured with PyArrow 25, one null in a 10,000-row batch costs, per batch: - about 124 µs for a Variant-like struct (1.28MB), - about 241 µs for `map<string,int64>` (0.74MB), - plus temporary memory the size of the result. A fast path would take about 0.5 µs. Also, PyArrow cannot construct a `MapArray` with null keys, so the key check only copies and can never fail. Cheaper alternatives: - Skip the check when the unfiltered child has `null_count == 0`, since it is a superset of the visible values. - Otherwise, check `pc.and_(pc.is_valid(parent), pc.is_null(child))` on bitmaps instead of filtering. ########## docs/sql-pyspark-inprocess-udf.md: ########## @@ -0,0 +1,621 @@ +--- +layout: global +title: In-Process Python UDFs +displayTitle: In-Process Python UDFs +license: | + Licensed to the Apache Software Foundation (ASF) under one or more + contributor license agreements. See the NOTICE file distributed with + this work for additional information regarding copyright ownership. + The ASF licenses this file to You under the Apache License, Version 2.0 + (the "License"); you may not use this file except in compliance with + the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +--- + +* Table of contents +{:toc} + +## Runtime and result contract + +Each executor owns a dedicated interpreter thread. The plugin initializes the +interpreter on that thread, and task calls and shutdown are dispatched to the +same thread. Calls from concurrent tasks are queued on the interpreter thread. +One task per executor is recommended for throughput, but is not a correctness requirement. +Application-level Python parallelism comes from multiple executor JVMs. +The plugin configures JEP's process-wide interpreter with hash seed `0`, matching +Spark's default Python worker seed. It must initialize before any other JEP user in +the JVM. The seed cannot change between SparkContexts in the same process; a custom +worker `PYTHONHASHSEED` does not override this embedded-runtime setting. + +Task cancellation cannot safely stop arbitrary native Python code. An interrupted +caller waits for the current invocation to finish before freeing the Arrow CDI +structures, then restores its interrupt status. A UDF that never returns can +therefore prevent its task from completing cancellation and block every subsequent +in-process UDF on that executor, including calls from other tasks, jobs, and sessions. +Recovery from a permanently hung invocation requires replacing the executor process. +Plugin shutdown stops accepting new calls and waits up to five seconds for the interpreter thread. If a call is +still running, cleanup stays queued behind it; its memory remains live until the +call returns or the process exits. Shutdown does not forcibly interrupt native +code. A new interpreter cannot start until the previous one has fully stopped. + +A scalar UDF must return a `pyarrow.Array` with exactly one element per input row. +The runtime checks the result type against the declared Spark type, including +nested fields, decimal scale, and timestamp unit/timezone. Value types must match +exactly: use an explicit PyArrow cast in the UDF for numeric or other conversions. +Nested field nullability may differ if the actual values satisfy the declared nullability. Sliced results, including nested +child slices, are copied to remove offsets that Arrow Java's CDI importer cannot +read. Compatible results retain zero-copy transfer. + +The API produces a regular `PythonUDF` expression with an in-process evaluation +type. Spark's existing `ArrowEvalPython` planning rules handle aggregation, +nested calls, nondeterminism, and filter/limit pushdown. `ArrowEvalPythonExec` selects +an in-process evaluator factory for this evaluation type, reusing the projection, +row queue, result join, and partition-evaluator path. Ordinary Python UDFs continue +to use Python workers. + +`maxRecordsPerBatch <= 0` means no row-count limit. The independent +`spark.sql.execution.arrow.maxBytesPerBatch` limit still applies when positive. +Only UDF arguments are converted to Arrow. Other columns stay in Spark rows, +buffered in a spillable queue until the results are joined back. Duplicate nested +field names in UDF arguments or declared results are rejected before Arrow Java +reads their buffers. + +Each batch uses fresh input buffers. A Python function may retain an input array; +later batches do not overwrite it. Retained arrays keep native memory alive, so +functions should release them when no longer needed. JVM input vectors and result +vectors are released on task completion, early termination and failure. + +UDF deserialization uses PySpark's bundled cloudpickle. Each task registers its +own function instance once and passes a small handle for subsequent batches. +Task completion queues release of the registered function and its closure state. Imported +Python modules still share executor-wide state. Extra site-packages paths are +processed with `site.addsitedir` before loading the runtime bridge, including `.pth` +files. Configured directories and newly discovered `.pth` paths precede system +paths. Already imported modules cannot be replaced by changing the search path. + +Spark broadcasts, accumulators, `SparkContext.addPyFile`, and Python `TaskContext` +are not supported by this embedded runtime. Captured broadcast and accumulator +objects are rejected during serialization; functions must not access them through +imported modules either. Install modules on executors before startup, optionally +using `spark.inprocess.python.sitePackages`. SQL registration through +`spark.udf.register` is not supported and is rejected at registration time. +Spark Connect does not support this execution mode; both client SQL registration and +server planning reject it. The decorator accepts a `DataType` or a DDL string; DDL +strings are parsed lazily with the active Spark session. It exposes `func`, `returnType`, +`evalType`, `deterministic`, and `asNondeterministic()` along with the function's name +and docstring. +Functions must receive at least one input column (a literal also works) to determine +the batch length. Positional and keyword arguments are supported. Functions are +serialized on first use, so globals can be defined or rebound after decoration +and before that first call. The driver's Python major.minor +version must match the embedded interpreter; registration checks this before +unpickling. Python exceptions, including `SystemExit` during deserialization or +execution, are converted into task failures. Tracebacks honor the query's +`spark.sql.execution.pyspark.udf.hideTraceback.enabled` and +`spark.sql.execution.pyspark.udf.simplifiedTraceback.enabled` settings. Native process +termination remains outside this exception handling. + +## Overview + +In-process Python UDFs embed CPython directly into the Spark executor JVM using +[jep (Java Embedded Python)](https://github.com/ninia/jep), eliminating the IPC overhead of +standard Python UDFs and pandas UDFs. Data is passed to Python as +[PyArrow](https://arrow.apache.org/docs/python/) arrays via the +[Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html) — zero-copy +for compatible input and output buffers. Row-to-Arrow conversion and normalization +of sliced results still copy data. + +**Use `inprocess_udf` when:** +- You are already using `pandas_udf` for vectorized transformations and want lower latency. +- Your UDF operates on Arrow/PyArrow arrays (e.g. using `pyarrow.compute`). +- You can deploy enough executor JVMs for Python parallelism (see [Requirements](#requirements)). + +**Stick with `pandas_udf` or `udf` when:** +- You need pandas Series semantics in your UDF logic. +- You need concurrent Python invocations within a single executor. +- You are not able to install jep on executors. + +--- + +## Quick Start + +### 1. Install dependencies + +```bash +pip install "jep>=4.3.2" pyarrow cloudpickle +``` + +JEP and `org.apache.arrow:arrow-c-data` are provided dependencies and are not +bundled with Spark. Supply their JARs on the driver/executor classpaths before +starting Spark, and make the JEP native library available. Use an `arrow-c-data` +version matching Spark's Arrow Java version. Installing the Python packages alone +does not supply the Arrow Java CDI JAR. + +Building JEP from source requires a JDK, a C compiler, and development headers for +the Python version being embedded (for example, `python3.12-dev` on Ubuntu with +Python 3.12). These headers are build dependencies; running a prebuilt compatible +JEP installation does not require the development package. The corresponding +Python shared library must remain available at runtime. + +### 2. Register the plugin + +```python +spark = SparkSession.builder \ + .config("spark.plugins", + "org.apache.spark.sql.execution.python.InProcessPythonPlugin") \ + .config("spark.executor.cores", "1") \ + .config("spark.task.cpus", "1") \ + .getOrCreate() +``` + +### 3. Write and call a UDF + +```python +import pyarrow.compute as pc +from pyspark.inprocess.udf import inprocess_udf +from pyspark.sql.types import LongType + +@inprocess_udf(return_type=LongType()) +def double(x): + return pc.multiply(x, 2) + +df = spark.range(10) +df.select(double(df["id"])).show() +``` + +The function receives a `pa.Array` for each input column and must return a `pa.Array`. + +--- + +## Examples + +### String transformation + +```python +import pyarrow.compute as pc +from pyspark.inprocess.udf import inprocess_udf +from pyspark.sql.types import StringType + +@inprocess_udf(return_type=StringType()) +def upper(s): + return pc.utf8_upper(s) + +df = spark.createDataFrame([("hello",), ("world",)], ["text"]) +df.select(upper(df["text"])).show() +# +------------+ +# |upper(text) | +# +------------+ +# |HELLO | +# |WORLD | +# +------------+ +``` + +### Multi-column UDF + +A UDF receives one `pa.Array` argument per input column: + +```python +import pyarrow.compute as pc +from pyspark.inprocess.udf import inprocess_udf +from pyspark.sql.types import DoubleType + +@inprocess_udf(return_type=DoubleType()) +def weighted_sum(x, y): + return pc.add(pc.multiply(x, 0.6), pc.multiply(y, 0.4)) + +df = spark.createDataFrame([(1.0, 2.0), (3.0, 4.0)], ["x", "y"]) +df.select(weighted_sum(df["x"], df["y"])).show() +``` + +### Closure capture + +Free variables are captured by cloudpickle and frozen into the serialized UDF. The captured +value is evaluated once at UDF definition time and shipped with the function to every executor: Review Comment: **Docs: this contradicts the lazy serialization.** This says the captured value is evaluated once at UDF definition time. However, `InProcessUDFWrapper` serializes on the first `__call__` and caches the result (`_serialize`), as lines 95-97 of this doc also say. For example: - rebinding `SCALE_FACTOR` after decoration but before the first call is picked up, - rebinding it after the first call is ignored for all later DataFrames. Could we say "at first use" here? -- 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]
