docs: reframe distributed-expression docs around the user goal

User-facing docs throughout this PR led with "pickle support":
filename-shaped headings, function docstrings describing how things
get cloudpickled into a Rust-side codec, etc. That's the
implementation pathway, not the user's goal.

The user's goal is to build an expression in a driver process and
ship it to worker processes for distributed evaluation. Pickle is the
mechanism Python provides to make that work; we hook into it. End
users typically don't care how the bytes are produced — they care
which references survive the trip and what they have to register on
each worker.

Reframe across user-facing surfaces:

* `docs/source/user-guide/io/distributing_expressions.rst` — leads
  with the worker-pool use case, drops `PythonUDFCodec` /
  cloudpickle vocabulary, presents "what travels with the
  expression" as the user contract.
* `datafusion.ipc` module docstring + `set_worker_ctx` /
  `clear_worker_ctx` / `get_worker_ctx` — describes what the user
  installs and why, not internal lookup details.
* `Expr.to_bytes` / `from_bytes` / `__reduce__` — describes what's
  shipped vs what travels by name; cross-references the user guide
  instead of repeating the codec story.
* `examples/ray_pickle_expr.py` header + comment + README entry —
  goal-first wording.
* Pickle test module docstrings — drop the dangling reference to
  `PythonUDFCodec` (also a stale name post-PR1).

Code behavior unchanged. 1088 tests still green.

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

Author: timsaucer

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

@@ -18,41 +18,36 @@
 Distributing expressions across processes
 =========================================
 
-DataFusion expressions (:py:class:`~datafusion.Expr`) can be serialized and
-shipped across process boundaries — useful for distributing work over a
-``multiprocessing.Pool``, a Ray actor pool, or any framework that supports a
-per-worker initialization hook.
+A common pattern is to build a DataFusion expression
+(:py:class:`~datafusion.Expr`) in a driver process, hand it to a pool of
+worker processes (``multiprocessing.Pool``, a Ray actor pool, or any other
+framework with a per-worker initialization hook), and have each worker
+evaluate the expression against its own slice of data.
 
-Pickle support
---------------
+DataFusion expressions support this directly: they can be sent through
+:py:mod:`pickle` like any other Python object. Python scalar UDFs ride along
+inside the pickled bytes — the receiver does not need to pre-register them.
 
-:py:class:`~datafusion.Expr` implements the pickle protocol directly. Call
-:py:func:`pickle.dumps` on an expression and ship the bytes; the receiver
-calls :py:func:`pickle.loads`. Python *scalar UDFs* are cloudpickled into the
-proto wire format by a Rust-side codec (``PythonUDFCodec``), so the blob is
-self-contained — the receiver does not need to pre-register the UDF.
+Basic worker-pool example
+-------------------------
 
 .. code-block:: python
 
     import multiprocessing as mp
     import pickle
 
     import pyarrow as pa
-    from datafusion import SessionContext, col, lit, udf
+    from datafusion import SessionContext, col, udf
 
-    def init_worker():
-        # Optional: install a worker context for aggregate / window UDFs,
-        # table providers, or Rust-side function registrations. Not needed
-        # for built-ins or Python scalar UDFs.
-        pass
 
     def evaluate(blob_and_batch):
         blob, batch = blob_and_batch
-        expr = pickle.loads(blob)
+        expr = pickle.loads(blob)  # Python scalar UDFs ride along inline.
         ctx = SessionContext()
         df = ctx.from_pydict({"a": batch})
         return df.with_column("result", expr).select("result").to_pydict()["result"]
 
+
     if __name__ == "__main__":
         double = udf(
             lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),
@@ -68,74 +63,92 @@ self-contained — the receiver does not need to pre-register the UDF.
             )
         print(results)  # [[2, 4, 6], [20, 40, 60]]
 
-Worker-scoped context
----------------------
 
-For references the codec cannot inline — aggregate UDFs, window UDFs, FFI
-capsule UDFs, or anything resolved through the
-:class:`SessionContext`'s function registry — set a worker-scoped context
-once per process using :py:func:`datafusion.ipc.set_worker_ctx`:
+What travels with the expression
+--------------------------------
+
+* **Built-in functions** (``abs``, ``length``, arithmetic, comparisons, etc.)
+  — fully portable. Worker needs nothing pre-registered.
+* **Python scalar UDFs** (defined with :py:func:`datafusion.udf`) — fully
+  portable. The callable and its signature travel inside the pickled bytes
+  and are reconstructed on the worker automatically.
+* **Aggregate UDFs**, **window UDFs**, **UDFs imported via the FFI capsule
+  protocol** — travel **by name only**. The worker must already have a
+  matching registration on its :py:class:`SessionContext`. Without that
+  registration, evaluation raises an error.
+
+Registering shared UDFs on workers
+----------------------------------
+
+When an expression references something that travels by name only (aggregate
+UDF, window UDF, FFI UDF), set up the worker's :py:class:`SessionContext`
+once per process and install it as the *worker context*:
 
 .. code-block:: python
 
     from datafusion import SessionContext
     from datafusion.ipc import set_worker_ctx
 
+
     def init_worker():
         ctx = SessionContext()
-        ctx.register_udaf(my_aggregate)  # if needed
+        ctx.register_udaf(my_aggregate)
         set_worker_ctx(ctx)
 
+
     with mp.get_context("forkserver").Pool(
         processes=4, initializer=init_worker
     ) as pool:
         ...
 
-Without a worker context, unpickling falls back to a fresh
-:py:class:`SessionContext`. Built-in functions resolve; Python scalar UDFs
-ride along inside the blob via the codec. References to aggregate / window
-UDFs or other registry-only entries raise an informative error if not
-registered on the worker.
+Inside a worker, expressions reconstructed by :py:func:`pickle.loads` 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 scalar UDFs,
+but anything by-name-only will fail to resolve.
 
 Python 3.14 default change
 --------------------------
 
 Python 3.14 changed the POSIX default start method for
-:py:mod:`multiprocessing` from ``fork`` to ``forkserver``. With ``fork``, a
-context set in the parent was visible in workers via copy-on-write; with
-``forkserver`` and ``spawn`` it is not. The codec + worker-init pattern works
-on every start method — prefer it over relying on inherited state.
-
-Trade-offs of inline UDFs
--------------------------
-
-* **Blob size** — cloudpickled callables add bytes per blob. A trivial
-  built-in expression is ~20 bytes; an expression referencing a Python scalar
-  UDF is hundreds of bytes (the cloudpickled callable + signature). Pre-register
-  shared UDFs on workers via :py:func:`~datafusion.ipc.set_worker_ctx` when
-  the same UDF is shipped many times and you want to avoid the overhead.
-* **Closure capture** — cloudpickle captures closure state. Surprises are
-  possible if the UDF closes over large objects, module-level mutable state,
-  or non-portable file paths.
-* **FFI scalar UDFs cannot be inlined** — PyCapsule-backed UDFs have no
-  Python callable to cloudpickle. The codec leaves their ``fun_definition``
-  empty; the receiver must have a matching registration.
-* **Aggregate and window UDFs cannot be inlined yet** — their Python state
-  is held inside opaque factory closures on the Rust side. Pre-register on
-  the worker.
+:py:mod:`multiprocessing` from ``fork`` to ``forkserver``. With ``fork``, any
+state set in the parent was visible in workers via copy-on-write; with
+``forkserver`` and ``spawn`` it is not. The
+:py:func:`~datafusion.ipc.set_worker_ctx` pattern works on every start
+method — prefer it over relying on inherited state.
+
+Practical considerations
+------------------------
+
+* **Pickled size scales with what travels inline.** A pickled expression of
+  just built-ins is small (tens of bytes). An expression carrying a Python
+  scalar UDF is hundreds of bytes (the callable and its signature). When the
+  same UDF is shipped many times, pre-registering it on each worker via
+  :py:func:`~datafusion.ipc.set_worker_ctx` and referring to it by name
+  cuts the per-blob overhead.
+* **Closure capture.** When a Python scalar UDF closes over surrounding
+  state — local variables, module-level objects, file paths — that state
+  is captured at pickling time. Surprises are possible if the captured
+  state is large, mutable, or not portable to the worker's environment.
+* **Aggregate and window UDFs always travel by name.** Their Python state
+  is held inside opaque factory closures that cannot be reconstructed from
+  bytes alone. Use :py:func:`~datafusion.ipc.set_worker_ctx` to register
+  them on each worker.
 
 Security
 --------
 
 .. warning::
 
-   Pickle blobs containing inlined UDFs deserialize via :py:mod:`cloudpickle`,
-   which executes arbitrary code on the receiver. Only :py:func:`pickle.loads`
-   blobs from trusted sources. For untrusted-source workflows, restrict the
-   sender to built-in functions and pre-registered Rust-side UDFs.
+   Reconstructing an expression containing a Python scalar UDF executes
+   arbitrary Python code on the receiver. Only :py:func:`pickle.loads`
+   expressions from trusted sources. For untrusted-source workflows,
+   restrict senders to built-in functions and pre-registered Rust-side
+   UDFs, and never feed externally supplied bytes through
+   :py:func:`pickle.loads`.
 
 See also
 --------
 
-* :py:mod:`datafusion.ipc` — module-level API reference.
+* :py:mod:`datafusion.ipc` — worker context API.
 * ``examples/ray_pickle_expr.py`` — runnable Ray actor example.

examples/README.md

@@ -46,7 +46,7 @@ Here is a direct link to the file used in the examples:
 
 ### Distributing DataFusion expressions
 
-- [Pickle expressions and send them to Ray actors](./ray_pickle_expr.py)
+- [Distribute expression evaluation across Ray actors](./ray_pickle_expr.py)
 
 ### Substrait Support
 

examples/ray_pickle_expr.py

@@ -17,13 +17,11 @@
 
 """Distribute DataFusion expressions to Ray actors.
 
-This example shows the worker-init pattern from the user guide adapted to
-Ray's actor model: each actor builds its own :class:`SessionContext`.
-Python scalar UDFs travel inside the pickle blob via the Rust-side
-``PythonUDFCodec`` — no actor-side pre-registration is required. The
-worker context (set via :func:`datafusion.ipc.set_worker_ctx`) is still
-useful for aggregate/window UDFs or other registry-only entries; we set
-one here to show the pattern.
+Build an expression in the driver, ship it to a pool of Ray actors, and have
+each actor evaluate it against its own slice of data. Each actor sets up
+its own :class:`SessionContext` once in `__init__` and registers any UDFs
+it needs to resolve by name. Python scalar UDFs travel with the shipped
+expression and need no actor-side pre-registration.
 
 Prerequisites:
     pip install ray
@@ -58,8 +56,9 @@ class DataFusionWorker:
     def __init__(self) -> None:
         ctx = SessionContext()
         ctx.register_udf(_build_double_udf())
-        # The worker context is what Expr.__setstate__ consults when
-        # pickled expressions arrive at this actor.
+        # Install the actor's SessionContext as its worker context;
+        # expressions reconstructed in this actor will resolve their
+        # by-name references against it.
         set_worker_ctx(ctx)
         self._ctx = ctx
 

python/datafusion/expr.py

@@ -434,30 +434,32 @@ def variant_name(self) -> str:
         return self.expr.variant_name()
 
     def to_bytes(self, ctx: SessionContext | None = None) -> bytes:
-        """Serialize this expression to protobuf bytes.
+        """Serialize this expression to bytes for shipping to another process.
 
-        Python scalar UDFs are inlined into the returned bytes — the
-        receiver does not need to pre-register them. Aggregate UDFs,
-        window UDFs, and UDFs imported via the FFI capsule protocol
-        are stored by name only; the receiver must have them
-        registered.
+        Use this — or :func:`pickle.dumps` — to send an expression to a
+        worker process for distributed evaluation.
 
-        When ``ctx`` is supplied, encoding also routes through the
-        session's installed codec stack.
+        Built-in functions and Python scalar UDFs travel inside the
+        returned bytes; the worker does not need to pre-register them.
+        Aggregate UDFs, window UDFs, and UDFs imported via the FFI
+        capsule protocol travel by name only and must be registered on
+        the worker. See :doc:`/user-guide/io/distributing_expressions`.
         """
         ctx_arg = ctx.ctx if ctx is not None else None
         return bytes(self.expr.to_bytes(ctx_arg))
 
     @classmethod
     def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr:
-        """Decode an expression from serialized protobuf bytes.
-
-        ``ctx`` is the receiver :class:`SessionContext` used to resolve
-        function references not inlined by the codec (aggregate UDFs,
-        window UDFs, FFI UDFs). When ``ctx`` is ``None`` the worker
-        context set via :func:`datafusion.ipc.set_worker_ctx` is
-        consulted; if no worker context is set, a fresh
-        :class:`SessionContext` is used.
+        """Reconstruct an expression from serialized bytes.
+
+        Accepts output of :meth:`to_bytes` or :func:`pickle.dumps`.
+        ``ctx`` is the :class:`SessionContext` used to resolve any
+        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
+        :class:`SessionContext` is used (sufficient for built-ins and
+        Python scalar UDFs).
         """
         from datafusion.ipc import _resolve_ctx
 
@@ -467,14 +469,12 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr:
     def __reduce__(self) -> tuple:
         """Pickle protocol hook.
 
-        Python scalar UDFs referenced by the expression are inlined
-        into the pickle blob, so the receiver does not need to
-        pre-register them. On unpickle the bytes are decoded against
-        the worker context set via
-        :func:`datafusion.ipc.set_worker_ctx` (or a fresh
-        :class:`SessionContext` if none) for any registry-resolved
-        references — aggregate UDFs, window UDFs, UDFs imported via
-        the FFI capsule protocol.
+        Lets expressions be shipped to worker processes via
+        :func:`pickle.dumps` / :func:`pickle.loads`. The worker's
+        :class:`SessionContext` for resolving by-name references is
+        looked up via :func:`datafusion.ipc.set_worker_ctx`, falling
+        back to a fresh empty :class:`SessionContext` if none has been
+        installed on the worker.
         """
         return (Expr._reconstruct, (self.to_bytes(),))
 

python/datafusion/ipc.py

@@ -15,28 +15,29 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Inter-process communication helpers for distributing DataFusion expressions.
+"""Worker-side setup for distributing DataFusion expressions.
 
-This module provides a worker-scoped :class:`SessionContext` slot that
-:meth:`Expr.__reduce__` consults when unpickling expressions across process
-boundaries. Set the worker context once per worker process (typically from a
-``multiprocessing.Pool`` initializer or a Ray actor ``__init__``):
+When a :class:`Expr` is shipped to a worker process (e.g. through
+:func:`multiprocessing.Pool` or a Ray actor), the worker reconstructs the
+expression against a :class:`SessionContext`. If the expression references
+aggregate UDFs, window UDFs, table providers, or UDFs imported via the FFI
+capsule protocol — anything the worker would otherwise resolve from its
+registered functions — install a configured :class:`SessionContext` once
+per worker:
 
 >>> # doctest: +SKIP
 >>> from datafusion import SessionContext
 >>> from datafusion.ipc import set_worker_ctx
 >>>
 >>> def init_worker():
 ...     ctx = SessionContext()
-...     # register Rust-backed UDFs / aggregates / window functions here
+...     ctx.register_udaf(my_aggregate)
 ...     set_worker_ctx(ctx)
 
-Python scalar UDFs do not need pre-registration: their definitions
-travel inside the pickled expression and are reconstructed on the
-receiver automatically. The worker context is only needed when the
-expression references aggregate UDFs, window UDFs, table providers,
-or UDFs imported via the FFI capsule protocol — anything the
-receiver would otherwise resolve from its registered functions.
+Built-in functions and Python scalar UDFs travel inside the shipped
+expression itself and do not need pre-registration on the worker.
+
+See :doc:`/user-guide/io/distributing_expressions` for the full pattern.
 """
 
 from __future__ import annotations
@@ -59,25 +60,30 @@
 
 
 def set_worker_ctx(ctx: SessionContext) -> None:
-    """Register the receiver :class:`SessionContext` for this worker.
-
-    Call once per worker process — typically from a ``Pool`` initializer or a
-    Ray actor ``__init__``. Idempotent: overwrites any previous value.
+    """Install this worker's :class:`SessionContext` for shipped expressions.
 
-    The worker context is stored in a thread-local slot, so each thread within
-    a worker can install its own context independently.
+    Call once per worker — typically from a ``multiprocessing.Pool``
+    initializer or a Ray actor ``__init__``. Idempotent: overwrites any
+    previous value. Stored in a thread-local slot, so each thread within a
+    worker may install its own context independently.
     """
     _local.ctx = ctx
 
 
 def clear_worker_ctx() -> None:
-    """Remove the worker context, restoring fresh-context fallback behavior."""
+    """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
+    scalar UDFs, but anything that travels by name only (aggregate UDFs,
+    window UDFs, FFI UDFs) will fail to resolve.
+    """
     if hasattr(_local, "ctx"):
         del _local.ctx
 
 
 def get_worker_ctx() -> SessionContext | None:
-    """Return the worker context if set, else ``None``."""
+    """Return this worker's installed :class:`SessionContext`, or ``None``."""
     return getattr(_local, "ctx", None)
 
 

python/tests/test_pickle_expr.py

@@ -17,10 +17,11 @@
 
 """In-process pickle round-trip tests for :class:`Expr`.
 
-The Rust-side ``PythonUDFCodec`` cloudpickles Python scalar UDF callables
-directly into the proto wire format, so pickle blobs are self-contained.
-The worker context (:mod:`datafusion.ipc`) is only needed for references
-the codec can't inline — aggregate UDFs, window UDFs, FFI capsule UDFs.
+Built-in functions and Python scalar UDFs travel with the pickled
+expression and do not need worker-side pre-registration. The worker
+context (:mod:`datafusion.ipc`) is only consulted for references that
+travel by name — aggregate UDFs, window UDFs, UDFs imported via the FFI
+capsule protocol.
 
 Cross-process tests live in ``test_pickle_multiprocessing.py``.
 """