feat: pickle/dill support for Expr

Add `to_bytes` / `from_bytes` on `Expr` (Python wrapper) and the underlying
`RawExpr` (Rust). Serialization uses `datafusion-proto`'s `Serializeable`
trait, encoding function references by name. The Python wrapper implements
`__getstate__` / `__setstate__` on top, so `pickle.dumps` / `dill.dumps` work
out of the box.

Reconstruction resolves function names against the process-wide global
`SessionContext` (introduced as settable in the previous commit). Built-in
functions always roundtrip; user-defined functions roundtrip when registered
on a context that has been installed via `SessionContext.set_as_global()`.

Adds `dill` to the dev dependency group and parametrized tests covering
both serializers across columns, literals, binary ops, casts, between,
aggregates, case/when, and a UDF with the global-ctx pattern.

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

Author: ntjohnson1

crates/core/src/expr.rs

@@ -31,9 +31,12 @@ use datafusion::logical_expr::{
     Between, BinaryExpr, Case, Cast, Expr, ExprFuncBuilder, ExprFunctionExt, Like, LogicalPlan,
     Operator, TryCast, WindowFunctionDefinition, col, lit, lit_with_metadata,
 };
+use datafusion_proto::bytes::Serializeable;
+use datafusion_python_util::get_global_ctx;
 use pyo3::IntoPyObjectExt;
 use pyo3::basic::CompareOp;
 use pyo3::prelude::*;
+use pyo3::types::PyBytes;
 use window::PyWindowFrame;
 
 use self::alias::PyAlias;
@@ -256,6 +259,29 @@ impl PyExpr {
         Ok(format!("Expr({})", self.expr))
     }
 
+    /// Serialize the underlying expression to bytes via the `datafusion-proto`
+    /// wire format. Used by the Python `Expr` wrapper to implement
+    /// `__getstate__` / `__setstate__`; also exposed directly so callers can
+    /// persist or transmit expressions without going through `pickle`.
+    fn to_bytes<'py>(&self, py: Python<'py>) -> PyDataFusionResult<Bound<'py, PyBytes>> {
+        let bytes = self.expr.to_bytes()?;
+        Ok(PyBytes::new(py, &bytes))
+    }
+
+    /// Reconstruct a `RawExpr` from bytes produced by [`PyExpr::to_bytes`].
+    ///
+    /// Function references (built-ins, UDFs, UDAFs, UDWFs) are resolved by
+    /// name against the process-wide global `SessionContext`. Built-in
+    /// functions are registered on every fresh context, so they always
+    /// roundtrip. To roundtrip user-defined functions, register them on a
+    /// context and call `SessionContext.set_as_global()` before unpickling.
+    #[staticmethod]
+    fn from_bytes(bytes: &[u8]) -> PyDataFusionResult<PyExpr> {
+        let ctx = get_global_ctx();
+        let expr = Expr::from_bytes_with_registry(bytes, ctx.as_ref())?;
+        Ok(expr.into())
+    }
+
     fn __add__(&self, rhs: PyExpr) -> PyResult<PyExpr> {
         Ok((self.expr.clone() + rhs.expr).into())
     }

pyproject.toml

@@ -188,6 +188,7 @@ ignore-words-list = ["IST", "ans"]
 dev = [
   "arro3-core==0.6.5",
   "codespell==2.4.1",
+  "dill>=0.3.8",
   "maturin>=1.8.1",
   "nanoarrow==0.8.0",
   "numpy>1.25.0;python_version<'3.14'",

python/datafusion/expr.py

@@ -410,6 +410,31 @@ def __init__(self, expr: expr_internal.RawExpr) -> None:
         """This constructor should not be called by the end user."""
         self.expr = expr
 
+    def to_bytes(self) -> bytes:
+        """Serialize this expression to bytes via the ``datafusion-proto`` wire format.
+
+        Function references (built-ins and UDFs/UDAFs/UDWFs) are encoded by
+        name; on :py:meth:`from_bytes` the names are resolved against the
+        process-wide global :py:class:`SessionContext`. Built-in functions
+        always roundtrip; for user-defined functions, register them on a
+        context and call :py:meth:`SessionContext.set_as_global` before
+        loading.
+        """
+        return self.expr.to_bytes()
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> Expr:
+        """Inverse of :py:meth:`to_bytes`. See that method for caveats."""
+        return cls(expr_internal.RawExpr.from_bytes(data))
+
+    def __getstate__(self) -> bytes:
+        """Serialize for ``pickle`` / ``dill``. Delegates to :py:meth:`to_bytes`."""
+        return self.to_bytes()
+
+    def __setstate__(self, state: bytes) -> None:
+        """Inverse of :py:meth:`__getstate__`."""
+        self.expr = expr_internal.RawExpr.from_bytes(state)
+
     def to_variant(self) -> Any:
         """Convert this expression into a python object if possible."""
         return self.expr.to_variant()

python/tests/test_pickle.py

@@ -0,0 +1,167 @@
+# 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.
+
+"""Pickle / dill roundtrip tests for :py:class:`datafusion.Expr`.
+
+The wire format is `datafusion-proto`'s ``LogicalExprNode``. Function
+references are encoded by name, so unpickling resolves them against the
+process-wide global :py:class:`SessionContext`. Tests that need a
+non-built-in function temporarily install a custom global context and
+restore the previous one.
+"""
+
+import pickle
+from contextlib import contextmanager
+
+import dill
+import pyarrow as pa
+import pytest
+from datafusion import SessionContext, col, lit, udf
+from datafusion import functions as f
+from datafusion.expr import Expr
+
+
+@pytest.fixture
+def ctx():
+    return SessionContext()
+
+
+@pytest.fixture
+def df(ctx):
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, None])],
+        names=["a", "b"],
+    )
+    return ctx.create_dataframe([[batch]], name="t")
+
+
+@contextmanager
+def temporary_global_ctx(new_ctx):
+    """Install ``new_ctx`` as the process-wide global and restore on exit."""
+    previous = SessionContext.global_ctx()
+    new_ctx.set_as_global()
+    try:
+        yield
+    finally:
+        previous.set_as_global()
+
+
+@pytest.mark.parametrize("dumper", [pickle, dill], ids=["pickle", "dill"])
+@pytest.mark.parametrize(
+    "build_expr",
+    [
+        pytest.param(lambda: col("a"), id="column"),
+        pytest.param(lambda: lit(42), id="literal_int"),
+        pytest.param(lambda: lit("hello"), id="literal_str"),
+        pytest.param(lambda: col("a") + lit(1), id="binary_add"),
+        pytest.param(lambda: (col("a") * lit(2)) - col("b"), id="binary_nested"),
+        pytest.param(lambda: col("a").alias("renamed"), id="alias"),
+        pytest.param(lambda: col("a").cast(pa.float64()), id="cast"),
+        pytest.param(lambda: col("a").is_null(), id="is_null"),
+        pytest.param(lambda: col("a").between(lit(1), lit(10)), id="between"),
+        pytest.param(lambda: ~(col("a") > lit(0)), id="not_gt"),
+        pytest.param(lambda: f.sum(col("a")), id="agg_sum"),
+        pytest.param(
+            lambda: f.case(col("a")).when(lit(1), lit("one")).end(),
+            id="case_when",
+        ),
+    ],
+)
+def test_builtin_roundtrip(build_expr, dumper):
+    """Built-in expressions roundtrip via pickle and dill."""
+    expr = build_expr()
+    restored = dumper.loads(dumper.dumps(expr))
+    assert isinstance(restored, Expr)
+    # canonical_name() gives a full string form including function names,
+    # so equal canonical names imply structural equivalence.
+    assert restored.canonical_name() == expr.canonical_name()
+
+
+@pytest.mark.parametrize("dumper", [pickle, dill], ids=["pickle", "dill"])
+def test_pickled_expr_executes(df, dumper):
+    """A roundtripped expression evaluates to the same result as the original."""
+    expr = (col("a") + lit(10)).alias("a_plus_ten")
+    restored = dumper.loads(dumper.dumps(expr))
+
+    original = df.select(expr).collect()[0].column(0)
+    after = df.select(restored).collect()[0].column(0)
+    assert original == after
+    assert original == pa.array([11, 12, 13], type=pa.int64())
+
+
+def test_udf_roundtrip_via_global_ctx():
+    """UDFs roundtrip when registered on the active global context.
+
+    Mirrors the documented usage of ``SessionContext.set_as_global``.
+    """
+    is_null = udf(
+        lambda x: x.is_null(),
+        [pa.int64()],
+        pa.bool_(),
+        volatility="immutable",
+        name="pickle_test_is_null",
+    )
+
+    custom_ctx = SessionContext()
+    custom_ctx.register_udf(is_null)
+
+    expr = is_null(col("b"))
+
+    with temporary_global_ctx(custom_ctx):
+        data = pickle.dumps(expr)
+        restored = pickle.loads(data)  # noqa: S301
+        assert restored.canonical_name() == expr.canonical_name()
+
+        # Also evaluate to confirm the UDF body is wired up post-roundtrip.
+        batch = pa.RecordBatch.from_arrays([pa.array([1, None, 3])], names=["b"])
+        df = custom_ctx.create_dataframe([[batch]], name="t_udf")
+        result = df.select(restored.alias("nul")).collect()[0].column(0)
+        assert result == pa.array([False, True, False])
+
+
+def test_udf_roundtrip_fails_without_registration():
+    """Without the UDF registered on the global context, unpickle errors out
+    rather than silently substituting a different implementation."""
+    is_null = udf(
+        lambda x: x.is_null(),
+        [pa.int64()],
+        pa.bool_(),
+        volatility="immutable",
+        name="pickle_test_unknown_udf",
+    )
+    expr = is_null(col("b"))
+
+    data = pickle.dumps(expr)
+    # The default global ctx does not have this UDF registered. Reconstruction
+    # must raise rather than silently substitute a placeholder. DataFusion
+    # surfaces this as a generic Python ``Exception`` whose message names the
+    # missing function, so match on the function name.
+    with pytest.raises(Exception, match="pickle_test_unknown_udf"):
+        pickle.loads(data)  # noqa: S301
+
+
+def test_getstate_returns_bytes():
+    """``__getstate__`` is exposed directly and returns raw bytes — useful for
+    callers that want to persist or transmit expressions without pickle."""
+    expr = col("a") + lit(1)
+    state = expr.__getstate__()
+    assert isinstance(state, bytes)
+    assert len(state) > 0
+
+    rebuilt = Expr.__new__(Expr)
+    rebuilt.__setstate__(state)
+    assert rebuilt.canonical_name() == expr.canonical_name()