Allow pickling PyExpr

.github/workflows/test.yml

@@ -29,6 +29,7 @@ env:
 jobs:
  test-matrix:
  runs-on: ubuntu-latest
+ timeout-minutes: 30
  strategy:
  fail-fast: false
  matrix:

conftest.py

@@ -17,12 +17,29 @@
 
 """Pytest configuration for doctest namespace injection."""
 
-import datafusion as dfn
-import numpy as np
-import pyarrow as pa
-import pytest
-from datafusion import col, lit
-from datafusion import functions as F
+import sys
+from pathlib import Path
+
+# Ensure ``python/`` is reachable by ``import``. The ``tests`` package lives at
+# ``python/tests`` and spawn-based multiprocessing tests need workers to be
+# able to resolve ``tests._pickle_multiprocessing_helpers`` by its real dotted
+# name when unpickling task args. Editable installs add this path via a .pth
+# file, but the wheel install used in CI does not, which led to spawn workers
+# dying with ``ModuleNotFoundError`` and ``Pool.map`` hanging.
+#
+# Append (don't prepend) so the wheel-installed ``datafusion`` in
+# site-packages still wins over the source tree at ``python/datafusion`` —
+# the source tree has no compiled ``_internal`` module on a fresh checkout.
+_python_dir = str(Path(__file__).parent / "python")
+if _python_dir not in sys.path:
+ sys.path.append(_python_dir)
+
+import datafusion as dfn # noqa: E402
+import numpy as np # noqa: E402
+import pyarrow as pa # noqa: E402
+import pytest # noqa: E402
+from datafusion import col, lit # noqa: E402
+from datafusion import functions as F # noqa: E402
 
 
 @pytest.fixture(autouse=True)

crates/core/src/context.rs

@@ -56,7 +56,7 @@ use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
 use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
 use datafusion_python_util::{
  create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx,
- get_tokio_runtime, spawn_future, wait_for_future,
+ get_tokio_runtime, set_global_ctx, spawn_future, wait_for_future,
 };
 use object_store::ObjectStore;
 use pyo3::IntoPyObjectExt;
@@ -407,11 +407,22 @@ impl PySessionContext {
  #[staticmethod]
  #[pyo3(signature = ())]
  pub fn global_ctx() -> PyResult<Self> {
- let ctx = get_global_ctx().clone();
+ let ctx = get_global_ctx();
  let logical_codec = Self::default_logical_codec(&ctx);
  Ok(Self { ctx, logical_codec })
  }
 
+ /// Replace the process-wide global `SessionContext` with this one.
+ ///
+ /// All subsequent callers of `SessionContext.global_ctx()` (and Rust
+ /// helpers that fall back to the global context, such as the
+ /// `read_parquet` / `read_csv` / etc. module-level helpers) will see this
+ /// context. Existing references already obtained from `global_ctx()` are
+ /// not affected.
+ pub fn set_as_global(&self) {
+ set_global_ctx(self.ctx.clone());
+ }
+
  /// Register an object store with the given name
  #[pyo3(signature = (scheme, store, host=None))]
  pub fn register_object_store(

crates/core/src/dataframe.rs

@@ -851,7 +851,7 @@ impl PyDataFrame {
  Some(f) => f
  .parse::<datafusion::common::format::ExplainFormat>()
  .map_err(|e| {
- PyDataFusionError::Common(format!("Invalid explain format '{}': {}", f, e))
+ PyDataFusionError::Common(format!("Invalid explain format '{f}': {e}"))
  })?,
  None => datafusion::common::format::ExplainFormat::Indent,
  };

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())
  }

crates/util/src/lib.rs

@@ -17,7 +17,7 @@
 
 use std::future::Future;
 use std::ptr::NonNull;
-use std::sync::{Arc, OnceLock};
+use std::sync::{Arc, OnceLock, RwLock};
 use std::time::Duration;
 
 use datafusion::datasource::TableProvider;
@@ -59,11 +59,29 @@ pub fn is_ipython_env(py: Python) -> &'static bool {
  })
 }
 
-/// Utility to get the Global Datafussion CTX
+fn global_ctx_slot() -> &'static RwLock<Arc<SessionContext>> {
+ static CTX: OnceLock<RwLock<Arc<SessionContext>>> = OnceLock::new();
+ CTX.get_or_init(|| RwLock::new(Arc::new(SessionContext::new())))
+}
+
+/// Utility to get the Global DataFusion CTX.
+///
+/// Returns an owned `Arc<SessionContext>` snapshot. The underlying slot can be
+/// replaced via [`set_global_ctx`]; existing snapshots are unaffected.
 #[inline]
-pub fn get_global_ctx() -> &'static Arc<SessionContext> {
- static CTX: OnceLock<Arc<SessionContext>> = OnceLock::new();
- CTX.get_or_init(|| Arc::new(SessionContext::new()))
+pub fn get_global_ctx() -> Arc<SessionContext> {
+ global_ctx_slot()
+ .read()
+ .expect("global SessionContext lock poisoned")
+ .clone()
+}
+
+/// Replace the Global DataFusion CTX. Subsequent calls to [`get_global_ctx`]
+/// will return the new context. Already-cloned `Arc`s are not affected.
+pub fn set_global_ctx(ctx: Arc<SessionContext>) {
+ *global_ctx_slot()
+ .write()
+ .expect("global SessionContext lock poisoned") = ctx;
 }
 
 /// Utility to collect rust futures with GIL released and respond to
@@ -224,3 +242,40 @@ pub fn ffi_logical_codec_from_pycapsule(obj: Bound<PyAny>) -> PyResult<FFI_Logic
 
  Ok(codec.clone())
 }
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ /// The global slot must round-trip a custom `SessionContext`. Since the
+ /// global is process-wide, this test only asserts identity through a
+ /// single set/get cycle and restores the prior value at the end so the
+ /// test is independent of ordering with other tests in the binary.
+ #[test]
+ fn set_global_ctx_replaces_default() {
+ let prior = get_global_ctx();
+ let custom = Arc::new(SessionContext::new());
+ let custom_ptr = Arc::as_ptr(&custom);
+
+ set_global_ctx(custom.clone());
+ let observed = get_global_ctx();
+ assert_eq!(
+ Arc::as_ptr(&observed),
+ custom_ptr,
+ "get_global_ctx should return the context installed by set_global_ctx",
+ );
+
+ // A snapshot taken before the swap should be unaffected after another
+ // set_global_ctx call, because get_global_ctx clones the Arc.
+ let snapshot = get_global_ctx();
+ let replacement = Arc::new(SessionContext::new());
+ set_global_ctx(replacement);
+ assert_eq!(
+ Arc::as_ptr(&snapshot),
+ custom_ptr,
+ "previously cloned snapshots must not be invalidated by set_global_ctx",
+ );
+
+ set_global_ctx(prior);
+ }
+}

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'",
@@ -196,6 +197,7 @@ dev = [
  "pyarrow>=19.0.0",
  "pygithub==2.5.0",
  "pytest-asyncio>=0.23.3",
+ "pytest-timeout>=2.3.1",
  "pytest>=7.4.4",
  "pyyaml>=6.0.3",
  "ruff>=0.15.1",

python/datafusion/context.py

@@ -566,6 +566,22 @@ def global_ctx(cls) -> SessionContext:
  wrapper.ctx = internal_ctx
  return wrapper
 
+ def set_as_global(self) -> None:
+ """Install this context as the process-wide global ``SessionContext``.
+
+ After this call, :meth:`SessionContext.global_ctx` (and the module-level
+ helpers in :mod:`datafusion.io` that fall back to the global context)
+ will return this context. Existing references already obtained from
+ ``global_ctx()`` are not invalidated.
+
+ Example::
+
+ ctx = SessionContext()
+ ctx.register_udf(my_udf)
+ ctx.set_as_global()
+ """
+ self.ctx.set_as_global()
+
  def enable_url_table(self) -> SessionContext:
  """Control if local files can be queried as tables.
 

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/_pickle_multiprocessing_helpers.py

@@ -0,0 +1,85 @@
+# 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.
+
+"""Helpers for :mod:`test_pickle_multiprocessing`.
+
+Spawn workers re-import the module that defines a pickled function by the
+function's ``__module__`` attribute. Pytest's ``--import-mode=importlib``
+loads test modules under synthetic names that the worker cannot resolve via
+the normal import machinery, which can cause ``Pool.map`` to hang waiting
+for a worker that died during unpickling.
+
+Keeping the helpers in this regular (non-test) module side-steps that: it
+is importable under its real dotted name (``tests._pickle_multiprocessing_helpers``)
+in both parent and worker, and the leading underscore keeps pytest from
+collecting it as a test module.
+"""
+
+from __future__ import annotations
+
+import pyarrow as pa
+import pyarrow.compute as pc
+from datafusion import SessionContext, udf
+
+UDF_NAME = "mp_pickle_add_ten"
+
+
+def add_ten_impl(array: pa.Array) -> pa.Array:
+ return pc.add(array, 10)
+
+
+def build_add_ten_udf():
+ return udf(
+ add_ten_impl,
+ [pa.int64()],
+ pa.int64(),
+ volatility="immutable",
+ name=UDF_NAME,
+ )
+
+
+def register_udf_on_global_ctx() -> None:
+ """Pool initializer: install a global ctx in the worker that knows the UDF.
+
+ ``Expr.__setstate__`` resolves UDF references by name against the
+ *global* context, so the registration must happen before any task arg is
+ unpickled — i.e. in the Pool's ``initializer``, not in the task body.
+ """
+ ctx = SessionContext()
+ ctx.register_udf(build_add_ten_udf())
+ ctx.set_as_global()
+
+
+def apply_builtin_expr(args: tuple) -> list:
+ expr, values = args
+ ctx = SessionContext()
+ batch = pa.RecordBatch.from_arrays([pa.array(values, type=pa.int64())], names=["a"])
+ df = ctx.create_dataframe([[batch]], name="t")
+ return df.select(expr.alias("out")).collect()[0].column(0).to_pylist()
+
+
+def apply_udf_expr(args: tuple) -> list:
+ expr, values = args
+ # Reuse the worker's global ctx so the UDF registered by the initializer
+ # is visible during execution as well as during arg unpickling. Omit the
+ # table name so each call gets a fresh auto-generated one — a worker may
+ # process multiple tasks, and reusing a fixed name on the shared ctx would
+ # collide on the second call.
+ ctx = SessionContext.global_ctx()
+ batch = pa.RecordBatch.from_arrays([pa.array(values, type=pa.int64())], names=["a"])
+ df = ctx.create_dataframe([[batch]])
+ return df.select(expr.alias("out")).collect()[0].column(0).to_pylist()