test: move pickle/multiprocessing helpers into a non-test module

Spawn workers re-import the module that defined a pickled function by its
``__module__`` attribute. Pytest's ``--import-mode=importlib`` (used in
CI) can load test modules under synthetic names that the worker's normal
import machinery cannot resolve, which manifests as ``Pool.map`` hanging
forever waiting on a worker that died during unpickling.

Move the worker-side helpers (``_apply_builtin_expr``, ``_apply_udf_expr``,
``_register_udf_on_global_ctx``, ``_build_add_ten_udf``) to a regular
``tests._pickle_multiprocessing_helpers`` module. The leading underscore
keeps pytest from collecting it; spawn workers import it under its real
dotted name in both parent and child.

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

Author: ntjohnson1

python/tests/_pickle_multiprocessing_helpers.py

@@ -0,0 +1,82 @@
+# 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.
+    ctx = SessionContext.global_ctx()
+    batch = pa.RecordBatch.from_arrays([pa.array(values, type=pa.int64())], names=["a"])
+    df = ctx.create_dataframe([[batch]], name="t_udf")
+    return df.select(expr.alias("out")).collect()[0].column(0).to_pylist()

python/tests/test_pickle_multiprocessing.py

@@ -32,64 +32,24 @@
 2. A custom Python UDF — the worker must register the UDF on its global
    context *before* unpickling, since ``Expr.__setstate__`` resolves
    function references by name against the global context.
+
+Worker-side helpers live in :mod:`tests._pickle_multiprocessing_helpers`
+rather than in this test module, so spawn workers can resolve them by their
+real dotted name regardless of how pytest imported the test module.
 """
 
 from __future__ import annotations
 
 import multiprocessing as mp
 
-import pyarrow as pa
-import pyarrow.compute as pc
-from datafusion import SessionContext, col, lit, udf
-
-# Module-scope helpers — must be importable by name so the `spawn` workers
-# can resolve them after re-importing this module.
-
-_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()
-
+from datafusion import col, lit
 
-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.
-    ctx = SessionContext.global_ctx()
-    batch = pa.RecordBatch.from_arrays([pa.array(values, type=pa.int64())], names=["a"])
-    df = ctx.create_dataframe([[batch]], name="t_udf")
-    return df.select(expr.alias("out")).collect()[0].column(0).to_pylist()
+from tests._pickle_multiprocessing_helpers import (
+    apply_builtin_expr,
+    apply_udf_expr,
+    build_add_ten_udf,
+    register_udf_on_global_ctx,
+)
 
 
 def test_builtin_expr_through_multiprocessing_pool() -> None:
@@ -99,7 +59,7 @@ def test_builtin_expr_through_multiprocessing_pool() -> None:
     chunks = [[1, 2, 3], [10, 20, 30]]
 
     with spawn_ctx.Pool(processes=2) as pool:
-        results = pool.map(_apply_builtin_expr, [(expr, c) for c in chunks])
+        results = pool.map(apply_builtin_expr, [(expr, c) for c in chunks])
 
     assert results == [[3, 5, 7], [21, 41, 61]]
 
@@ -108,11 +68,11 @@ def test_udf_expr_through_multiprocessing_pool() -> None:
     """A UDF-backed ``Expr`` survives ``Pool.map`` when the worker registers
     the UDF on its global context via the Pool initializer."""
     spawn_ctx = mp.get_context("spawn")
-    add_ten = _build_add_ten_udf()
+    add_ten = build_add_ten_udf()
     expr = add_ten(col("a"))
     chunks = [[1, 2, 3], [10, 20, 30]]
 
-    with spawn_ctx.Pool(processes=2, initializer=_register_udf_on_global_ctx) as pool:
-        results = pool.map(_apply_udf_expr, [(expr, c) for c in chunks])
+    with spawn_ctx.Pool(processes=2, initializer=register_udf_on_global_ctx) as pool:
+        results = pool.map(apply_udf_expr, [(expr, c) for c in chunks])
 
     assert results == [[11, 12, 13], [20, 30, 40]]