python/tests/test_pickle_expr.py

# 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.

"""In-process pickle round-trip tests for :class:`Expr`.

Built-in functions and Python UDFs (scalar, aggregate, window) travel
with the pickled expression and do not need worker-side pre-registration.
The worker context (:mod:`datafusion.ipc`) is only consulted for UDFs
imported via the FFI capsule protocol.

Cross-process tests live in ``test_pickle_multiprocessing.py``.
"""

from __future__ import annotations

import pickle
import threading

import pyarrow as pa
import pytest
from datafusion import Expr, SessionContext, col, lit, udf
from datafusion.ipc import (
    clear_sender_ctx,
    clear_worker_ctx,
    get_sender_ctx,
    get_worker_ctx,
    set_sender_ctx,
    set_worker_ctx,
)


@pytest.fixture(autouse=True)
def _reset_worker_ctx():
    """Ensure every test starts with no worker or sender context installed."""
    clear_worker_ctx()
    clear_sender_ctx()
    yield
    clear_worker_ctx()
    clear_sender_ctx()


def _double_udf():
    return udf(
        lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),
        [pa.int64()],
        pa.int64(),
        volatility="immutable",
        name="double",
    )


class TestProtoRoundTrip:
    def test_builtin_round_trip(self):
        e = col("a") + lit(1)
        blob = pickle.dumps(e)
        decoded = pickle.loads(blob)  # noqa: S301
        assert decoded.canonical_name() == e.canonical_name()

    def test_to_bytes_from_bytes(self):
        e = col("x") * lit(7)
        blob = e.to_bytes()
        assert isinstance(blob, bytes)
        decoded = Expr.from_bytes(blob)
        assert decoded.canonical_name() == e.canonical_name()

    def test_explicit_ctx_used(self, ctx):
        e = col("a") + lit(1)
        decoded = Expr.from_bytes(e.to_bytes(), ctx=ctx)
        assert decoded.canonical_name() == e.canonical_name()


class TestUDFCodec:
    """Python scalar UDFs ride inside the proto blob via the Rust codec.

    No worker context needed on the receiver — the cloudpickled callable is
    embedded in ``fun_definition`` and reconstructed automatically.
    """

    def test_udf_self_contained_blob(self):
        e = _double_udf()(col("a"))
        blob = pickle.dumps(e)
        # The codec inlines the callable, so the blob is much bigger than a
        # pure built-in blob but doesn't depend on receiver-side registration.
        assert len(blob) > 200

    def test_udf_decodes_into_fresh_ctx(self):
        e = _double_udf()(col("a"))
        blob = e.to_bytes()
        fresh = SessionContext()
        decoded = Expr.from_bytes(blob, ctx=fresh)
        assert "double" in decoded.canonical_name()

    def test_udf_decodes_via_pickle_with_no_worker_ctx(self):
        e = _double_udf()(col("a"))
        blob = pickle.dumps(e)
        decoded = pickle.loads(blob)  # noqa: S301
        assert "double" in decoded.canonical_name()

    def test_udf_decodes_via_pickle_with_worker_ctx(self):
        set_worker_ctx(SessionContext())
        e = _double_udf()(col("a"))
        blob = pickle.dumps(e)
        decoded = pickle.loads(blob)  # noqa: S301
        assert "double" in decoded.canonical_name()

    def test_closure_capturing_udf_names_match(self):
        captured_multiplier = 7

        def fn(arr):
            return pa.array([(v.as_py() or 0) * captured_multiplier for v in arr])

        u = udf(
            fn,
            [pa.int64()],
            pa.int64(),
            volatility="immutable",
            name="times_seven",
        )
        e = u(col("a"))
        blob = pickle.dumps(e)
        decoded = pickle.loads(blob)  # noqa: S301
        # Round-trip names match; functional verification of captured state
        # happens in test_pickle_multiprocessing via an actual UDF call.
        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 TestAggregateUDFCodec:
    """Python aggregate UDFs travel inline like scalar UDFs."""

    def _build_aggregate_udf(self):
        from datafusion import udaf
        from datafusion.user_defined import Accumulator

        class CountAcc(Accumulator):
            def __init__(self):
                self._count = 0

            def state(self):
                return [pa.scalar(self._count, type=pa.int64())]

            def update(self, values):
                self._count += len(values)

            def merge(self, states):
                partition_counts = states[0]
                for i in range(len(partition_counts)):
                    self._count += partition_counts[i].as_py()

            def evaluate(self):
                return pa.scalar(self._count, type=pa.int64())

        return udaf(
            CountAcc,
            [pa.int64()],
            pa.int64(),
            [pa.int64()],
            "immutable",
            name="count_all",
        )

    def test_agg_udf_self_contained_blob(self):
        u = self._build_aggregate_udf()
        e = u(col("a"))
        blob = pickle.dumps(e)
        assert len(blob) > 200

    def test_agg_udf_decodes_into_fresh_ctx(self):
        u = self._build_aggregate_udf()
        e = u(col("a"))
        blob = e.to_bytes()
        fresh = SessionContext()
        decoded = Expr.from_bytes(blob, ctx=fresh)
        assert "count_all" in decoded.canonical_name()

    def test_agg_udf_decodes_via_pickle_with_no_worker_ctx(self):
        u = self._build_aggregate_udf()
        e = u(col("a"))
        blob = pickle.dumps(e)
        decoded = pickle.loads(blob)  # noqa: S301
        assert "count_all" in decoded.canonical_name()

    def test_agg_udf_evaluates_after_roundtrip(self):
        """End-to-end: the decoded aggregate UDF runs and merges across
        partitions, exercising the round-tripped state-field schema."""
        u = self._build_aggregate_udf()
        e = u(col("a"))
        decoded = pickle.loads(pickle.dumps(e))  # noqa: S301

        ctx = SessionContext()
        schema = pa.schema([pa.field("a", pa.int64())])
        batch1 = pa.record_batch([pa.array([1, 2, 3], type=pa.int64())], schema=schema)
        batch2 = pa.record_batch([pa.array([4, 5], type=pa.int64())], schema=schema)
        df = ctx.create_dataframe([[batch1], [batch2]])
        out = df.aggregate([], [decoded.alias("n")]).to_pydict()
        assert out["n"] == [5]


class TestWindowUDFCodec:
    """Python window UDFs travel inline like scalar UDFs."""

    def _build_window_udf(self):
        from datafusion import udwf
        from datafusion.user_defined import WindowEvaluator

        class CountUpEvaluator(WindowEvaluator):
            def evaluate_all(self, values, num_rows):
                return pa.array(list(range(num_rows)))

        return udwf(
            CountUpEvaluator,
            [pa.int64()],
            pa.int64(),
            "immutable",
            name="count_up",
        )

    def test_window_udf_self_contained_blob(self):
        u = self._build_window_udf()
        e = u(col("a"))
        blob = pickle.dumps(e)
        assert len(blob) > 200

    def test_window_udf_decodes_into_fresh_ctx(self):
        u = self._build_window_udf()
        e = u(col("a"))
        blob = e.to_bytes()
        fresh = SessionContext()
        decoded = Expr.from_bytes(blob, ctx=fresh)
        assert "count_up" in decoded.canonical_name()

    def test_window_udf_decodes_via_pickle_with_no_worker_ctx(self):
        u = self._build_window_udf()
        e = u(col("a"))
        blob = pickle.dumps(e)
        decoded = pickle.loads(blob)  # noqa: S301
        assert "count_up" in decoded.canonical_name()

    def test_window_udf_evaluates_after_roundtrip(self):
        """End-to-end: decoded window UDF runs and emits per-row values
        produced by the round-tripped evaluator factory."""
        from datafusion.expr import WindowFrame

        u = self._build_window_udf()
        e = u(col("a"))
        decoded = pickle.loads(pickle.dumps(e))  # noqa: S301

        ctx = SessionContext()
        df = ctx.from_pydict({"a": [1, 2, 3, 4, 5]})
        framed = (
            decoded.window_frame(WindowFrame("rows", None, None)).build().alias("c")
        )
        out = df.select(framed).to_pydict()
        assert out["c"] == [0, 1, 2, 3, 4]


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"")

    def test_cross_version_error_message(self):
        """Decoding a payload stamped with a different Python minor
        version raises a clear, actionable error rather than an opaque
        marshal/unpickle failure.

        The wire frame inside the protobuf is:
        ``DFPYUDF (7) | version (1) | py_major (1) | py_minor (1) | cloudpickle``.
        We locate the frame inside the outer protobuf and patch the
        minor byte at offset 9.
        """
        import sys

        e = _double_udf()(col("a"))
        blob = e.to_bytes()

        idx = blob.find(b"DFPYUDF")
        assert idx >= 0, "DFPYUDF frame not found in payload"

        different_minor = (sys.version_info.minor + 1) % 256
        tampered = bytearray(blob)
        tampered[idx + 9] = different_minor

        with pytest.raises(
            Exception, match="not portable across Python minor versions"
        ):
            Expr.from_bytes(bytes(tampered))


class TestPythonUdfInliningToggle:
    """`SessionContext.with_python_udf_inlining(enabled=False)` opts out of
    inline Python UDF encoding for both encode and decode paths."""

    def _build_double_udf(self):
        return udf(
            lambda arr: pa.array([(v.as_py() or 0) * 2 for v in arr]),
            [pa.int64()],
            pa.int64(),
            volatility="immutable",
            name="double",
        )

    def test_strict_encoder_omits_inline_payload(self):
        """Strict mode emits the by-name wire form: no `DFPYUDF` magic
        in the blob, no cloudpickled callable. Semantic check is
        sharper than a size-ratio heuristic — a renamed UDF or a
        smaller-than-expected closure would still flip the magic
        bytes, but might not move the size by 4x.
        """
        ctx_inline = SessionContext()
        ctx_strict = ctx_inline.with_python_udf_inlining(enabled=False)
        u = self._build_double_udf()
        e = u(col("a"))

        blob_inline = e.to_bytes(ctx_inline)
        blob_strict = e.to_bytes(ctx_strict)

        # `DFPYUDF` is the scalar Python-UDF family prefix; see
        # `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs.
        assert b"DFPYUDF" in blob_inline
        assert b"DFPYUDF" not in blob_strict

    def test_toggle_off_then_on_restores_inline_encoding(self):
        """`with_python_udf_inlining` is per-call clone semantics:
        flipping off and then on must produce a context that emits the
        same inline form as a fresh default context, byte-for-byte.

        Guards against a regression where the off→on transition leaves
        the codec in a sticky strict state (e.g. by mutating shared
        codec state instead of cloning).
        """
        u = self._build_double_udf()
        e = u(col("a"))

        baseline = SessionContext()
        toggled = (
            SessionContext()
            .with_python_udf_inlining(enabled=False)
            .with_python_udf_inlining(enabled=True)
        )

        blob_baseline = e.to_bytes(baseline)
        blob_toggled = e.to_bytes(toggled)

        assert blob_baseline == blob_toggled

        # Sanity check the decoded form against a fresh ctx — the
        # toggled-back blob should be self-contained inline, not a
        # strict by-name payload that needs registry resolution.
        decoded = Expr.from_bytes(blob_toggled, ctx=SessionContext())
        assert "double" in decoded.canonical_name()

    def test_strict_roundtrip_via_registry(self):
        """When both sender and receiver disable inlining, the UDF
        travels by name only and the receiver resolves it from its
        registered functions."""
        strict_sender = SessionContext().with_python_udf_inlining(enabled=False)
        u = self._build_double_udf()
        blob = u(col("a")).to_bytes(strict_sender)

        receiver = SessionContext().with_python_udf_inlining(enabled=False)
        receiver.register_udf(u)
        restored = Expr.from_bytes(blob, ctx=receiver)
        assert "double" in restored.canonical_name()

    def test_strict_decoder_refuses_inline_payload(self):
        """An inline-encoded blob fed to a strict receiver raises with a
        clear error rather than silently invoking cloudpickle.loads.

        The receiver is intentionally *not* given a matching
        registration: the codec refusal must trip before the registry
        is ever consulted, so registering the UDF here would only mask
        a regression that moved the check after registry lookup.
        """
        sender = SessionContext()
        u = self._build_double_udf()
        blob = u(col("a")).to_bytes(sender)

        strict_receiver = SessionContext().with_python_udf_inlining(enabled=False)
        # `RuntimeError` (not bare `Exception`): the codec refusal is
        # surfaced through `parse_expr` → `PyRuntimeError`. Tightening
        # the assertion catches a regression that swallows the refusal
        # as a different error type.
        with pytest.raises(RuntimeError, match="inlining is disabled"):
            Expr.from_bytes(blob, ctx=strict_receiver)

    def test_sender_ctx_propagates_through_pickle(self):
        """`set_sender_ctx` makes `pickle.dumps` use a strict codec.

        Without a sender context, pickle defaults to the inline codec
        and the blob contains the `DFPYUDF` family prefix. With a
        strict sender context installed, the callable encodes by name
        and the prefix is absent.
        """
        u = self._build_double_udf()
        e = u(col("a"))

        blob_default = pickle.dumps(e)

        strict_sender = SessionContext().with_python_udf_inlining(enabled=False)
        set_sender_ctx(strict_sender)
        try:
            blob_strict = pickle.dumps(e)
        finally:
            clear_sender_ctx()

        assert b"DFPYUDF" in blob_default
        assert b"DFPYUDF" not in blob_strict

    def test_sender_ctx_strict_roundtrip_via_pickle(self):
        """End-to-end pickle round-trip with strict mode on both sides.

        Driver installs a strict sender context. Worker installs a
        matching strict context with the UDF registered. The UDF
        travels by name through `pickle.dumps` / `pickle.loads`.
        """
        u = self._build_double_udf()
        e = u(col("a"))

        strict_sender = SessionContext().with_python_udf_inlining(enabled=False)
        set_sender_ctx(strict_sender)
        try:
            blob = pickle.dumps(e)
        finally:
            clear_sender_ctx()

        worker = SessionContext().with_python_udf_inlining(enabled=False)
        worker.register_udf(u)
        set_worker_ctx(worker)
        try:
            decoded = pickle.loads(blob)  # noqa: S301
        finally:
            clear_worker_ctx()

        assert "double" in decoded.canonical_name()

    def test_sender_ctx_strict_pickle_accepted_by_inline_worker_with_registry(self):
        """A strict-encoded blob still decodes fine on an inline worker
        because the wire format is the same default-codec by-name form.
        Sanity check: cross-config works as long as the receiver can
        resolve the name."""
        u = self._build_double_udf()
        e = u(col("a"))

        strict_sender = SessionContext().with_python_udf_inlining(enabled=False)
        set_sender_ctx(strict_sender)
        try:
            blob = pickle.dumps(e)
        finally:
            clear_sender_ctx()

        worker = SessionContext()
        worker.register_udf(u)
        set_worker_ctx(worker)
        try:
            decoded = pickle.loads(blob)  # noqa: S301
        finally:
            clear_worker_ctx()

        assert "double" in decoded.canonical_name()


class TestWorkerCtxLifecycle:
    def test_set_and_clear(self):
        assert get_worker_ctx() is None
        ctx = SessionContext()
        set_worker_ctx(ctx)
        assert get_worker_ctx() is ctx
        clear_worker_ctx()
        assert get_worker_ctx() is None

    def test_clear_when_unset_is_noop(self):
        clear_worker_ctx()  # no error
        assert get_worker_ctx() is None

    def test_thread_local_isolation(self):
        main_ctx = SessionContext()
        set_worker_ctx(main_ctx)

        seen_in_thread: list = []

        def worker():
            seen_in_thread.append(get_worker_ctx())
            set_worker_ctx(SessionContext())
            seen_in_thread.append(get_worker_ctx())

        t = threading.Thread(target=worker)
        t.start()
        t.join()

        # Thread saw no ctx initially (thread-local), then its own.
        assert seen_in_thread[0] is None
        assert seen_in_thread[1] is not main_ctx
        # Main thread's ctx is unchanged by the thread's actions.
        assert get_worker_ctx() is main_ctx


class TestSenderCtxLifecycle:
    def test_set_and_clear(self):
        assert get_sender_ctx() is None
        ctx = SessionContext()
        set_sender_ctx(ctx)
        assert get_sender_ctx() is ctx
        clear_sender_ctx()
        assert get_sender_ctx() is None

    def test_clear_when_unset_is_noop(self):
        clear_sender_ctx()  # no error
        assert get_sender_ctx() is None

    def test_thread_local_isolation(self):
        main_ctx = SessionContext()
        set_sender_ctx(main_ctx)

        seen_in_thread: list = []

        def worker():
            seen_in_thread.append(get_sender_ctx())
            set_sender_ctx(SessionContext())
            seen_in_thread.append(get_sender_ctx())

        t = threading.Thread(target=worker)
        t.start()
        t.join()

        assert seen_in_thread[0] is None
        assert seen_in_thread[1] is not main_ctx
        assert get_sender_ctx() is main_ctx