feat: sender-side session context for outbound pickles

`SessionContext.with_python_udf_inlining(False)` previously had no
effect on `pickle.dumps(expr)` because `__reduce__` called
`to_bytes()` without a context. Add a thread-local sender slot
(`datafusion.ipc.set_sender_ctx`) that `__reduce__` consults, so the
inlining toggle flows through pickle. Symmetric to the existing
worker slot.

Also switch the worker-side decode fallback from a freshly
constructed `SessionContext` to `SessionContext.global_ctx()` so
expressions arriving with no explicit/worker context can still
resolve names registered on the global singleton.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: timsaucer

docs/source/user-guide/io/distributing_work.rst

@@ -116,6 +116,85 @@ What travels with the expression
   :py:class:`SessionContext`. Without that registration, evaluation
   raises an error.
 
+Session contexts at a glance
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+There is only one type — :py:class:`SessionContext`. It can occupy
+up to four *slots* in a running program:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 12 18 40 30
+
+   * - Slot
+     - Lifetime
+     - Purpose
+     - Set how
+   * - User-held
+     - Local variable / attribute
+     - Build and run queries
+     - ``ctx = SessionContext(...)``
+   * - Global
+     - Process singleton (lazy-init)
+     - Backs module-level
+       :py:func:`~datafusion.io.read_parquet`,
+       :py:func:`~datafusion.io.read_csv`,
+       :py:func:`~datafusion.io.read_json`,
+       :py:func:`~datafusion.io.read_avro`; final fallback for
+       :py:meth:`Expr.from_bytes`
+     - Implicit; access via
+       :py:meth:`SessionContext.global_ctx`
+   * - Sender
+     - Thread-local on the driver
+     - Codec settings for outbound :py:func:`pickle.dumps` /
+       :py:meth:`Expr.to_bytes` without ``ctx``
+     - :py:func:`~datafusion.ipc.set_sender_ctx`
+   * - Worker
+     - Thread-local on the worker
+     - Function registry for inbound :py:func:`pickle.loads` /
+       :py:meth:`Expr.from_bytes` without ``ctx``
+     - :py:func:`~datafusion.ipc.set_worker_ctx`
+
+The same :py:class:`SessionContext` object may occupy more than one
+slot simultaneously — installing it into a slot is a reference, not
+a copy.
+
+**Non-distributed user.** One user-held context. The global slot is
+invisible unless you call top-level ``read_*`` helpers. Sender and
+worker slots are unused.
+
+**Distributed user.** Two questions to answer:
+
+1. *Driver side — what wire format do I want?* The default (Python UDF
+   inlining on) is self-contained; you do not need a sender context.
+   To opt into the strict format,
+   :py:func:`~datafusion.ipc.set_sender_ctx`
+   with a session built via
+   :py:meth:`SessionContext.with_python_udf_inlining(False)
+   <datafusion.SessionContext.with_python_udf_inlining>`.
+
+2. *Worker side — what registrations does decode need?* For built-ins
+   and inline Python UDFs, nothing. For FFI-capsule UDFs (or
+   strict-mode round-trips that travel by name), call
+   :py:func:`~datafusion.ipc.set_worker_ctx` once per worker with a
+   context that has the relevant registrations.
+
+Resolution order on the worker side is *explicit argument →
+worker context → global context.* Explicit ``ctx=`` on
+:py:meth:`Expr.from_bytes` always wins; the sender slot is ignored
+on decode and the worker slot is ignored on encode.
+
+Sharp edges:
+
+* Sender and worker slots are **thread-local**. Background threads
+  on either side see ``None`` until they install their own.
+* The global slot persists across ``fork`` workers (copy-on-write
+  memory inherit) but not across ``spawn`` / ``forkserver`` workers
+  (fresh process — register or install a worker context on
+  start-up).
+* The inlining toggle is per-context state, not a global switch.
+  Two contexts with different toggles can coexist in one process.
+
 Registering shared UDFs on workers
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -143,9 +222,10 @@ as the *worker context*:
 
 Inside a worker, expressions arriving from the driver resolve their
 by-name references against the installed worker context. If no worker
-context is installed, a fresh empty :py:class:`SessionContext` is
-used — fine for expressions that only reference built-ins and Python
-UDFs, but FFI-capsule-backed registrations will fail to resolve.
+context is installed, the global :py:class:`SessionContext` is used —
+fine for expressions that only reference built-ins and Python UDFs,
+but FFI-capsule-backed registrations must be installed on the global
+context to resolve.
 
 Python 3.14 default change
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -198,6 +278,27 @@ Mismatched configurations raise a descriptive error: an inline blob
 fed to a strict receiver fails fast rather than silently dropping
 into ``cloudpickle.loads``.
 
+To make the toggle apply through :py:func:`pickle.dumps` (which
+calls :py:meth:`Expr.to_bytes` with no context), install the strict
+session as the driver's *sender context*:
+
+.. code-block:: python
+
+    from datafusion import SessionContext
+    from datafusion.ipc import set_sender_ctx
+
+    set_sender_ctx(SessionContext().with_python_udf_inlining(False))
+    # Every subsequent pickle.dumps(expr) on this thread encodes
+    # without inlining the Python callable.
+
+Pair with a matching strict worker context
+(:py:func:`~datafusion.ipc.set_worker_ctx`) so the ``pickle.loads``
+side also refuses inline payloads. Explicit
+:py:meth:`Expr.to_bytes(ctx) <Expr.to_bytes>` and
+:py:meth:`Expr.from_bytes(blob, ctx=ctx) <Expr.from_bytes>` calls
+honor the supplied ``ctx`` directly and ignore the sender / worker
+contexts.
+
 Note that :py:func:`pickle.loads` itself remains unsafe on untrusted
 input regardless of this setting — an attacker producing the outer
 pickle envelope can execute arbitrary code before the codec ever

python/datafusion/context.py

@@ -1786,6 +1786,16 @@ def with_python_udf_inlining(self, enabled: bool) -> SessionContext:
           :meth:`Expr.from_bytes` input that may come from an untrusted
           source — ``cloudpickle.loads`` will not be invoked.
 
+        The toggle applies directly to :meth:`Expr.to_bytes` /
+        :meth:`Expr.from_bytes` calls that pass this session as their
+        ``ctx`` argument. To make the toggle apply through
+        :func:`pickle.dumps` (which calls :meth:`Expr.to_bytes` with no
+        context), install this session as the driver's sender context
+        via :func:`datafusion.ipc.set_sender_ctx` — and install it as
+        the worker's context via
+        :func:`datafusion.ipc.set_worker_ctx` for the corresponding
+        :func:`pickle.loads`.
+
         ``pickle.loads`` on untrusted bytes remains unsafe regardless of
         this setting (see the `pickle module security warning
         <https://docs.python.org/3/library/pickle.html#module-pickle>`_

python/datafusion/expr.py

@@ -457,9 +457,10 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr:
         function references that travel by name — aggregate UDFs, window
         UDFs, FFI UDFs. When ``ctx`` is ``None`` the worker context
         installed via :func:`datafusion.ipc.set_worker_ctx` is consulted;
-        if no worker context is installed, a fresh
+        if no worker context is installed, the global
         :class:`SessionContext` is used (sufficient for built-ins and
-        Python scalar UDFs).
+        Python scalar UDFs, plus any UDFs registered on the global
+        context).
         """
         from datafusion.ipc import _resolve_ctx
 
@@ -475,10 +476,17 @@ def __reduce__(self) -> tuple:
         UDFs require pre-registration on the worker. The worker's
         :class:`SessionContext` for resolving those references is
         looked up via :func:`datafusion.ipc.set_worker_ctx`, falling
-        back to a fresh empty :class:`SessionContext` if none has been
+        back to the global :class:`SessionContext` if none has been
         installed on the worker.
+
+        The encoding side honors a driver-side sender context installed
+        via :func:`datafusion.ipc.set_sender_ctx` — that is how
+        :meth:`SessionContext.with_python_udf_inlining` propagates
+        through ``pickle.dumps``.
         """
-        return (Expr._reconstruct, (self.to_bytes(),))
+        from datafusion.ipc import get_sender_ctx
+
+        return (Expr._reconstruct, (self.to_bytes(get_sender_ctx()),))
 
     @classmethod
     def _reconstruct(cls, proto_bytes: bytes) -> Expr:

python/datafusion/ipc.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Worker-side setup for distributing DataFusion expressions.
+"""Driver- and worker-side setup for distributing DataFusion expressions.
 
 When a :class:`Expr` is shipped to a worker process (e.g. through
 :func:`multiprocessing.Pool` or a Ray actor), the worker reconstructs the
@@ -38,6 +38,24 @@
 inside the shipped expression itself and do not need pre-registration
 on the worker.
 
+On the driver side, call :func:`set_sender_ctx` to control how
+:func:`pickle.dumps` encodes expressions — for example, to apply
+:meth:`SessionContext.with_python_udf_inlining` to every pickled
+expression on this thread:
+
+>>> # doctest: +SKIP
+>>> from datafusion import SessionContext
+>>> from datafusion.ipc import set_sender_ctx
+>>>
+>>> driver_ctx = SessionContext().with_python_udf_inlining(False)
+>>> set_sender_ctx(driver_ctx)
+>>> pickle.dumps(expr)  # encoded with inlining disabled
+
+Without a sender context the default codec is used (Python UDF
+inlining on). The sender context only affects pickle / ``to_bytes``
+encoding; explicit ``expr.to_bytes(ctx)`` calls still use the supplied
+``ctx``.
+
 See :doc:`/user-guide/io/distributing_work` for the full pattern.
 """
 
@@ -51,8 +69,11 @@
 
 
 __all__ = [
+    "clear_sender_ctx",
     "clear_worker_ctx",
+    "get_sender_ctx",
     "get_worker_ctx",
+    "set_sender_ctx",
     "set_worker_ctx",
 ]
 
@@ -74,10 +95,10 @@ def set_worker_ctx(ctx: SessionContext) -> None:
 def clear_worker_ctx() -> None:
     """Remove this worker's installed :class:`SessionContext`.
 
-    After clearing, expressions reconstructed in this worker fall back to a
-    fresh empty :class:`SessionContext` — adequate for built-ins and Python
+    After clearing, expressions reconstructed in this worker fall back to
+    the global :class:`SessionContext` — adequate for built-ins and Python
     UDFs (scalar, aggregate, window), but anything imported via the FFI
-    capsule protocol will fail to resolve.
+    capsule protocol must be registered on the global context to resolve.
     """
     if hasattr(_local, "ctx"):
         del _local.ctx
@@ -88,12 +109,48 @@ def get_worker_ctx() -> SessionContext | None:
     return getattr(_local, "ctx", None)
 
 
+def set_sender_ctx(ctx: SessionContext) -> None:
+    """Install this driver's :class:`SessionContext` for outbound pickles.
+
+    Controls how :func:`pickle.dumps` encodes :class:`Expr` instances on
+    this thread. The most useful application is propagating a session
+    configured with
+    :meth:`SessionContext.with_python_udf_inlining` so the toggle takes
+    effect through pickle (which otherwise calls
+    :meth:`Expr.to_bytes` with no context and uses the default codec).
+
+    Idempotent: overwrites any previous value. Stored in a thread-local
+    slot, so worker threads on the driver may install their own contexts.
+    Does not affect :meth:`Expr.to_bytes` calls that pass an explicit
+    ``ctx`` — those continue to use the supplied context.
+    """
+    _local.sender_ctx = ctx
+
+
+def clear_sender_ctx() -> None:
+    """Remove this driver's installed sender :class:`SessionContext`.
+
+    After clearing, pickled expressions fall back to the default codec
+    (Python UDF inlining on).
+    """
+    if hasattr(_local, "sender_ctx"):
+        del _local.sender_ctx
+
+
+def get_sender_ctx() -> SessionContext | None:
+    """Return this driver's installed sender :class:`SessionContext`, or ``None``."""
+    return getattr(_local, "sender_ctx", None)
+
+
 def _resolve_ctx(
     explicit_ctx: SessionContext | None = None,
 ) -> SessionContext:
     """Resolve a context for Expr reconstruction.
 
-    Priority: explicit argument > worker context > fresh context.
+    Priority: explicit argument > worker context > global context.
+    Falling back to the global :class:`SessionContext` (instead of a
+    freshly constructed one) preserves any registrations the user has
+    installed on it.
     """
     if explicit_ctx is not None:
         return explicit_ctx
@@ -102,4 +159,4 @@ def _resolve_ctx(
         return worker
     from datafusion.context import SessionContext  # noqa: PLC0415
 
-    return SessionContext()
+    return SessionContext.global_ctx()