docs(pickle): add cloudpickle security warnings, docstring examples, edge-case tests

Inline `.. warning::` blocks on `Expr.to_bytes`, `Expr.from_bytes`, and
`Expr.__reduce__` so the cloudpickle / arbitrary-code-execution caveat is
visible at the public API surface in advance of the user-guide page that
lands in PR 4.

Add doctest-style `Examples:` blocks to `datafusion.ipc` functions
(`set_worker_ctx`, `clear_worker_ctx`, `get_worker_ctx`, `_resolve_ctx`),
`ScalarUDF.name`, and the new `Expr` pickle methods, per CLAUDE.md.

Tighten `Expr.__reduce__` return annotation to
`tuple[Callable[[bytes], Expr], tuple[bytes]]`.

Tests: multi-arg UDF round-trip (covers synthetic `arg_{i}` schema-field
loop in the codec) plus malformed-bytes paths through `Expr.from_bytes`.

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

Author: timsaucer

python/datafusion/expr.py

@@ -46,7 +46,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import pyarrow as pa
@@ -447,6 +447,19 @@ def to_bytes(self, ctx: SessionContext | None = None) -> bytes:
         returned bytes; the worker does not need to pre-register them.
         UDFs imported via the FFI capsule protocol travel by name only
         and must be registered on the worker.
+
+        .. warning::
+            Bytes returned here may embed a cloudpickled Python
+            callable (when the expression carries a Python scalar UDF).
+            Reconstructing them via :meth:`from_bytes` or
+            :func:`pickle.loads` executes arbitrary Python on the
+            receiver. Only accept payloads from trusted sources.
+
+        Examples:
+            >>> from datafusion import col, lit
+            >>> blob = (col("a") + lit(1)).to_bytes()
+            >>> isinstance(blob, bytes)
+            True
         """
         ctx_arg = ctx.ctx if ctx is not None else None
         return self.expr.to_bytes(ctx_arg)
@@ -463,13 +476,25 @@ def from_bytes(cls, buf: bytes, ctx: SessionContext | None = None) -> Expr:
         context is installed, the global :class:`SessionContext` is used
         (sufficient for built-ins and Python scalar UDFs, plus any UDFs
         registered on the global context).
+
+        .. warning::
+            Decoding may invoke ``cloudpickle.loads`` on bytes embedded
+            in the payload, which executes arbitrary Python code. Treat
+            ``buf`` as code, not data — only decode bytes you produced
+            yourself or received from a trusted sender.
+
+        Examples:
+            >>> from datafusion import Expr, col, lit
+            >>> blob = (col("a") + lit(1)).to_bytes()
+            >>> Expr.from_bytes(blob).canonical_name()
+            'a + Int64(1)'
         """
         from datafusion.ipc import _resolve_ctx
 
         resolved = _resolve_ctx(ctx)
         return cls(expr_internal.RawExpr.from_bytes(resolved.ctx, buf))
 
-    def __reduce__(self) -> tuple:
+    def __reduce__(self) -> tuple[Callable[[bytes], Expr], tuple[bytes]]:
         """Pickle protocol hook.
 
         Lets expressions be shipped to worker processes via
@@ -480,12 +505,32 @@ def __reduce__(self) -> tuple:
         is looked up via :func:`datafusion.ipc.set_worker_ctx`, falling
         back to the global :class:`SessionContext` if none has been
         installed on the worker.
+
+        .. warning::
+            :func:`pickle.loads` on the returned tuple executes
+            arbitrary Python on the receiver, including any
+            cloudpickled UDF callable embedded in the payload. Only
+            unpickle expressions from trusted sources.
+
+        Examples:
+            >>> import pickle
+            >>> from datafusion import col, lit
+            >>> e = col("a") * lit(2)
+            >>> pickle.loads(pickle.dumps(e)).canonical_name()
+            'a * Int64(2)'
         """
         return (Expr._reconstruct, (self.to_bytes(),))
 
     @classmethod
     def _reconstruct(cls, proto_bytes: bytes) -> Expr:
-        """Internal entry point used by :meth:`__reduce__` on unpickle."""
+        """Internal entry point used by :meth:`__reduce__` on unpickle.
+
+        Examples:
+            >>> from datafusion import Expr, col, lit
+            >>> blob = (col("a") + lit(1)).to_bytes()
+            >>> Expr._reconstruct(blob).canonical_name()
+            'a + Int64(1)'
+        """
         return cls.from_bytes(proto_bytes)
 
     def __richcmp__(self, other: Expr, op: int) -> Expr:

python/datafusion/ipc.py

@@ -65,6 +65,14 @@ def set_worker_ctx(ctx: SessionContext) -> None:
     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.
+
+    Examples:
+        >>> from datafusion import SessionContext
+        >>> from datafusion.ipc import set_worker_ctx, get_worker_ctx, clear_worker_ctx
+        >>> set_worker_ctx(SessionContext())
+        >>> get_worker_ctx() is not None
+        True
+        >>> clear_worker_ctx()
     """
     _local.ctx = ctx
 
@@ -76,13 +84,28 @@ def clear_worker_ctx() -> None:
     the global :class:`SessionContext` — adequate for built-ins and Python
     scalar UDFs, but anything imported via the FFI capsule protocol must
     be registered on the global context to resolve.
+
+    Examples:
+        >>> from datafusion import SessionContext
+        >>> from datafusion.ipc import set_worker_ctx, clear_worker_ctx, get_worker_ctx
+        >>> set_worker_ctx(SessionContext())
+        >>> clear_worker_ctx()
+        >>> get_worker_ctx() is None
+        True
     """
     if hasattr(_local, "ctx"):
         del _local.ctx
 
 
 def get_worker_ctx() -> SessionContext | None:
-    """Return this worker's installed :class:`SessionContext`, or ``None``."""
+    """Return this worker's installed :class:`SessionContext`, or ``None``.
+
+    Examples:
+        >>> from datafusion.ipc import get_worker_ctx, clear_worker_ctx
+        >>> clear_worker_ctx()
+        >>> get_worker_ctx() is None
+        True
+    """
     return getattr(_local, "ctx", None)
 
 
@@ -95,6 +118,16 @@ def _resolve_ctx(
     Falling back to the global :class:`SessionContext` (instead of a
     freshly constructed one) preserves any registrations the user has
     installed on it.
+
+    Examples:
+        >>> from datafusion import SessionContext
+        >>> from datafusion.ipc import _resolve_ctx, clear_worker_ctx
+        >>> clear_worker_ctx()
+        >>> isinstance(_resolve_ctx(), SessionContext)
+        True
+        >>> ctx = SessionContext()
+        >>> _resolve_ctx(ctx) is ctx
+        True
     """
     if explicit_ctx is not None:
         return explicit_ctx

python/datafusion/user_defined.py

@@ -148,6 +148,19 @@ def name(self) -> str:
         For UDFs imported via the FFI capsule protocol, this is the
         name the capsule itself reports — not the ``name`` argument
         passed to the constructor (which is ignored on the FFI path).
+
+        Examples:
+            >>> import pyarrow as pa
+            >>> from datafusion import udf
+            >>> double = udf(
+            ...     lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),
+            ...     [pa.int64()],
+            ...     pa.int64(),
+            ...     volatility="immutable",
+            ...     name="double",
+            ... )
+            >>> double.name
+            'double'
         """
         return self._udf.name
 

python/tests/test_pickle_expr.py

@@ -125,3 +125,33 @@ def fn(arr):
         blob = pickle.dumps(e)
         decoded = pickle.loads(blob)  # noqa: S301
         assert decoded.canonical_name() == e.canonical_name()
+
+    def test_multi_arg_udf_round_trip(self):
+        """Wire format builds synthetic `arg_{i}` fields per input — exercise
+        with a 2-arg UDF spanning two distinct DataTypes."""
+        add_scaled = udf(
+            lambda a, b: pa.array(
+                [
+                    (x.as_py() or 0) + (y.as_py() or 0.0)
+                    for x, y in zip(a, b, strict=False)
+                ]
+            ),
+            [pa.int64(), pa.float64()],
+            pa.float64(),
+            volatility="immutable",
+            name="add_scaled",
+        )
+        e = add_scaled(col("a"), col("b"))
+        decoded = pickle.loads(pickle.dumps(e))  # noqa: S301
+        assert decoded.canonical_name() == e.canonical_name()
+        assert "add_scaled" in decoded.canonical_name()
+
+
+class TestErrorPaths:
+    def test_from_bytes_rejects_garbage(self):
+        with pytest.raises(Exception):  # noqa: B017
+            Expr.from_bytes(b"not a valid protobuf payload")
+
+    def test_from_bytes_rejects_empty(self):
+        with pytest.raises(Exception):  # noqa: B017
+            Expr.from_bytes(b"")