python/pyspark/sql/tests/test_udf_transpile_hypothesis.py - master - githubmirror/spark

python/pyspark/sql/tests/test_udf_transpile_hypothesis.py
711 строк28 KB
Holden Karau
[SPARK-55206][PYTHON][SQL] Transpilation minimal functional implementation with python
03 авг 2026, 19:13
03 авг 2026, 19:131680141
#
# 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.
#
"""
Differential / property-based tests for UDF transpilation (SPARK-54783).

These tests run a fixed set of small Python UDFs twice -- once with
``spark.sql.experimental.optimizer.transpilePyUDFs`` enabled (so the catalyst
transpiler in :mod:`pyspark.sql.transpile` rewrites them into native
expressions) and once without -- and assert that the two runs produce the
same results for inputs generated by Hypothesis.

The transpiler is intentionally minimal at this point so we expect this
suite to surface bugs (e.g. truthiness / NULL-handling mismatches between
Python's ``if x:`` semantics and SQL's ``CASE WHEN``). Failures here should
be treated as real correctness gaps in the transpiler, not as test bugs to
silence.

The suite is gated on two things, both required:

* the ``RUN_HYPOTHESIS`` env var must be set to a truthy value
  (``1``, ``true``, or ``yes``, case-insensitive), and
* the ``hypothesis`` package must be installed.

The gate is value-based rather than presence-based because CI always sets
``RUN_HYPOTHESIS`` (to ``"true"`` or ``"false"``) via the transpile
precondition in ``build_and_test.yml``; mere presence must not opt in, or
the slow suite would run on every PySpark job.

If either gate is unmet the entire suite is skipped cleanly so it never
becomes a CI tax for folks who haven't opted in. In CI, this opt-in suite
is wired through ``.github/workflows/build_and_test.yml``, which flips both
gates on for the relevant job when PR changes touch the transpiler or this
test file.

To run locally::

    pip install hypothesis
    RUN_HYPOTHESIS=1 RUN_HYPOTHESIS_MAX_EXAMPLES=1000 \
        python/run-tests --testnames pyspark.sql.tests.test_udf_transpile_hypothesis

Set ``RUN_HYPOTHESIS_MAX_EXAMPLES`` to override the per-test example count
(default 1000). Each generated example runs two full Spark jobs (a
transpiled-vs-interpreted differential), so CI caps this at 50 via
``build_and_test.yml`` to stay under ``PYSPARK_TEST_TIMEOUT``; the explicit
``@example`` edge seeds always run on top of the generated ones regardless.
"""

import os
import unittest
import warnings
from typing import Optional

from pyspark.sql import Row
from pyspark.sql.types import (
    BooleanType,
    LongType,
    StructField,
    StructType,
)
from pyspark.sql.udf import UserDefinedFunction
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import have_package
from pyspark.util import is_remote_only


# Sentinel value used by ``_run`` to mark "this side raised". A unique
# object is sufficient because we only ever compare it against itself
# inside the helper.
_SENTINEL_RAISED = object()


_HYPOTHESIS_ENV = "RUN_HYPOTHESIS"
_have_hypothesis = have_package("hypothesis")


def _env_opts_in(value: Optional[str]) -> bool:
    """Value-based opt-in: only 1/true/yes (case-insensitive) enable the suite.

    CI always sets ``RUN_HYPOTHESIS`` -- to ``"true"`` or ``"false"`` -- via the
    transpile precondition in ``build_and_test.yml``, so a presence-based check
    would run this very slow suite on every PySpark job regardless of the
    gating decision.
    """
    return value is not None and value.strip().lower() in ("1", "true", "yes")


_hypothesis_enabled = _env_opts_in(os.environ.get(_HYPOTHESIS_ENV))
# Transpilation is only supported in regular (non-Connect) Spark for now,
# so the hypothesis suite skips cleanly under a pyspark-client-only install.
_regular_spark = not is_remote_only()
_skip_reason = (
    f"Set {_HYPOTHESIS_ENV}=1 (or true/yes) to run; hypothesis must also be installed, "
    "and the suite only runs under regular (non-Connect) Spark."
)


if _have_hypothesis:
    from hypothesis import HealthCheck, example, given, settings, strategies as st

    _DEFAULT_MAX_EXAMPLES = int(os.environ.get("RUN_HYPOTHESIS_MAX_EXAMPLES", "1000"))

    # The ``function_scoped_fixture` health check is suppressed because we intentionally reuse the
    # class-level SparkSession across examples; the per-example ``deadline`` is disabled because
    # Spark task execution is much slower than hypothesis's default budget.
    _hyp_settings = settings(
        max_examples=_DEFAULT_MAX_EXAMPLES,
        deadline=None,
        suppress_health_check=[HealthCheck.function_scoped_fixture],
    )

    # Full 64-bit signed range -- used by comparison and equality tests where
    # no arithmetic can overflow.
    _LONG_BOUND = 2**63 - 1
    _long_strategy = st.one_of(
        st.none(), st.integers(min_value=-_LONG_BOUND, max_value=_LONG_BOUND)
    )

    # Narrower range for arithmetic tests (+4, -2, *3, +7, x+y).  Python's
    # arithmetic never overflows, but Spark's ANSI mode raises on LongType
    # overflow.  Worse, the Python UDF runner silently wraps out-of-range
    # return values even with ANSI=True, so "both raised" never fires for
    # overflow values and the test sees a spurious mismatch instead.
    # 2**61 is safe for all operations: (2**61)*3 ~= 6.9e18 < 9.2e18 = Long.MAX.
    _LONG_ARITH_BOUND = 2**61
    _long_arith_strategy = st.one_of(
        st.none(),
        st.integers(min_value=-_LONG_ARITH_BOUND, max_value=_LONG_ARITH_BOUND),
    )

    _bool_strategy = st.one_of(st.none(), st.booleans())

    # 32-bit signed boundaries. Values round-trip through LongType, but the
    # int32 limits are where narrowing / off-by-one bugs in parameter-index
    # plumbing and boundary handling tend to hide, so we always seed them.
    _INT32_MAX = 2**31 - 1
    _INT32_MIN = -(2**31)

    # ---- Edge-case seeds (scalacheck-style) -----------------------------
    #
    # Hypothesis already biases toward "interesting" boundary values, but
    # explicit ``@example`` decorators make a regression on a specific
    # value -- e.g. NULL, zero, the type's max -- deterministic across
    # runs. These are the values we always want to try, before random
    # generation kicks in.
    _LONG_EDGES = (None, 0, 1, -1, 7, -7, _INT32_MAX, _INT32_MIN, _LONG_BOUND, -_LONG_BOUND)
    _LONG_ARITH_EDGES = (None, 0, 1, -1, 7, -7, _LONG_ARITH_BOUND, -_LONG_ARITH_BOUND)
    # Bool space is exhaustive (only three values) so the @example
    # decorators here serve more as documentation of the NULL handling
    # we care about than as new coverage on top of hypothesis's
    # generator.
    _BOOL_EDGES = (None, True, False)
    # Multi-arg edges -- nulls plus the four sign-combos for non-zero
    # values. Catches off-by-one errors in parameter-index plumbing
    # better than random generation alone.
    _LONG_PAIR_EDGES = (
        (None, None),
        (None, 0),
        (0, None),
        (0, 0),
        (1, -1),
        (-1, 1),
        (_INT32_MAX, _INT32_MIN),
        (_INT32_MIN, _INT32_MAX),
        (_LONG_BOUND, 1),
        (1, -_LONG_BOUND),
    )
    _LONG_ARITH_PAIR_EDGES = (
        (None, None),
        (None, 0),
        (0, None),
        (0, 0),
        (1, -1),
        (-1, 1),
        (_LONG_ARITH_BOUND, 1),
        (1, -_LONG_ARITH_BOUND),
    )
    # Sign-combo edges (plus NULL combinations) for the boolean tests.
    # The bodies (``x > 0 and y > 0`` / ``x > 0 or y > 0``) raise in
    # pure Python on a None input (``TypeError``), and the transpiler's
    # NULL-guarded Compare also raises -- so the ``_run`` helper's "both
    # raised" equivalence covers the NULL cases here.
    _BOOLEAN_PAIR_EDGES = (
        (None, None),
        (None, 0),
        (0, None),
        (0, 0),
        (1, -1),
        (-1, 1),
        (1, 1),
        (-1, -1),
        (_LONG_BOUND, 1),
        (1, -_LONG_BOUND),
    )

    def _seed_examples(values, key="value"):
        """Stack one ``@example`` decorator per seed value."""

        def wrapper(method):
            for v in reversed(values):
                method = example(**{key: v})(method)
            return method

        return wrapper

    def _seed_pair_examples(pairs, keys=("x", "y")):
        def wrapper(method):
            for v0, v1 in reversed(pairs):
                method = example(**{keys[0]: v0, keys[1]: v1})(method)
            return method

        return wrapper


# ---- The UDF templates we exercise --------------------------------------
#
# We keep these as module-level callables so that ``inspect.getsource`` works
# (the transpiler reads source via inspection). They are deliberately written
# the way a user would: idiomatic Python, including ``if x is not None``
# guards and bare ``if x:`` truthiness checks.


def plus_four(x):
    if x is not None:
        return x + 4


def plus_four_unsafe(x):
    return x + 4


def plus_four_with_else(x):
    if x is not None:
        return x + 4
    else:
        return 0


def is_none_branch(x):
    if x is None:
        return -1
    else:
        return x


def truthy_bool_branch(x):
    # The transpiler currently mishandles ``if x:`` for nullable bool inputs:
    # Python treats ``None`` as falsy and takes the else branch, but a naive
    # SQL lowering can produce NULL. This test is the canonical regression.
    if x:
        return 1
    else:
        return 2


def add_then_mod(x):
    if x is not None:
        return (x + 7) % 5


def minus_two(x):
    # Exercises ast.Sub.
    if x is not None:
        return x - 2


def times_three(x):
    # Exercises ast.Mult.
    if x is not None:
        return x * 3


def negate_truthy(x):
    # Exercises ast.UnaryOp(Not). Same NULL-as-falsy semantics as
    # ``truthy_bool_branch`` above, just inverted.
    if not x:
        return 0
    else:
        return 1


def both_positive(x, y):
    # Exercises ast.BoolOp(And) over Compare operands -- both operands
    # are statically boolean, so the transpiler should lower to `&`.
    # Kept as a single-statement body since the transpiler doesn't yet
    # support multi-statement function bodies; NULL inputs flow through
    # `>` to NULL on the Spark side and to a raise on the Python side,
    # so the strategy below skips None.
    return x > 0 and y > 0


def either_positive(x, y):
    # Exercises ast.BoolOp(Or) over Compare operands.
    return x > 0 or y > 0


def add_two(x, y):
    # Multi-arg UDF -- exercises the parameter-index plumbing for
    # functions with more than one positional argument.
    if x is not None and y is not None:
        return x + y
    else:
        return 0


def eq_zero(x):
    # Exercises ast.Compare(Eq) via _lower_eq with a non-None left and
    # a literal 0 on the right. The None guard keeps the comparison
    # itself away from NULL operands so the test stays single-path.
    if x is not None:
        return x == 0


def neq_zero(x):
    # Exercises ast.Compare(NotEq) through _lower_eq.
    if x is not None:
        return x != 0


def eq_pair(x, y):
    # Two-arg ``x == y`` exercising the full _lower_eq four-branch when
    # chain that reproduces Python's None-equality semantics:
    #   None == None -> True;  None == 0 -> False;  0 == None -> False.
    # Note: no None guard, so every NULL combination runs through the
    # transpiler's lowering.
    return x == y


def neq_pair(x, y):
    # Sister of ``eq_pair`` for ast.Compare(NotEq).
    return x != y


# A lambda captured at module scope so ``inspect.getsource`` can read
# its definition. Exercises the ``ast.Lambda`` branch in
# ``_get_function_from_ast``.
lambda_plus_four = lambda x: x + 4 if x is not None else 0  # noqa: E731


# ----------------------------------------------------------------------------


@unittest.skipUnless(_have_hypothesis and _hypothesis_enabled and _regular_spark, _skip_reason)
class UDFTranspileHypothesisTests(ReusedSQLTestCase):
    """Compare transpiled vs. interpreted Python UDF output on Hypothesis-generated inputs."""

    # Markers we treat as "transpilation didn't actually happen" -- if any
    # warning matches one of these, the differential comparison would
    # collapse to interpreted-vs-interpreted and pass meaninglessly, so we
    # fail loudly.
    _BAD_TRANSPILE_WARNING_MARKERS = (
        "Unable to transpile",
        "Errors encountered during transpilation",
        "Exception transpiling",
        "ANSI mode",
    )

    def _run(self, func, return_type, df, *udf_arg_columns, kwargs=None):
        """Run ``func`` as a UDF with transpilation on and off, return both rows.

        ``udf_arg_columns`` and ``kwargs`` mirror what a caller would
        write at the dataframe API: positional column names go in
        ``udf_arg_columns`` and named-argument bindings go in ``kwargs``
        (e.g. ``kwargs={"y": "b", "x": "a"}`` to bind UDF parameter ``y``
        to column ``b`` and ``x`` to column ``a``). Use either, or both.

        Asserts the transpiled code path was actually exercised:
        ``transpiled`` must be non-empty after construction, and no
        transpilation-related warning may fire. Without these checks both
        runs could silently fall back to interpreted Python and the
        differential assertion would succeed for the wrong reason.
        """
        func_name = getattr(func, "__name__", repr(func))
        kwargs = kwargs or {}

        transpile_on_conf = {
            "spark.sql.experimental.optimizer.transpilePyUDFs": True,
            # Transpilation requires ANSI; pin it on so the test result
            # doesn't depend on the surrounding session default.
            "spark.sql.ansi.enabled": True,
        }
        transpiled_error: Optional[Exception] = None
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            with self.sql_conf(transpile_on_conf):
                transpiled_udf = UserDefinedFunction(func, return_type)
                self.assertTrue(
                    transpiled_udf.transpiled,
                    f"transpilation produced no Catalyst expression for "
                    f"{func_name!r} -- the differential comparison would be "
                    "meaningless without it",
                )
                try:
                    transpiled_value = df.select(
                        transpiled_udf(*udf_arg_columns, **kwargs)
                    ).collect()[0][0]
                except Exception as e:
                    transpiled_value = _SENTINEL_RAISED
                    transpiled_error = e
            bad = [
                w
                for w in caught
                if any(marker in str(w.message) for marker in self._BAD_TRANSPILE_WARNING_MARKERS)
            ]
            self.assertFalse(
                bad,
                f"unexpected transpile warnings for {func_name!r}: {[str(w.message) for w in bad]}",
            )

        interpreted_error: Optional[Exception] = None
        # Pin ANSI on for the interpreted path too so both sides see the same
        # overflow semantics. Without this, the interpreted path would run with
        # the ambient session default (likely False), causing LongType overflow
        # to silently wrap in Python UDF results while ANSI raises on the
        # transpiled path.
        with self.sql_conf(
            {
                "spark.sql.experimental.optimizer.transpilePyUDFs": False,
                "spark.sql.ansi.enabled": True,
            }
        ):
            interpreted_udf = UserDefinedFunction(func, return_type)
            try:
                interpreted_value = df.select(
                    interpreted_udf(*udf_arg_columns, **kwargs)
                ).collect()[0][0]
            except Exception as e:
                interpreted_value = _SENTINEL_RAISED
                interpreted_error = e

        # If the transpiled path raises an exception we also need the interpreted path to raise one,
        # however if the Python code (that in the interpreted path) raises an exception, the transpiled
        # path may return a valid value.
        if transpiled_error is not None:
            self.assertIsNotNone(
                interpreted_error,
                f"{func_name!r}: transpiled raised {transpiled_error!r} but interpreted did not",
            )
        elif interpreted_error is not None:
            interpreted_value = transpiled_value

        return transpiled_value, interpreted_value

    def _single_arg_df(self, value, dtype):
        schema = StructType([StructField("a", dtype, nullable=True)])
        return self.spark.createDataFrame([Row(a=value)], schema=schema)

    def _two_long_arg_df(self, x, y):
        schema = StructType(
            [
                StructField("a", LongType(), nullable=True),
                StructField("b", LongType(), nullable=True),
            ]
        )
        return self.spark.createDataFrame([Row(a=x, b=y)], schema=schema)

    if _have_hypothesis:

        @_hyp_settings
        @given(value=_long_arith_strategy)
        @_seed_examples(_LONG_ARITH_EDGES)
        def test_plus_four_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(plus_four, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"plus_four mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_arith_strategy)
        @_seed_examples(_LONG_ARITH_EDGES)
        def test_plus_four_unsafe_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(plus_four_unsafe, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"plus_four mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_arith_strategy)
        @_seed_examples(_LONG_ARITH_EDGES)
        def test_plus_four_with_else_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(plus_four_with_else, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"plus_four_with_else mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_strategy)
        @_seed_examples(_LONG_EDGES)
        def test_is_none_branch_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(is_none_branch, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"is_none_branch mismatch on {value!r}")

        @_hyp_settings
        @given(value=_bool_strategy)
        @_seed_examples(_BOOL_EDGES)
        def test_truthy_bool_branch_falls_back(self, value):
            # `if x:` on a bare parameter name is a bare truthiness test whose
            # type is unknown at transpile time. The transpiler must refuse to
            # lower it (Spark's coalesce(x, false) is unsound for non-boolean
            # columns) and fall back to interpreted Python instead.
            df = self._single_arg_df(value, BooleanType())
            with self.sql_conf(
                {
                    "spark.sql.experimental.optimizer.transpilePyUDFs": True,
                    "spark.sql.ansi.enabled": True,
                }
            ):
                pudf = UserDefinedFunction(truthy_bool_branch, LongType())
                self.assertEqual(
                    [],
                    pudf.transpiled,
                    "truthy_bool_branch: bare truthiness test must NOT transpile",
                )
                interpreted = df.select(pudf("a")).collect()[0][0]
            expected = 1 if value else 2
            self.assertEqual(interpreted, expected, f"truthy_bool_branch mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_arith_strategy)
        # add_then_mod is the case that surfaced the Python-vs-SQL mod
        # sign mismatch; the seed values cover the four sign combinations
        # of `(x + 7) % 5` so we always re-prove the pmod fix on every
        # run regardless of the random seed.
        @_seed_examples((*_LONG_ARITH_EDGES, -2, -8, 8, 100, -100))
        def test_add_then_mod_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(add_then_mod, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"add_then_mod mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_arith_strategy)
        @_seed_examples(_LONG_ARITH_EDGES)
        def test_minus_two_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(minus_two, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"minus_two mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_arith_strategy)
        @_seed_examples(_LONG_ARITH_EDGES)
        def test_times_three_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(times_three, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"times_three mismatch on {value!r}")

        @_hyp_settings
        @given(value=_bool_strategy)
        @_seed_examples(_BOOL_EDGES)
        def test_negate_truthy_falls_back(self, value):
            # `if not x:` where x is a bare parameter name is unknown-type at
            # transpile time. The transpiler must refuse (Spark's `~` is
            # bitwise, not Python truthiness) and fall back to interpreted Python.
            df = self._single_arg_df(value, BooleanType())
            with self.sql_conf(
                {
                    "spark.sql.experimental.optimizer.transpilePyUDFs": True,
                    "spark.sql.ansi.enabled": True,
                }
            ):
                pudf = UserDefinedFunction(negate_truthy, LongType())
                self.assertEqual(
                    [],
                    pudf.transpiled,
                    "negate_truthy: bare `not x` must NOT transpile",
                )
                interpreted = df.select(pudf("a")).collect()[0][0]
            expected = 0 if not value else 1
            self.assertEqual(interpreted, expected, f"negate_truthy mismatch on {value!r}")

        @_hyp_settings
        @given(x=_long_arith_strategy, y=_long_arith_strategy)
        @_seed_pair_examples(_LONG_ARITH_PAIR_EDGES)
        def test_add_two_matches_python(self, x, y):
            df = self._two_long_arg_df(x, y)
            transpiled, interpreted = self._run(add_two, LongType(), df, "a", "b")
            self.assertEqual(transpiled, interpreted, f"add_two mismatch on (x={x!r}, y={y!r})")

        @_hyp_settings
        @given(x=_long_arith_strategy, y=_long_arith_strategy)
        @_seed_pair_examples(_LONG_ARITH_PAIR_EDGES)
        def test_add_two_named_args_matches_python(self, x, y):
            # Same UDF as above but called with kwargs (and intentionally
            # in reversed order) to exercise the named-argument codepath
            # documented in the udf() reference. The Python side resolves
            # the kwargs to the function's positional params; the
            # transpiled side has to align ``_udf_param_0`` /
            # ``_udf_param_1`` with the same resolved positions, so any
            # mistake here produces a swapped-argument bug.
            df = self._two_long_arg_df(x, y)
            transpiled, interpreted = self._run(
                add_two,
                LongType(),
                df,
                kwargs={"y": "b", "x": "a"},
            )
            self.assertEqual(
                transpiled,
                interpreted,
                f"add_two named-args mismatch on (x={x!r}, y={y!r})",
            )

        @_hyp_settings
        @given(x=_long_strategy, y=_long_strategy)
        @_seed_pair_examples(_BOOLEAN_PAIR_EDGES)
        def test_both_positive_matches_python(self, x, y):
            df = self._two_long_arg_df(x, y)
            transpiled, interpreted = self._run(both_positive, BooleanType(), df, "a", "b")
            self.assertEqual(
                transpiled, interpreted, f"both_positive mismatch on (x={x!r}, y={y!r})"
            )

        @_hyp_settings
        @given(x=_long_strategy, y=_long_strategy)
        @_seed_pair_examples(_BOOLEAN_PAIR_EDGES)
        def test_either_positive_matches_python(self, x, y):
            df = self._two_long_arg_df(x, y)
            transpiled, interpreted = self._run(either_positive, BooleanType(), df, "a", "b")
            self.assertEqual(
                transpiled, interpreted, f"either_positive mismatch on (x={x!r}, y={y!r})"
            )

        @_hyp_settings
        @given(x=_long_strategy, y=_long_strategy)
        @_seed_pair_examples(_LONG_PAIR_EDGES)
        def test_eq_pair_matches_python(self, x, y):
            # Python's ``==`` has different NULL semantics from SQL ``=``:
            # ``None == None`` is True, ``None == n`` is False. The
            # transpiler's _lower_eq reproduces those semantics, so the
            # transpiled and interpreted runs must agree on every NULL
            # combination as well as the non-NULL cases.
            df = self._two_long_arg_df(x, y)
            transpiled, interpreted = self._run(eq_pair, BooleanType(), df, "a", "b")
            self.assertEqual(transpiled, interpreted, f"eq_pair mismatch on (x={x!r}, y={y!r})")

        @_hyp_settings
        @given(x=_long_strategy, y=_long_strategy)
        @_seed_pair_examples(_LONG_PAIR_EDGES)
        def test_neq_pair_matches_python(self, x, y):
            # Sister of ``test_eq_pair_matches_python`` for the NotEq arm.
            df = self._two_long_arg_df(x, y)
            transpiled, interpreted = self._run(neq_pair, BooleanType(), df, "a", "b")
            self.assertEqual(transpiled, interpreted, f"neq_pair mismatch on (x={x!r}, y={y!r})")

        @_hyp_settings
        @given(value=_long_strategy)
        @_seed_examples(_LONG_EDGES)
        def test_eq_zero_matches_python(self, value):
            # Single-arg ``x == 0`` with a None guard, exercising _lower_eq's
            # non-None-on-both-sides arm.
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(eq_zero, BooleanType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"eq_zero mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_strategy)
        @_seed_examples(_LONG_EDGES)
        def test_neq_zero_matches_python(self, value):
            # Sister of ``test_eq_zero_matches_python`` for the NotEq arm.
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(neq_zero, BooleanType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"neq_zero mismatch on {value!r}")

        @_hyp_settings
        @given(value=_long_arith_strategy)
        @_seed_examples(_LONG_ARITH_EDGES)
        def test_lambda_plus_four_matches_python(self, value):
            df = self._single_arg_df(value, LongType())
            transpiled, interpreted = self._run(lambda_plus_four, LongType(), df, "a")
            self.assertEqual(transpiled, interpreted, f"lambda_plus_four mismatch on {value!r}")


class UDFTranspileHypothesisGatingTests(unittest.TestCase):
    """Smoke tests that always run, regardless of the env gate.

    These don't talk to Spark; they just verify the gating / skipping logic
    so a misconfigured environment doesn't silently skip everything forever.
    """

    def test_env_var_name_is_documented(self):
        # If we ever rename the env var, the docstring needs to follow.
        self.assertIn(_HYPOTHESIS_ENV, __doc__)

    def test_skip_reason_mentions_env_var(self):
        self.assertIn(_HYPOTHESIS_ENV, _skip_reason)

    def test_env_gate_is_value_based(self):
        # CI always sets RUN_HYPOTHESIS (to "true" or "false") via the
        # transpile precondition in build_and_test.yml. If this gate ever
        # regresses to presence-based, the very slow suite silently runs on
        # every PySpark job. Pin the contract from both directions.
        for opted_in in ("1", "true", "TRUE", "yes", " True "):
            self.assertTrue(_env_opts_in(opted_in), f"{opted_in!r} must opt in")
        for opted_out in (None, "", "0", "false", "FALSE", "no", "off"):
            self.assertFalse(_env_opts_in(opted_out), f"{opted_out!r} must NOT opt in")


if __name__ == "__main__":
    from pyspark.testing import main

    main()