From d839a27c9fec56f3bf39d5fcb8e607bb55686c11 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Fri, 7 Aug 2026 19:13:07 +0000 Subject: [PATCH 01/18] [SPARK-58456][PYTHON][SQL] Support chained comparisons and in / not in in the Python UDF transpiler Teaches the experimental Python UDF transpiler to lower two more comparison forms, scoped to the default system types it already handles: - Chained comparisons (`0 <= x <= 100`, `a < b < c`), up to `_MAX_CHAIN_OPS` links. - `in` / `not in` over a literal tuple / list / set of constants, up to `_MAX_IN_ELEMENTS` elements. Both fold through new `_lazy_and` / `_lazy_or` helpers that nest `CASE WHEN`s instead of folding with `&` / `|`. That is required rather than stylistic: the comparison lowerings emit a `raise_error` guard to mirror Python's `TypeError` on a NULL operand, and `splitConjunctivePredicates` decomposes `And` while `PushPredicateThroughJoin` gates pushdown on `!cond.throwable` -- which `RaiseError` does not override. An `And` fold therefore lets the guard be pushed across a join and fire on rows Python never evaluated. `and` / `or` move onto the same lazy fold, which fixes that latent divergence for them too, and their operand gate tightens from `_is_definitely_boolean` to the new `_is_never_null_boolean` since a NULL `CASE WHEN` condition reads as not-taken and would turn Python's `None` into `False`. Also fixes `_category` returning "numeric" for an `ast.Compare` whose operands are not basic types (a ternary or nested compare subject). That let the return-type match accept a numeric declared return type for a boolean body, where the lowered option returned `cast( as bigint)` (1/0) but interpreted Python returns NULL. Anything outside the supported shapes fails closed and falls back to interpreted Python: `in` on a str/bytes (substring, not equality), dict membership, runtime-valued containers, `in` inside a chain, and over-cap sizes. Chained comparisons and membership tests are common in real UDF bodies, and previously forced the whole UDF back to interpreted Python. The `and` / `or` pushdown hazard is a pre-existing correctness bug on the transpiled path. Only under the experimental, off-by-default `spark.sql.experimental.optimizer.transpilePyUDFs` (which also requires ANSI). More UDF bodies transpile to native Catalyst expressions. One body that previously transpiled now falls back: `(True if c else None) and (y > 0)`, whose old lowering diverged from Python. `test_udf_transpile_unit.py`: chained/`in` lowering and plan elision, a short-circuit test run with whole-stage codegen both on and off, all NULL combinations for `x == y == z`, empty-container behaviour pinned under both values of `spark.sql.legacy.nullInEmptyListBehavior`, cap boundaries at N and N+1, a must-not-transpile set that also asserts each query still returns Python's answer, and a plan-shape regression test asserting the raise guard is not pushed below a Join. `test_udf_transpile_hypothesis.py`: nine new UDF templates and property tests comparing transpiled against interpreted results, seeded with all eight NULL combinations for the three-operand chains. Passed at 100 examples per test, and previously at 1000. Two tests are DELIBERATELY FAILING and left red on purpose, pinning a wrong-results bug whose fix is a companion change: `test_udf_transpile_nondeterministic_argument_falls_back` and `ResolveTranspiledPythonUDFOptionsSuite`'s non-deterministic case. Because `resolveUDFParams` substitutes a call-site argument at every `_udf_param_N` occurrence, a non-deterministic argument becomes several independent expressions; with lazy evaluation they desynchronize, so `50 < x < 60` over `rand(7) * 100` yields ~30% true where Python yields ~10%. This is not fixable from Python (transpilation predates any call site) nor in the JVM builder (arguments are still unresolved there); it needs an analyzer-side guard. Scoping that honestly: `and` / `or` were already wrong this way, but a single comparison is not (its duplicated operand only feeds an `IS NULL` test, whose outcome is value-independent), and chained comparisons previously fell back to interpreted Python and were therefore correct. So this PR widens the bug's reach rather than merely coexisting with it, which argues for landing the analyzer guard before or with it. Separately, moving `and` / `or` off `&` / `|` costs predicate pushdown for non-nullable columns: the old fold simplified to a plain `And` that `splitConjunctivePredicates` could split, while `CaseWhen(Seq((c1, c2)), Some(false))` has no such rewrite and stays opaque. That case is also the one where `And` was already safe (the guard folds away, so no `raise_error` remains), so overriding `throwable` on `RaiseError` would recover both properties; tracked separately. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../tests/test_udf_transpile_hypothesis.py | 215 +++++- .../sql/tests/test_udf_transpile_unit.py | 631 +++++++++++++++++- python/pyspark/sql/transpile.py | 533 ++++++++++++--- ...solveTranspiledPythonUDFOptionsSuite.scala | 38 +- 4 files changed, 1281 insertions(+), 136 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py index c0db082821ca6..38a3a22cb978b 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -206,6 +206,30 @@ def _env_opts_in(value: Optional[str]) -> bool: (1, -_LONG_BOUND), ) + # Triple edges for the chained-comparison tests. The eight NULL + # combinations are the point: a chain short-circuits, so whether a NULL + # operand's raise guard is reached depends on the earlier links. The + # ordering chains raise in pure Python on any NULL they reach (covered by + # ``_run``'s "both raised" equivalence), while ``x == y == z`` has a + # defined answer for every combination because Python's ``==`` never + # raises on None. + _LONG_TRIPLE_EDGES = ( + (None, None, None), + (None, None, 0), + (None, 0, None), + (0, None, None), + (None, 0, 1), + (0, None, 1), + (0, 1, None), + (1, 0, None), + (0, 0, 0), + (0, 1, 2), + (2, 1, 0), + (1, 1, 1), + (_INT32_MIN, 0, _INT32_MAX), + (-_LONG_BOUND, 0, _LONG_BOUND), + ) + def _seed_examples(values, key="value"): """Stack one ``@example`` decorator per seed value.""" @@ -216,14 +240,20 @@ def wrapper(method): return wrapper - def _seed_pair_examples(pairs, keys=("x", "y")): + def _seed_tuple_examples(tuples, keys=("x", "y", "z")): + """Stack one ``@example`` decorator per seed tuple, zipped onto ``keys``.""" + def wrapper(method): - for v0, v1 in reversed(pairs): - method = example(**{keys[0]: v0, keys[1]: v1})(method) + for values in reversed(tuples): + method = example(**dict(zip(keys, values)))(method) return method return wrapper + def _seed_pair_examples(pairs, keys=("x", "y")): + """Two-argument ``_seed_tuple_examples``.""" + return _seed_tuple_examples(pairs, keys) + # ---- The UDF templates we exercise -------------------------------------- # @@ -343,6 +373,73 @@ def neq_pair(x, y): return x != y +def chained_bounds(x): + # Single-parameter chained comparison: ``ast.Compare`` with two ops, which + # Python evaluates as ``(0 < x) and (x < 10)`` with ``x`` evaluated once. + # No None guard, so a NULL input reaches the first link's raise guard -- + # matching Python's TypeError -- and ``_run``'s "both raised" equivalence + # covers it. + return 0 < x < 10 + + +def chained_lt(x, y, z): + # Three-column ordering chain. The interesting rows are the ones where an + # earlier link is False and a later operand is NULL: Python short-circuits + # and returns False without ever comparing against the NULL, so the + # transpiled form must not raise there either. + return x < y < z + + +def chained_ge(x, y, z): + # Reversed-direction chain, so the short-circuit fires on a different set + # of rows than ``chained_lt``. + return x >= y >= z + + +def chained_eq(x, y, z): + # ``x == y == z`` -- equality never raises in Python, so every one of the + # eight NULL combinations has a defined boolean answer that the + # four-branch ``_lower_eq`` per link plus the short-circuit fold must + # reproduce exactly. This is the strongest NULL-semantics probe of the + # chained lowering. + return x == y == z + + +def mixed_chain(x, y, z): + # Mixed operators in one chain (``ast.LtE`` then ``ast.Gt``), so the two + # links go through different lowering helpers. + return x <= y > z + + +def in_literal(x): + # ``ast.Compare(In)`` over a constant tuple. Python's ``in`` is + # equality-based and never raises on a None probe (``None in (1, 2, 3)`` is + # False), unlike Spark's ``IN``, which returns NULL for a NULL probe. + return x in (1, 2, 3) + + +def not_in_literal(x): + # Sister of ``in_literal`` for ast.Compare(NotIn); also exercises the + # negation, which is only exact because the ``in`` lowering never yields + # NULL. + return x not in (1, 2, 3) + + +def in_with_none(x): + # A None element in the container: ``None in (1, None)`` is True in Python, + # while Spark's ``IN`` would return NULL for any non-matching probe if the + # NULL element were left in the list. + return x in (1, None) + + +def in_and_compare(x, y): + # ``in`` as an ``and`` operand, so the container membership has to be + # recognized as a non-NULL boolean by the BoolOp gate. When the membership + # test is False, Python never evaluates ``y > 0``, so a NULL ``y`` must not + # raise. + return x in (1, 2) and y > 0 + + # A lambda captured at module scope so ``inspect.getsource`` can read # its definition. Exercises the ``ast.Lambda`` branch in # ``_get_function_from_ast``. @@ -466,6 +563,16 @@ def _two_long_arg_df(self, x, y): ) return self.spark.createDataFrame([Row(a=x, b=y)], schema=schema) + def _three_long_arg_df(self, x, y, z): + schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + StructField("c", LongType(), nullable=True), + ] + ) + return self.spark.createDataFrame([Row(a=x, b=y, c=z)], schema=schema) + if _have_hypothesis: @_hyp_settings @@ -664,6 +771,108 @@ def test_lambda_plus_four_matches_python(self, value): transpiled, interpreted = self._run(lambda_plus_four, LongType(), df, "a") self.assertEqual(transpiled, interpreted, f"lambda_plus_four mismatch on {value!r}") + # ---- Chained comparisons ---------------------------------------- + # + # Python's ``a OP b OP c`` short-circuits: when the first link is falsy + # the second is never evaluated, so a NULL operand it would have + # touched must not raise. ``_run`` catches the dangerous direction of a + # regression here -- if the transpiled form raises where interpreted + # Python does not, ``_run`` asserts on it. (The reverse direction, + # Python raising where the transpiled form returns a value, is + # tolerated by ``_run``, so it is pinned by explicit assertions in + # test_udf_transpile_unit.py instead.) + + @_hyp_settings + @given(value=_long_strategy) + @_seed_examples(_LONG_EDGES) + def test_chained_bounds_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(chained_bounds, BooleanType(), df, "a") + self.assertEqual(transpiled, interpreted, f"chained_bounds mismatch on {value!r}") + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy, z=_long_strategy) + @_seed_tuple_examples(_LONG_TRIPLE_EDGES) + def test_chained_lt_matches_python(self, x, y, z): + df = self._three_long_arg_df(x, y, z) + transpiled, interpreted = self._run(chained_lt, BooleanType(), df, "a", "b", "c") + self.assertEqual( + transpiled, interpreted, f"chained_lt mismatch on ({x!r}, {y!r}, {z!r})" + ) + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy, z=_long_strategy) + @_seed_tuple_examples(_LONG_TRIPLE_EDGES) + def test_chained_ge_matches_python(self, x, y, z): + df = self._three_long_arg_df(x, y, z) + transpiled, interpreted = self._run(chained_ge, BooleanType(), df, "a", "b", "c") + self.assertEqual( + transpiled, interpreted, f"chained_ge mismatch on ({x!r}, {y!r}, {z!r})" + ) + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy, z=_long_strategy) + @_seed_tuple_examples(_LONG_TRIPLE_EDGES) + def test_chained_eq_matches_python(self, x, y, z): + # Equality never raises in Python, so every NULL combination has a + # defined boolean answer on both sides -- no "both raised" + # escape hatch can mask a divergence here. + df = self._three_long_arg_df(x, y, z) + transpiled, interpreted = self._run(chained_eq, BooleanType(), df, "a", "b", "c") + self.assertEqual( + transpiled, interpreted, f"chained_eq mismatch on ({x!r}, {y!r}, {z!r})" + ) + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy, z=_long_strategy) + @_seed_tuple_examples(_LONG_TRIPLE_EDGES) + def test_mixed_chain_matches_python(self, x, y, z): + df = self._three_long_arg_df(x, y, z) + transpiled, interpreted = self._run(mixed_chain, BooleanType(), df, "a", "b", "c") + self.assertEqual( + transpiled, interpreted, f"mixed_chain mismatch on ({x!r}, {y!r}, {z!r})" + ) + + # ---- `in` / `not in` --------------------------------------------- + + @_hyp_settings + @given(value=_long_strategy) + @_seed_examples(_LONG_EDGES) + def test_in_literal_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(in_literal, BooleanType(), df, "a") + self.assertEqual(transpiled, interpreted, f"in_literal mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_strategy) + @_seed_examples(_LONG_EDGES) + def test_not_in_literal_matches_python(self, value): + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(not_in_literal, BooleanType(), df, "a") + self.assertEqual(transpiled, interpreted, f"not_in_literal mismatch on {value!r}") + + @_hyp_settings + @given(value=_long_strategy) + @_seed_examples(_LONG_EDGES) + def test_in_with_none_matches_python(self, value): + # ``None in (1, None)`` is True in Python. The lowering answers the + # NULL-probe case with a literal and drops the NULL element from + # the emitted ``IN`` list, so this must agree for a NULL probe as + # well as for matching and non-matching values. + df = self._single_arg_df(value, LongType()) + transpiled, interpreted = self._run(in_with_none, BooleanType(), df, "a") + self.assertEqual(transpiled, interpreted, f"in_with_none mismatch on {value!r}") + + @_hyp_settings + @given(x=_long_strategy, y=_long_strategy) + @_seed_pair_examples(_BOOLEAN_PAIR_EDGES) + def test_in_and_compare_matches_python(self, x, y): + df = self._two_long_arg_df(x, y) + transpiled, interpreted = self._run(in_and_compare, BooleanType(), df, "a", "b") + self.assertEqual( + transpiled, interpreted, f"in_and_compare mismatch on (x={x!r}, y={y!r})" + ) + class UDFTranspileHypothesisGatingTests(unittest.TestCase): """Smoke tests that always run, regardless of the env gate. diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index dce2e9eeefc4c..69465973f260b 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -270,8 +270,9 @@ def func_closure_capture(x): def test_udf_transpile_boolean_and_or_lowered(self): # When `and`/`or` operands are syntactically boolean (Compare - # results in this case), the transpiler should lower to bitwise - # `&`/`|` and produce results matching the interpreted UDF. + # results in this case), the transpiler should lower to a nested + # CASE WHEN that reproduces Python's short-circuit and produce + # results matching the interpreted UDF. # Each UDF is a single top-level statement (the transpiler # doesn't support multi-statement bodies yet). from pyspark.sql.types import StructField, StructType @@ -295,11 +296,11 @@ def either_positive(x, y): "spark.sql.ansi.enabled": True, } ): - # NULL inputs propagate through `>` to NULL, which then - # passes through `&` / `|` per SQL three-valued logic. We - # only assert on non-NULL inputs here since Python's - # interpreted `x > 0 and y > 0` would raise on None; the - # NULL handling itself is covered by the hypothesis suite. + # A NULL operand makes the lowered `>` raise (mirroring Python's + # TypeError), unless an earlier operand already short-circuited the + # result. We only assert on non-NULL inputs here; the NULL handling + # itself is covered by test_udf_transpile_boolop_short_circuits and + # the hypothesis suite. for func, x, y, expected in [ (both_positive, 1, 2, True), (both_positive, 1, -1, False), @@ -491,40 +492,437 @@ def gte_zero(x): self.assertEqual(row_lte[0], lte_expected) self.assertEqual(row_gte[0], gte_expected) - def test_udf_transpile_chained_comparison_falls_back(self): - # ``a < b < c`` is a chained comparison: Python evaluates it as - # ``(a < b) and (b < c)``. The transpiler refuses chained Compare - # nodes (``len(ops) != 1``) and must fall back to interpreted Python. - import warnings as _warnings - + def test_udf_transpile_chained_comparison_lowered(self): + # ``0 < x < 10`` is a chained comparison: Python evaluates it as + # ``(0 < x) and (x < 10)`` with ``x`` evaluated once. The transpiler + # folds the links into a nested CASE WHEN, so it transpiles and matches + # interpreted Python. A NULL input still raises through the ordering + # guard on the first link (Python would raise TypeError), so the values + # asserted here are the non-NULL ones. from pyspark.sql.types import StructField, StructType def chained(x): return 0 < x < 10 schema = StructType([StructField("a", LongType(), nullable=False)]) - with self.sql_conf( - { - "spark.sql.experimental.optimizer.transpilePyUDFs": True, - "spark.sql.ansi.enabled": True, - } - ): - with _warnings.catch_warnings(record=True) as caught: - _warnings.simplefilter("always") - pudf = UserDefinedFunction(chained, BooleanType()) - self.assertEqual([], pudf.transpiled, "chained comparison must NOT transpile") - fallback = [ - w - for w in caught - if "Unable to transpile" in str(w.message) or "Errors encountered" in str(w.message) - ] - self.assertTrue(fallback, "expected a fallback warning") - for value, expected in [(5, True), (0, False), (10, False), (-3, False)]: + with self.sql_conf(_TRANSPILE_ON): + pudf = UserDefinedFunction(chained, BooleanType()) + self.assertTrue(pudf.transpiled, "chained comparison should transpile") + for value, expected in [(5, True), (0, False), (10, False), (-3, False), (9, True)]: with self.subTest(value=value): df = self.spark.createDataFrame([Row(a=value)], schema=schema) [row] = df.select(pudf("a")).collect() self.assertEqual(row[0], expected) + def test_udf_transpile_chain_short_circuits(self): + # THE test for the chained lowering. Python's chain short-circuits: when + # ``a < b`` is false, ``b < c`` is never evaluated, so a NULL ``c`` + # cannot raise. The transpiler must reproduce that -- which is why the + # links are folded into a nested CASE WHEN (whose branches Catalyst + # evaluates lazily and which predicate splitting cannot take apart) + # rather than `&`-ed together. + # + # Note the literal form (`1 < 0 < None`) proves nothing here: a literal + # None has no column category, so `_lower_value_compare` refuses the + # whole UDF and it falls back. Nullable *columns* are required. + # + # Run with whole-stage codegen both on and off: `CaseWhen` implements + # laziness twice over (`eval` walks branches until one is TRUE, and the + # single-branch codegen path defers to `If.doGenCode`, which emits each + # branch's code inside the generated `if`/`else`). A regression in + # either implementation must fail this test. + from pyspark.sql.types import StructField, StructType + + def chained(a, b, c): + return a < b < c + + schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + StructField("c", LongType(), nullable=True), + ] + ) + for wholestage in (True, False): + conf = dict(_TRANSPILE_ON) + conf["spark.sql.codegen.wholeStage"] = wholestage + with self.subTest(wholestage=wholestage), self.sql_conf(conf): + pudf = UserDefinedFunction(chained, BooleanType()) + self.assertTrue(pudf.transpiled, "3-column chain should transpile") + df = self.spark.createDataFrame([Row(a=1, b=0, c=None)], schema=schema) + [row] = df.select(pudf("a", "b", "c")).collect() + self.assertEqual( + False, + row[0], + "first link is false, so the NULL second link must not be evaluated", + ) + # ... but when the first link IS true the second link is reached + # and its NULL guard fires, exactly as Python raises TypeError. + df = self.spark.createDataFrame([Row(a=0, b=1, c=None)], schema=schema) + with self.assertRaises(Exception) as ctx: + df.select(pudf("a", "b", "c")).collect() + self.assertIn("cannot compare NULL", str(ctx.exception)) + + def test_udf_transpile_boolop_short_circuits(self): + # The `and` / `or` counterpart of the chain short-circuit, and the + # regression test for the pushdown hazard the CASE WHEN lowering closes. + # + # `x > 0 and y > 0` used to lower to `And(caseWhen1, caseWhen2)`. + # `splitConjunctivePredicates` decomposes `And` and nothing else, and + # `PushPredicateThroughJoin` treats a conjunct as pushable when + # `deterministic && !throwable` -- but `RaiseError` does not override + # `Expression.throwable`, so the NULL guard looks harmless and the + # second conjunct gets pushed to the far side of a join, where it is + # evaluated on rows the first conjunct excluded. That can raise where + # Python short-circuits to False. A single CASE WHEN is atomic to + # predicate splitting, so it cannot be separated that way. + # + # The guard below is a PLAN-SHAPE assertion, deliberately, because the + # runtime symptom is scheduling-dependent: with the old `&` lowering the + # pushed-down `raise_error` fires only on some runs (whether the failing + # task gets far enough before the join yields no rows). Asserting on the + # optimized plan instead is deterministic -- the raise guard must stay + # ABOVE the join rather than being pushed beneath it. + def both_positive(a, c): + return a > 0 and c > 0 + + def either_positive(a, c): + return a > 0 or c > 0 + + with self.sql_conf(_TRANSPILE_ON): + and_udf = UserDefinedFunction(both_positive, BooleanType()) + or_udf = UserDefinedFunction(either_positive, BooleanType()) + self.assertTrue(and_udf.transpiled, "and over compares should transpile") + self.assertTrue(or_udf.transpiled, "or over compares should transpile") + + # Plain select: the second operand is NULL but the first already + # decided the result, so Python's answer must come back unraised. + df = self.spark.createDataFrame([(-5, None)], "a long, c long") + self.assertEqual([False], [r[0] for r in df.select(and_udf("a", "c")).collect()]) + df = self.spark.createDataFrame([(5, None)], "a long, c long") + self.assertEqual([True], [r[0] for r in df.select(or_udf("a", "c")).collect()]) + + # Same thing across a join, where the NULL operand lives entirely on + # the right side and would be a pushdown candidate if the predicate + # decomposed into conjuncts. + left = self.spark.createDataFrame([(1, -5)], "k long, a long") + right = self.spark.createDataFrame([(1, None)], "k long, c long") + joined = left.join(right, "k") + self.assertEqual([False], [r[0] for r in joined.select(and_udf("a", "c")).collect()]) + + filtered = joined.filter(and_udf("a", "c")) + plan = filtered._jdf.queryExecution().optimizedPlan().toString() + # An optimized plan prints root-first, depth-first, so anything + # pushed below the join appears on a LATER line than the Join node. + lines = plan.split("\n") + join_line = next( + (i for i, line in enumerate(lines) if "Join" in line), + None, + ) + self.assertIsNotNone(join_line, f"expected a Join in the plan:\n{plan}") + pushed = [ + line for i, line in enumerate(lines) if i > join_line and "raise_error" in line + ] + self.assertEqual( + [], + pushed, + "the NULL raise guard was pushed below the Join, so it can be " + f"evaluated on rows the other operand excluded:\n{plan}", + ) + # And it still returns Python's answer rather than raising. + self.assertEqual([], filtered.select("k").collect()) + + def test_udf_transpile_chain_operator_mix_and_eq_nulls(self): + # Mixed operators in one chain, and `x == y == z` across every NULL + # combination. Equality never raises in Python (`None == None` is True, + # `None == 0` is False), so all eight combinations have a defined answer + # and the transpiled result must match it exactly. + from pyspark.sql.types import StructField, StructType + + schema = StructType( + [ + StructField("a", LongType(), nullable=True), + StructField("b", LongType(), nullable=True), + StructField("c", LongType(), nullable=True), + ] + ) + + def chained_eq(a, b, c): + return a == b == c + + # (a, b, c, expected) -- the eight NULL combinations plus value cases. + eq_cases = [ + (None, None, None, True), + (None, None, 1, False), + (None, 1, None, False), + (None, 1, 2, False), + (1, None, None, False), + (1, None, 2, False), + (1, 2, None, False), + (1, 1, None, False), + (1, 1, 1, True), + (1, 1, 2, False), + (1, 2, 3, False), + ] + with self.sql_conf(_TRANSPILE_ON): + pudf = UserDefinedFunction(chained_eq, BooleanType()) + self.assertTrue(pudf.transpiled, "a == b == c should transpile") + for a, b, c, expected in eq_cases: + with self.subTest(a=a, b=b, c=c): + df = self.spark.createDataFrame([Row(a=a, b=b, c=c)], schema=schema) + [row] = df.select(pudf("a", "b", "c")).collect() + self.assertEqual(row[0], expected, f"a == b == c on ({a}, {b}, {c})") + + # Operator mixes, on non-NULL rows so the ordering guards stay quiet. + B = BooleanType() + lt_eq = lambda a, b, c: a < b == c # noqa: E731 + eq_lt = lambda a, b, c: a == b < c # noqa: E731 + gt_gte = lambda a, b, c: a > b >= c # noqa: E731 + lte_lte = lambda a, b, c: a <= b <= c # noqa: E731 + four_ops = lambda a, b, c, d, e: a < b < c < d < e # noqa: E731 + for func, rows, expected in [ + (lt_eq, [(1, 2, 2), (1, 2, 3), (2, 1, 1)], [True, False, False]), + (eq_lt, [(1, 1, 2), (1, 1, 0), (1, 2, 3)], [True, False, False]), + (gt_gte, [(3, 2, 2), (3, 2, 5), (1, 2, 2)], [True, False, False]), + (lte_lte, [(1, 1, 2), (1, 2, 2), (2, 1, 3)], [True, True, False]), + ]: + with self.subTest(func=str(func)): + self.assertEqual(self._vals(func, B, "a long, b long, c long", rows), expected) + self.assertEqual( + self._vals( + four_ops, + B, + "a long, b long, c long, d long, e long", + [(1, 2, 3, 4, 5), (1, 2, 3, 5, 4)], + ), + [True, False], + "a chain of four operators is at the cap and must transpile", + ) + + def test_udf_transpile_in_lowered(self): + # ``in`` / ``not in`` over a literal container is equality-based in + # Python and never raises on a None probe: ``None in (1, 2)`` is False + # and ``None in (1, None)`` is True. Spark's ``IN`` returns NULL for a + # NULL probe or a NULL list element, so the lowering answers the + # NULL-probe case directly and drops NULL elements from the list. + from pyspark.sql.types import StructField, StructType + + schema = StructType([StructField("a", LongType(), nullable=True)]) + in_none = lambda x: x in (1, None) # noqa: E731 + not_in_none = lambda x: x not in (1, None) # noqa: E731 + in_plain = lambda x: x in (1, 2) # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + for func, cases in [ + (in_plain, [(None, False), (1, True), (2, True), (3, False)]), + (in_none, [(None, True), (1, True), (2, False)]), + (not_in_none, [(None, False), (1, False), (2, True)]), + ]: + pudf = UserDefinedFunction(func, BooleanType()) + self.assertTrue(pudf.transpiled, f"{func} should transpile") + for value, expected in cases: + with self.subTest(func=str(func), value=value): + df = self.spark.createDataFrame([Row(a=value)], schema=schema) + [row] = df.select(pudf("a")).collect() + self.assertEqual(row[0], expected) + + def test_udf_transpile_in_empty_container(self): + # ``x in ()`` is always False in Python, including for a None probe. + # The lowering spells that out with a literal rather than emitting + # ``IN ()``, whose NULL-probe result depends on + # ``spark.sql.legacy.nullInEmptyListBehavior``. Pin it under both + # settings so the answer cannot drift with that flag. + from pyspark.sql.types import StructField, StructType + + empty_in = lambda x: x in () # noqa: E731 + empty_not_in = lambda x: x not in () # noqa: E731 + schema = StructType([StructField("a", LongType(), nullable=True)]) + for legacy in (False, True): + conf = dict(_TRANSPILE_ON) + conf["spark.sql.legacy.nullInEmptyListBehavior"] = legacy + with self.sql_conf(conf): + for func, expected in [(empty_in, False), (empty_not_in, True)]: + with self.subTest(legacy=legacy, func=str(func)): + pudf = UserDefinedFunction(func, BooleanType()) + self.assertTrue(pudf.transpiled, f"{func} should transpile") + df = self.spark.createDataFrame([Row(a=None)], schema=schema) + [row] = df.select(pudf("a")).collect() + self.assertEqual(row[0], expected) + + def test_udf_transpile_comparison_fallbacks(self): + # Comparison forms the transpiler must REFUSE, each for its own reason. + # Every one asserts both that nothing transpiled and that a query using + # the UDF still runs and returns Python's answer -- the second half + # matters because a surviving option is type-checked by CheckAnalysis as + # a child of TranspiledPythonUDF, so a bad lowering fails the whole + # query instead of falling back (e.g. `x in (1, "a")` would raise + # DATA_DIFF_TYPES). + import warnings as _warnings + + # `in` on a str/bytes is substring/subsequence containment, not + # equality; dict membership tests keys; a container of runtime values + # would need runtime NULL analysis; `in` cannot be chained. + str_in = lambda x: x in "abc" # noqa: E731 + bytes_in = lambda x: x in b"ab" # noqa: E731 + dict_in = lambda x: x in {1: "a"} # noqa: E731 + runtime_in = lambda x, y: x in (y, 2) # noqa: E731 + chained_in = lambda a, b: a < b in (1, 2) # noqa: E731 + mixed_in = lambda x: x in (1, "a") # noqa: E731 + bool_in = lambda x: x in (True, False) # noqa: E731 + cross_chain = lambda x: 0 < x < "z" # noqa: E731 + five_ops = lambda a, b, c, d, e, f: a < b < c < d < e < f # noqa: E731 + chain_long_rt = lambda x: 0 < x < 10 # noqa: E731 + in_long_rt = lambda x: x in (1, 2) # noqa: E731 + six_longs = "a long, b long, c long, d long, e long, f long" + + # (func, schema, rows, expected). `expected` is what interpreted Python + # produces, so the query must still run and return it. + cases = [ + (str_in, "a string", [("b",), ("z",)], [True, False]), + (bytes_in, "a long", [(97,), (1,)], [True, False]), + (dict_in, "a long", [(1,), (2,)], [True, False]), + (runtime_in, "a long, b long", [(1, 1), (3, 5)], [True, False]), + (chained_in, "a long, b long", [(0, 1), (1, 0)], [True, False]), + (mixed_in, "a long", [(1,), (3,)], [True, False]), + (bool_in, "a long", [(1,), (5,)], [True, False]), + (five_ops, six_longs, [(1, 2, 3, 4, 5, 6)], [True]), + ] + with self.sql_conf(_TRANSPILE_ON): + for func, schema, rows, expected in cases: + with self.subTest(func=str(func)): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, BooleanType()) + self.assertEqual([], pudf.transpiled, f"{func} must NOT transpile") + df = self.spark.createDataFrame(rows, schema) + self.assertEqual( + [r[0] for r in df.select(pudf(*df.columns)).collect()], + expected, + f"{func}: interpreted result diverged", + ) + # `0 < x < "z"` raises in interpreted Python too (cross-type + # ordering), so assert only the fallback for it. + # + # A chained / `in` body is boolean, so a non-BooleanType declared + # return type must drop every variant rather than emit a cast the + # interpreted path would never perform. + for func, rt, label in [ + (cross_chain, BooleanType(), "a cross-category chain"), + (chain_long_rt, LongType(), "a chain with a LongType return type"), + (in_long_rt, LongType(), "an `in` with a LongType return type"), + ]: + with self.subTest(label=label): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + self.assertEqual( + [], + UserDefinedFunction(func, rt).transpiled, + f"{label} must NOT transpile", + ) + + def test_udf_transpile_boolean_body_requires_boolean_return_type(self): + # A comparison body is boolean no matter what its operands look like, so + # it must not satisfy a numeric declared return type. + # + # `_is_definitely_boolean`'s Compare arm additionally requires basic-type + # operands, so it answers False for a comparison whose operand is itself + # a ternary / comparison / boolean op. Those used to fall through + # `_category`'s numeric catch-all, which let the return-type match in + # `_transpile_from_ast` accept `LongType`: the lowered option returned + # `cast( as bigint)` -- 1 or 0 -- while the interpreted + # SQL_BATCHED_UDF path returns NULL for a Python bool declared LongType + # (EvaluatePython.makeFromJava nulls type-mismatched results). That is a + # silent divergence, exactly what the return-type match exists to stop. + # `_category` now classifies every `ast.Compare` as "bool" directly. + import warnings as _warnings + + def ternary_in(c, x, y): + # `in` whose subject is a ternary -- reachable only since `in` is + # lowered at all, hence pinned here. + return (x if c > 0 else y) in (1, 2) + + def eq_true(a, b): + # The same gap via `_lower_eq`; pre-dates `in` support. + return (a < b) == True # noqa: E712 + + def chain_of_compare(a, b, c): + # ... and via a chained comparison over a comparison operand. + return (a < b) <= (b < c) + + cases = [ + (ternary_in, "c long, x long, y long", [(1, 1, 9)]), + (eq_true, "a long, b long", [(1, 2)]), + (chain_of_compare, "a long, b long, c long", [(1, 2, 3)]), + ] + with self.sql_conf(_TRANSPILE_ON): + for func, schema, rows in cases: + with self.subTest(func=func.__name__): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, LongType()) + self.assertEqual( + [], + pudf.transpiled, + f"{func.__name__}: a boolean body must not satisfy LongType", + ) + # The interpreted path returns NULL for a bool declared + # LongType; falling back reproduces that faithfully. + df = self.spark.createDataFrame(rows, schema) + self.assertEqual( + [None], + [r[0] for r in df.select(pudf(*df.columns)).collect()], + f"{func.__name__}: interpreted result diverged", + ) + # Declared BooleanType, the same bodies transpile and are correct. + df = self.spark.createDataFrame([(1, 1, 9)], "c long, x long, y long") + pudf = UserDefinedFunction(ternary_in, BooleanType()) + self.assertTrue(pudf.transpiled, "ternary_in should transpile for BooleanType") + self.assertEqual([True], [r[0] for r in df.select(pudf("c", "x", "y")).collect()]) + + def test_udf_transpile_lowering_caps(self): + # The chain-operator and container-element caps bound plan growth: an + # interior chain operand's subtree is emitted once per adjacent pair, + # and every input-category variant carries its own copy for + # CheckAnalysis to walk. Exercise both boundaries directly rather than + # writing a 65-element tuple into the test source (a container built + # with `eval` has no retrievable source, so the UDF would fall back for + # the wrong reason and the test would pass vacuously). + import ast as _ast + + from pyspark.errors import UnsupportedOperationException + from pyspark.sql.transpile import ( + CatalystTranspiler, + _MAX_CHAIN_OPS, + _MAX_IN_ELEMENTS, + ) + + transpiler = CatalystTranspiler() + transpiler._param_categories = {0: "numeric"} + + def compare_node(src): + return _ast.parse(src, mode="eval").body + + def in_src(n): + return "x in (%s)" % ", ".join(str(i) for i in range(n)) + + def chain_src(n): + # n operators over n + 1 operands: x < 0 < 1 < ... < n - 1 + return " < ".join(["x"] + [str(i) for i in range(n)]) + + with self.sql_conf(_TRANSPILE_ON): + # At each cap the lowering succeeds. + node = compare_node(in_src(_MAX_IN_ELEMENTS)) + transpiler._lower_in(["x"], node.left, node.comparators[0], negate=False) + transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_CHAIN_OPS))) + + # One past it, it refuses so the UDF falls back. + with self.assertRaises(UnsupportedOperationException) as ctx: + node = compare_node(in_src(_MAX_IN_ELEMENTS + 1)) + transpiler._lower_in(["x"], node.left, node.comparators[0], negate=False) + self.assertIn("elements", str(ctx.exception)) + with self.assertRaises(UnsupportedOperationException) as ctx: + transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_CHAIN_OPS + 1))) + self.assertIn("operators", str(ctx.exception)) + def test_udf_transpile_multi_row(self): # Every other transpile test uses a 1-row DataFrame; this one runs # the same arithmetic transpile on a multi-row input to catch any @@ -1026,9 +1424,13 @@ def wrapper(x): self.assertEqual(result[0], 60, "must run the wrapper, not the wrapped source") def test_udf_transpile_falls_back_for_none_in_boolop(self): - # Python's `None and x` short-circuits to None; Spark's three-valued - # `null AND false` is false. A literal None operand must force a - # fallback rather than silently diverge. + # Python's `None and x` short-circuits to None; neither Spark's + # three-valued `null AND false` (false) nor the nested CASE WHEN the + # transpiler emits (a NULL condition is not-taken, so also false) + # reproduces that. A literal None operand -- and, more generally, any + # operand `_is_never_null_boolean` cannot prove non-NULL, such as a + # ternary with a None branch -- must force a fallback rather than + # silently diverge. import warnings as _warnings from pyspark.sql.types import StructField, StructType @@ -1036,12 +1438,22 @@ def test_udf_transpile_falls_back_for_none_in_boolop(self): def none_and(x): return None and (x > 0) + def nullable_ternary_and(x): + return (True if x > 0 else None) and (x > 0) + long_schema = StructType([StructField("a", LongType(), nullable=True)]) with self.sql_conf(_TRANSPILE_ON): - with _warnings.catch_warnings(record=True): - _warnings.simplefilter("always") - pudf = UserDefinedFunction(none_and, BooleanType()) - self.assertEqual([], pudf.transpiled, "literal None in and/or must not transpile") + for func in (none_and, nullable_ternary_and): + with self.subTest(func=func.__name__): + with _warnings.catch_warnings(record=True): + _warnings.simplefilter("always") + pudf = UserDefinedFunction(func, BooleanType()) + self.assertEqual( + [], + pudf.transpiled, + f"{func.__name__}: nullable and/or operand must not transpile", + ) + pudf = UserDefinedFunction(none_and, BooleanType()) df = self.spark.createDataFrame([Row(a=-5)], schema=long_schema) [result] = df.select(pudf("a")).collect() self.assertIsNone(result[0]) @@ -1954,6 +2366,23 @@ def test_udf_transpile_lowers_operators(self): none_eq = lambda x: None == x # noqa: E711,E731 col_lt = lambda a, b: (a < b) if a is not None and b is not None else None # noqa: E731 assigned = lambda v: v + 1 # noqa: E731 + bounded = lambda x: 0 <= x <= 100 # noqa: E731 + str_chain = lambda x: "a" < x < "z" # noqa: E731 + in_tuple = lambda x: x in (1, 2, 3) # noqa: E731 + not_in_tuple = lambda x: x not in (1, 2) # noqa: E731 + in_negative = lambda x: x in (-1, 2) # noqa: E731 + in_list = lambda x: x in [1, 2] # noqa: E731 + in_set = lambda x: x in {1, 2} # noqa: E731 + in_dup = lambda x: x in (1, 1) # noqa: E731 + in_str = lambda x: x in ("a", "b") # noqa: E731 + # `not (x in ...)` is deliberately spelled with an explicit `not` rather + # than `not in`: it parses to UnaryOp(Not) over Compare(In), which is a + # different lowering path than Compare(NotIn) and is covered by + # `not_in_tuple` above. Hence the E713 suppression. + not_in_chain = lambda x: not (x in (1, 2)) # noqa: E713,E731 + in_and = lambda x, y: (x in (1, 2)) and (y > 0) # noqa: E731 + chain_and = lambda a, b, c: (0 < a < 10) and (b == c) # noqa: E731 + chain_ternary = lambda x: (0 < x < 10) if x is not None else None # noqa: E731 def if_elif_else(x): if x is None: @@ -1963,6 +2392,12 @@ def if_elif_else(x): else: return 1 + def if_in(x): + if x in (1, 2): + return 10 + else: + return 20 + # (func, return_type, schema, rows, expected); arg columns come from the schema. cases = [ (modulo, L, "a long, b long", [(7, 3), (7, -3), (-7, 3), (-7, -3)], [1, -2, 2, -1]), @@ -1980,6 +2415,37 @@ def if_elif_else(x): (col_lt, B, "a long, b long", [(1, 2), (2, 1), (1, 1)], [True, False, False]), (if_elif_else, L, "a long", [(None,), (0,), (5,), (-3,)], [-1, 0, 1, 1]), (assigned, L, "a long", [(1,), (10,)], [2, 11]), + # Chained comparisons. + ( + bounded, + B, + "a long", + [(0,), (50,), (100,), (-1,), (101,)], + [True, True, True, False, False], + ), + (str_chain, B, "a string", [("m",), ("a",), ("z",)], [True, False, False]), + # A chain as a boolean operand and inside a ternary: the fold's + # result is a non-NULL boolean, so it is accepted by the `and`/`or` + # gate and unifies with a None ternary branch. + ( + chain_and, + B, + "a long, b long, c long", + [(5, 1, 1), (5, 1, 2), (50, 1, 1)], + [True, False, False], + ), + (chain_ternary, B, "a long", [(5,), (50,), (None,)], [True, False, None]), + # `in` / `not in` over container literals. + (in_tuple, B, "a long", [(1,), (3,), (4,), (None,)], [True, True, False, False]), + (not_in_tuple, B, "a long", [(1,), (3,), (None,)], [False, True, True]), + (in_negative, B, "a long", [(-1,), (2,), (1,)], [True, True, False]), + (in_list, B, "a long", [(1,), (3,)], [True, False]), + (in_set, B, "a long", [(2,), (3,)], [True, False]), + (in_dup, B, "a long", [(1,), (2,)], [True, False]), + (in_str, B, "a string", [("a",), ("c",), (None,)], [True, False, False]), + (not_in_chain, B, "a long", [(1,), (3,), (None,)], [False, True, True]), + (in_and, B, "a long, b long", [(1, 5), (1, -5), (9, 5)], [True, False, False]), + (if_in, L, "a long", [(1,), (2,), (9,), (None,)], [10, 10, 20, 20]), ( nested, B, @@ -2033,6 +2499,77 @@ def test_udf_transpile_plan_elision(self): self.assertEqual(chained.first()[0], 7.0) # ((10 + 3) + 1) / 2 self.assertEqual(2, self._eval_python_count(chained)) + # A chained comparison and an `in` are both fully lowered, so the + # EvalPython node disappears for them too. + bounded = lambda x: 0 < x < 10 # noqa: E731 + member = lambda x: x in (1, 2, 3) # noqa: E731 + for func, expected in [(bounded, [1, 7]), (member, [1])]: + with self.subTest(func=str(func)): + f = UserDefinedFunction(func, BooleanType()) + self.assertTrue(f.transpiled) + fdf = self.spark.createDataFrame([(7,), (11,), (1,)], "a long").filter(f("a")) + self.assertEqual(sorted(r[0] for r in fdf.collect()), expected) + self.assertEqual(0, self._eval_python_count(fdf)) + + def test_udf_transpile_nondeterministic_argument_falls_back(self): + # EXPECTED TO FAIL until the companion analyzer fix lands. This test is + # deliberately red: it pins a real wrong-results bug so the gap stays + # visible rather than being silently tolerated. Do not skip or xfail it. + # + # A non-deterministic ARGUMENT expression must suppress transpilation at + # the call site, even though the UDF itself transpiled fine. (Distinct + # from test_udf_transpile_skips_nondeterministic, which covers a UDF + # marked non-deterministic.) + # + # `resolveUDFParams` substitutes the argument expression at every + # `_udf_param_N` occurrence, and the lowerings reference an operand two + # to three times -- four for a two-link chain. For a deterministic + # argument that is only plan bloat, but a non-deterministic one becomes + # several independent expressions with independent internal state. Since + # `and` / `or` and chained comparisons evaluate later copies only on the + # rows where earlier links held, the copies advance at different rates + # and stop describing the same value: `50 < x < 60` over `rand(7) * 100` + # returns ~30% true instead of the correct ~10%. + # + # The bug pre-dates chained comparisons (`and` / `or` duplicate operands + # too) and cannot be fixed in Python: transpilation runs at UDF + # construction, before any call site exists. It also cannot be fixed in + # the JVM builder, where an argument is still unresolved (`rand(7)` + # reports deterministic=true as an UnresolvedFunction). The fix belongs + # in ResolveTranspiledPythonUDFOptions, which runs post-resolution. + # + # Assert on the plan rather than on observed proportions: a statistical + # assertion would be flaky, whereas leaving EvalPython in place means the + # interpreted UDF runs, which evaluates each argument exactly once per + # row and so is correct by construction. + from pyspark.sql.functions import col, rand + + band = lambda x: 50 < x < 60 # noqa: E731 + band_or = lambda x: x < 50 or x > 60 # noqa: E731 + with self.sql_conf(_TRANSPILE_ON): + for func in (band, band_or): + pudf = UserDefinedFunction(func, BooleanType()) + self.assertTrue(pudf.transpiled, f"{func} should transpile on its own") + base = self.spark.range(0, 8).select((col("id") * 10).alias("a")) + + # Deterministic argument -> transpiled and elided. + det = base.select(pudf("a")) + self.assertEqual( + 0, + self._eval_python_count(det), + f"{func}: a deterministic argument should still be elided", + ) + + # Non-deterministic argument -> must fall back. + nondet = base.select(pudf(rand(7) * 100)) + self.assertEqual( + 1, + self._eval_python_count(nondet), + f"{func}: a non-deterministic argument must NOT be transpiled", + ) + # And it still produces results. + self.assertEqual(8, len(nondet.collect())) + def test_udf_transpile_config_toggle_no_stale_nodes(self): # Built with the flags on, executed with them off -> clean fallback to # interpreted Python (the optimizer drops the transpiled node), no error. @@ -2430,6 +2967,26 @@ def test_udf_transpile_known_value_divergences(self): self._vals(eq_strlit, BooleanType(), "a long", [(5,), (3,)], require_lowered=False), [False, False], ) + # Chained comparisons inherit the same NaN divergence, one link at a + # time -- no new mechanism, but easier to hit. `a < b < NaN` reaches the + # second link whenever `a < b`, and Spark orders NaN above every value + # (Python's `b < nan` is False); `a == b == NaN` likewise, since Spark + # treats `NaN = NaN` as true where Python's `nan == nan` is False. + nan_chain = lambda a, b, c: a < b < c # noqa: E731 + nan_eq_chain = lambda a, b, c: a == b == c # noqa: E731 + nan = float("nan") + self.assertEqual( + self._vals(nan_chain, BooleanType(), "a double, b double, c double", [(1.0, 2.0, nan)]), + [True], + "Spark orders NaN highest; Python would return False", + ) + self.assertEqual( + self._vals( + nan_eq_chain, BooleanType(), "a double, b double, c double", [(nan, nan, nan)] + ), + [True], + "Spark treats NaN = NaN as true; Python would return False", + ) def test_udf_transpile_overflow_and_modulo_zero_raise(self): # Transpiled arithmetic that raises at runtime: `*` overflow raises under diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 21289a15a7e15..89219e404a623 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -90,6 +90,12 @@ expression. A draw is never inlined, because those rules check determinism. An argument the body never reads is not computed at all. + +Constructs that short-circuit in Python -- ``and`` / ``or`` and chained +comparisons such as ``0 <= x <= 100`` -- lower to nested ``CASE WHEN``s +rather than to ``&`` / ``|``, so that an operand Python never evaluates +cannot raise; see ``_lazy_and``. Their operands must also be provably +non-NULL booleans; see ``_is_never_null_boolean``. """ import ast @@ -132,6 +138,16 @@ from pyspark.sql._typing import DataTypeOrString +# Caps on how large a comparison we will lower, to bound plan growth: longer +# ones fall back to interpreted Python. Both are generous next to what shows up +# in real code (``0 <= x <= 100``, an allowlist of a few dozen values), and both +# matter because the lowered subtree is duplicated -- once per adjacent pair for +# a chain operand, once per input-category variant for every option -- and +# CheckAnalysis walks each copy. +_MAX_CHAIN_OPS = 4 +_MAX_IN_ELEMENTS = 64 + + class AbstractTranspiler(object): """Base class for transpilers. All experimental.""" @@ -198,9 +214,17 @@ def _is_definitely_boolean(node: ast.AST) -> bool: match node: case ast.Constant(value=v): return v is None or isinstance(v, bool) - case ast.Compare(left=left, comparators=comparators): - # All comparison operators of simple types bool - return all(_is_definitely_basic_type(v) for v in comparators + [left]) + case ast.Compare(left=left, ops=ops, comparators=comparators): + # Every Python comparison produces a bool, chained ones and + # ``in`` / ``not in`` included. Don't inspect an ``in`` right + # operand: it is a CONTAINER, which ``_is_definitely_basic_type`` + # rejects, yet the position still produces a bool. Being + # over-permissive is safe -- ``_convert_chunk`` fails closed on + # anything it cannot lower. + return _is_definitely_basic_type(left) and all( + isinstance(op, (ast.In, ast.NotIn)) or _is_definitely_basic_type(c) + for op, c in zip(ops, comparators) + ) case ast.BoolOp(values=values): return all(_is_definitely_boolean(v) for v in values) case ast.UnaryOp(op=ast.Not()): @@ -233,6 +257,148 @@ def _truthiness_col(cat: Optional[str], c: Column) -> Optional[Column]: return None +def _is_never_null_boolean(node: ast.AST) -> bool: + """Return True when ``node`` lowers to a boolean column that either raises + or yields True/False -- never NULL. + + Stricter than ``_is_definitely_boolean``, which also accepts operands that + may evaluate to NULL (a literal ``None``, a ternary with a ``None`` branch). + ``_lazy_and`` / ``_lazy_or`` need this: a CASE WHEN treats a NULL condition + as not-taken, so a NULL operand would silently become ``False`` where Python + returns ``None``. Bodies that fail this go back to interpreted Python. + ``not`` and if/ternary *tests* keep using ``_is_definitely_boolean``, since + they wrap the operand in ``coalesce`` and so handle NULL explicitly. + """ + match node: + case ast.Constant(value=v): + # A literal None is "boolean" for coalesce purposes but is NULL. + return isinstance(v, bool) + case ast.Compare(): + # Every comparison lowering is total (see ``_lower_compare_pair``), + # so only the operand shapes still need checking. + return _is_definitely_boolean(node) + case ast.UnaryOp(op=ast.Not()): + # Lowered as ``coalesce(~operand, lit(True))``, which is total. + return True + case ast.BoolOp(values=values): + return all(_is_never_null_boolean(v) for v in values) + case ast.IfExp(body=body, orelse=orelse): + return _is_never_null_boolean(body) and _is_never_null_boolean(orelse) + case _: + return False + + +def _lazy_and(cols: List[Column]) -> Column: + """Fold ``cols`` into a boolean column with Python's ``and`` semantics, + evaluating each element only when every earlier one was True. + + Nests CASE WHENs instead of folding with ``&``. The reason is NOT that + ``And`` fails to short-circuit -- both ``And.eval`` and ``And.doGenCode`` + skip the right operand when the left is false. It is the optimizer: + ``splitConjunctivePredicates`` decomposes ``And`` and nothing else, and + ``PushPredicateThroughJoin`` gates pushdown on + ``cond.deterministic && !cond.throwable``. ``Expression.throwable`` defaults + to ``children.exists(_.throwable)`` and ``RaiseError`` does not override it + (its children are literals), so Catalyst believes the NULL guard that + ``_lower_value_compare`` emits cannot throw, and will push one conjunct to + the far side of a join where it sees rows another conjunct excluded. + + A CaseWhen resists that: it is atomic to predicate splitting, evaluates + branches in order (``CaseWhen.eval`` stops at the first TRUE condition), + stays non-foldable (``ConditionalExpression.foldable`` needs every child + foldable and ``RaiseError.foldable`` is false), and hides its lazy tail from + subexpression elimination -- ``EquivalentExpressions`` reaches the tail only + via ``branchGroups``, which intersects the branch values, and one of ours is + a ``Literal`` that ``updateExprTree`` skips, leaving that intersection + empty. + + SCOPE, precisely: this stops predicate splitting from tearing ONE lowered + predicate into independently-pushable conjuncts. It does NOT stop the whole + predicate being pushed as a single conjunct, since ``throwable`` is still + false for it -- ``left.join(right, "k").filter(f(col("c")))`` with ``c`` on + the right side still pushes below the join, so the guard fires on rows that + never join and raises where interpreted Python never saw them. That + divergence pre-dates chained comparisons (a single comparison already lowers + to one guarded CASE WHEN); only overriding ``throwable`` on ``RaiseError`` + in Catalyst closes it, which is a separate change because it shifts pushdown + behavior for every user of ``raise_error()`` / ``assert_true()``. + + Callers must have proven every element non-NULL (``_is_never_null_boolean``), + which is what lets this skip ``coalesce``. + + COST, and it is not free. When every operand column is non-nullable the + ``&`` fold used to simplify all the way down: ``NullPropagation`` folds + ``IsNull(e)`` to false for a non-nullable ``e``, so the guard disappears and + what remains is a plain ``And`` that ``splitConjunctivePredicates`` can tear + into conjuncts for datasource pushdown, partition pruning and + ``InferFiltersFromConstraints``. ``CaseWhen(Seq((c1, c2)), Some(false))`` has + no matching rule in ``SimplifyConditionals`` (which only collapses the + ``TrueLiteral`` / ``FalseLiteral`` branch shapes), so it survives as an + opaque predicate and none of that applies. Measured on a non-nullable + column: ``x > 20 and x < 80`` optimized to ``((id > 20) AND (id < 80))`` + under the old fold and stays ``CASE WHEN (id > 20) THEN (id < 80) ELSE + false END`` under this one. + + The trade is deliberate but narrow. The pushdown loss only bites when the + guard would have folded away -- i.e. non-nullable columns -- and that is + exactly the case where the ``And`` form was already SAFE, because with no + ``raise_error`` left there is nothing for the optimizer to relocate. When + the guard does survive (nullable columns) it blocks datasource pushdown in + either form, so nothing is lost. Overriding ``throwable`` on ``RaiseError`` + would let this go back to ``&`` / ``|`` and get both properties at once; + that is the right end state and is tracked separately. + """ + result = cols[-1] + for c in reversed(cols[:-1]): + result = when(c, result).otherwise(lit(False)) + return result + + +def _lazy_or(cols: List[Column]) -> Column: + """Fold ``cols`` with Python's ``or`` semantics, evaluating each element only + when every earlier one was False. The mirror of ``_lazy_and``; see there for + why this nests CASE WHENs rather than folding with ``|``. + """ + result = cols[-1] + for c in reversed(cols[:-1]): + result = when(c, lit(True)).otherwise(result) + return result + + +def _const_container_elements(node: ast.AST) -> Optional[List[ast.expr]]: + """Element nodes of a constant container literal, or ``None`` when ``node`` + is not one. + + Recognises ``ast.Tuple`` / ``ast.List`` / ``ast.Set`` displays whose every + element is an ``ast.Constant`` or a unary ``+`` / ``-`` on a non-bool numeric + ``ast.Constant`` (``x in (-1, 2)`` parses ``-1`` as a ``UnaryOp``). Anything + else -- a ``Name``, a ``Call``, a starred or nested element, a ``dict`` / + ``str`` / ``bytes`` right operand -- returns ``None`` so the caller falls + back; see ``CatalystTranspiler._lower_in`` for why element values must be + known at transpile time. + """ + if not isinstance(node, (ast.Tuple, ast.List, ast.Set)): + return None + for e in node.elts: + if isinstance(e, ast.Constant): + continue + if ( + isinstance(e, ast.UnaryOp) + and isinstance(e.op, (ast.USub, ast.UAdd)) + and isinstance(e.operand, ast.Constant) + and isinstance(e.operand.value, (int, float)) + and not isinstance(e.operand.value, bool) + ): + continue + return None + return list(node.elts) + + +def _is_none_constant(node: ast.AST) -> bool: + """Return True for the literal ``None``.""" + return isinstance(node, ast.Constant) and node.value is None + + class CatalystTranspiler(AbstractTranspiler): """Transpiler that attempts to convert a Python UDF into native Spark SQL expressions.""" @@ -288,7 +454,7 @@ def _safe_category(self, params: List[str], node: Optional[ast.AST]) -> Optional if body_c is not None and else_c is not None and body_c != else_c: return None return body_c if body_c is not None else else_c - if isinstance(node, ast.Constant) and node.value is None: + if _is_none_constant(node): return None # Comparisons / ``not`` / boolean ops produce a boolean column; classify # them as "bool" (``_category``'s catch-all would mislabel them numeric). @@ -451,6 +617,162 @@ def _lower_value_compare( ) return when(null_guard, raise_error(err)).otherwise(op(left_col, right_col)) + def _lower_compare_pair( + self, + params: List[str], + op: ast.cmpop, + left_node: ast.AST, + right_node: ast.AST, + ) -> Column: + """Lower ONE adjacent comparison of a (possibly chained) ``ast.Compare``. + + TOTALITY INVARIANT, which callers depend on: when the returned column + returns a value at all, that value is never NULL. + + * ``is`` / ``is not`` emit ``isNull`` / ``isNotNull``, non-nullable by + construction. + * ``_lower_eq`` answers every NULL combination with a boolean literal, + reaching its ``otherwise`` only when both operands are non-NULL. + * ``_lower_value_compare`` raises rather than returning NULL. + + Equal operand categories plus the JVM-side category contract in + ``ResolveTranspiledPythonUDFOptions`` (non-decimal ``NumericType`` for + "numeric", UTF8_BINARY ``StringType`` for "string") leave numeric + widening as the only coercion Catalyst inserts, and under ANSI a failing + cast raises instead of yielding NULL. + + That invariant is what lets ``_lazy_and`` and ``_lower_in``'s negation + use plain ``when`` / ``~`` with no ``coalesce``. + """ + match op: + case ast.Is() | ast.IsNot(): + # Only lower `x is None` / `None is x` (and their + # `is not` variants) to isNull/isNotNull. For any + # other comparator (e.g. `x is 0`, `x is y`) Python + # performs an object-identity check that has no SQL + # equivalent, so we must fall back to interpreted + # Python rather than silently emitting a null check. + is_none_left = _is_none_constant(left_node) + is_none_right = _is_none_constant(right_node) + if not (is_none_left or is_none_right): + raise UnsupportedOperationException( + "`is`/`is not` is only supported when one " + "operand is the literal None; other identity " + "checks (e.g. `x is 0`, `x is y`) cannot be " + "lowered to SQL and the UDF falls back to " + "interpreted Python" + ) + subject_node = right_node if is_none_left else left_node + subject_col = self._convert_chunk(params, subject_node) + if isinstance(op, ast.Is): + return subject_col.isNull() + else: + return subject_col.isNotNull() + case ast.Eq(): + return self._lower_eq(params, left_node, right_node, equal=True) + case ast.NotEq(): + return self._lower_eq(params, left_node, right_node, equal=False) + case ast.Lt(): + return self._lower_value_compare( + params, left_node, right_node, lambda l, r: l < r, "<" + ) + case ast.LtE(): + return self._lower_value_compare( + params, left_node, right_node, lambda l, r: l <= r, "<=" + ) + case ast.Gt(): + return self._lower_value_compare( + params, left_node, right_node, lambda l, r: l > r, ">" + ) + case ast.GtE(): + return self._lower_value_compare( + params, left_node, right_node, lambda l, r: l >= r, ">=" + ) + case _: + raise UnsupportedOperationException( + f"comparison operator {type(op).__name__} is not supported by the transpiler" + ) + + def _lower_in( + self, + params: List[str], + subject_node: ast.AST, + container_node: ast.AST, + negate: bool, + ) -> Column: + """Lower ``x in (c1, ..., cn)`` / ``x not in (...)`` over a literal + tuple / list / set of constants. + + Python's ``in`` over a container is EQUALITY-based and never raises on a + ``None`` probe: ``None in (1, 2)`` is False and ``None in (1, None)`` is + True. Spark's ``IN`` returns NULL when the probe is NULL *or* when any + list element is NULL, so we answer the NULL-probe case directly with a + literal and DROP the NULL elements from the ``isin`` list. Leaving one in + would make the whole expression NULL for a non-matching probe, breaking + the totality that ``not in``'s ``~`` relies on. + + The empty container is spelled out as ``lit(False)`` rather than + delegating to ``IN ()``: ``In.eval`` returns false for an empty list only + while ``spark.sql.legacy.nullInEmptyListBehavior`` is off; with it on, + ``OptimizeIn`` rewrites ``IN ()`` to ``If(IsNotNull(v), false, null)`` + and a NULL probe would yield NULL. + + Every element must share the SUBJECT's category. That guard is + load-bearing: ``InTypeCoercion`` finds no wider common type for e.g. + bigint-vs-string under ANSI, so ``In.checkInputDataTypes`` fails with + DATA_DIFF_TYPES -- and because the option rides along as a child of + ``TranspiledPythonUDF``, that fails the WHOLE query at CheckAnalysis + instead of falling back. Python's cross-type ``in`` is simply False, so + coercing would buy nothing anyway. + """ + elements = _const_container_elements(container_node) + if elements is None: + raise UnsupportedOperationException( + "`in` / `not in` is only supported over a literal tuple, list " + "or set of constants (e.g. `x in (1, 2)`): `in` on a str or " + "bytes is substring/subsequence containment rather than " + "equality (`'b' in 'abc'` is True while `'b' == 'abc'` is " + "False, and `b'ab'` iterates ints so `97 in b'ab'` is True), " + "dict membership tests keys, and a container holding runtime " + "values would need runtime NULL analysis; the transpiler falls " + "back to interpreted Python" + ) + if len(elements) > _MAX_IN_ELEMENTS: + raise UnsupportedOperationException( + f"`in` / `not in` container has {len(elements)} elements, more " + f"than the {_MAX_IN_ELEMENTS} the transpiler will lower; the " + "UDF falls back to interpreted Python to bound plan growth " + "across the input-type option matrix" + ) + # Raises for a literal-None or complex subject, dropping this variant. + subject_cat = self._category(params, subject_node) + has_none = False + value_nodes: List[ast.expr] = [] + for element in elements: + if _is_none_constant(element): + has_none = True + continue + # Classify a unary +/- element by its operand rather than leaning on + # `_category`'s catch-all, which happens to answer "numeric" for an + # `ast.UnaryOp` today. `_const_container_elements` has already + # restricted the operand to a non-bool numeric constant. + classify = element.operand if isinstance(element, ast.UnaryOp) else element + element_cat = self._category(params, classify) + if element_cat != subject_cat: + raise UnsupportedOperationException( + f"`in` / `not in` element category ({element_cat}) does not " + f"match the subject's ({subject_cat}); Python compares " + "across types as unequal while Spark would coerce or fail " + "analysis, so the transpiler falls back to interpreted " + "Python" + ) + value_nodes.append(element) + subject_col = self._convert_chunk(params, subject_node) + element_cols = [self._convert_chunk(params, e) for e in value_nodes] + membership = subject_col.isin(*element_cols) if element_cols else lit(False) + result = when(subject_col.isNull(), lit(has_none)).otherwise(membership) + return result.__invert__() if negate else result + def _category(self, params: List[str], node: ast.AST) -> str: """Infer ``"numeric"`` or ``"string"`` for ``node`` under the current ``self._param_categories`` assumption (set per input-type variant). @@ -530,7 +852,7 @@ def _category_uncached(self, params: List[str], node: ast.AST) -> str: # unifies with any type in the lowered CASE WHEN); mismatched or # all-None branches raise so the variant is dropped. def branch_category(b: ast.AST) -> Optional[str]: - if isinstance(b, ast.Constant) and b.value is None: + if _is_none_constant(b): return None return self._category(params, b) @@ -547,9 +869,23 @@ def branch_category(b: ast.AST) -> Optional[str]: "ternary with all-None branches has no usable column category" ) return result_cat + case ast.Compare(): + # EVERY Python comparison evaluates to a bool -- `in` / `not in` + # and chained forms included -- whatever its operands look like. + # `_is_definitely_boolean` can't answer this because its Compare + # arm also requires basic-type operands, so it says False for + # e.g. `(x if c else y) in (1, 2)` or `(a < b) == True`. Those + # then fell through to the numeric catch-all below, which let the + # return-type match in `_transpile_from_ast` accept a numeric + # declared return type for a boolean body: the lowered option + # returned `cast( as bigint)` (1/0) where the + # interpreted SQL_BATCHED_UDF path returns NULL, a silent + # divergence. Classify on the node kind instead, which is a fact + # about Python rather than an operand analysis. + return "bool" case _ if _is_definitely_boolean(node): - # Comparisons, `not`, and boolean ops produce a boolean column. - # Labeling them "numeric" (the old catch-all) let booleans into + # `not` and boolean ops produce a boolean column. Labeling them + # "numeric" (the old catch-all) let booleans into # arithmetic/equality lowerings where ANSI analysis fails (e.g. # `(x > 0) + 1`, valid Python) instead of falling back. return "bool" @@ -607,48 +943,36 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: case ast.BoolOp(op=op, values=values): # Python `and` / `or` short-circuit and return one of the # operands rather than a strict boolean. For the booleans - # produced by Compare / UnaryOp(Not) / nested BoolOps this - # maps cleanly onto Spark Column `&` / `|`. For - # non-boolean operands (including bare parameter names whose - # runtime type is unknown) the right semantics would require - # Python's truthiness rules (0 / "" / None / [] all - # falsy), which we can't faithfully reproduce without the - # input column types -- Spark's `&` / `|` would silently - # do bitwise instead. Require all operands to be statically - # known boolean so the caller falls back to interpreted - # Python rather than producing a plan whose results diverge. - if not all(_is_definitely_boolean(v) for v in values): - raise UnsupportedOperationException( - "`and` / `or` operand type is not statically known " - "to be boolean; Spark's `&` / `|` are bitwise, not " - "Python truthiness, so the transpiler refuses to " - "lower this and the UDF falls back to interpreted " - "Python" - ) - # A literal None operand short-circuits differently: Python's - # `None and (x > 0)` returns None regardless of x, but Spark's - # three-valued `null AND false` is false (and `null OR true` is - # true), so the lowered form diverges. `_is_definitely_boolean` - # accepts None for `not`/if-test contexts where coalesce handles - # it; here it must force a fallback instead. - if any(isinstance(v, ast.Constant) and v.value is None for v in values): + # produced by Compare / UnaryOp(Not) / nested BoolOps that maps + # cleanly onto a nested CASE WHEN. For non-boolean operands + # (including bare parameter names whose runtime type is unknown) + # the right semantics would need Python's truthiness rules + # (0 / "" / None / [] all falsy), which we can't reproduce + # without the input column types -- Spark's `&` / `|` would do + # bitwise instead. Require statically known booleans so the + # caller falls back rather than emitting a divergent plan. + # + # The gate is `_is_never_null_boolean`, not + # `_is_definitely_boolean`, because `_lazy_and` / `_lazy_or` need + # non-NULL operands: a NULL CASE WHEN condition is not-taken, so + # `and` would yield False where Python returns None. That + # refuses a literal None operand (`None and (x > 0)` is None in + # Python) and a ternary with a None branch -- both of which + # diverge under Spark's three-valued `&` / `|` as well. + if not all(_is_never_null_boolean(v) for v in values): raise UnsupportedOperationException( - "literal None operand in `and` / `or` cannot be lowered: " - "Spark's three-valued logic diverges from Python's " - "short-circuit-return-operand semantics, so the UDF " + "`and` / `or` operand is not statically known to be a " + "non-NULL boolean; Python's short-circuit returns the " + "operand itself (so `None and x` is None) and Spark's " + "`&` / `|` are bitwise rather than Python truthiness, " + "so the transpiler refuses to lower this and the UDF " "falls back to interpreted Python" ) cols = [self._convert_chunk(params, v) for v in values] if isinstance(op, ast.And): - result = cols[0] - for c in cols[1:]: - result = result & c - return result + return _lazy_and(cols) if isinstance(op, ast.Or): - result = cols[0] - for c in cols[1:]: - result = result | c - return result + return _lazy_or(cols) raise UnsupportedOperationException(f"BoolOp operator {op} is not supported") case ast.IfExp(test=test, body=body_expr, orelse=orelse_expr): # Ternary `body if test else orelse` -- shares the @@ -673,60 +997,79 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: orelse[0] if orelse else None, ) case ast.Compare(left, ops, comps): - if len(ops) != 1 or len(comps) != 1: - raise UnsupportedOperationException( - "chained comparisons (e.g. `a < b < c`) are not supported by the transpiler" - ) - comp = comps[0] - match ops[0]: - case ast.Is() | ast.IsNot(): - # Only lower `x is None` / `None is x` (and their - # `is not` variants) to isNull/isNotNull. For any - # other comparator (e.g. `x is 0`, `x is y`) Python - # performs an object-identity check that has no SQL - # equivalent, so we must fall back to interpreted - # Python rather than silently emitting a null check. - is_none_left = isinstance(left, ast.Constant) and left.value is None - is_none_right = isinstance(comp, ast.Constant) and comp.value is None - if not (is_none_left or is_none_right): - raise UnsupportedOperationException( - "`is`/`is not` is only supported when one " - "operand is the literal None; other identity " - "checks (e.g. `x is 0`, `x is y`) cannot be " - "lowered to SQL and the UDF falls back to " - "interpreted Python" - ) - subject_node = comp if is_none_left else left - subject_col = self._convert_chunk(params, subject_node) - if isinstance(ops[0], ast.Is): - return subject_col.isNull() - else: - return subject_col.isNotNull() - case ast.Eq(): - return self._lower_eq(params, left, comp, equal=True) - case ast.NotEq(): - return self._lower_eq(params, left, comp, equal=False) - case ast.Lt(): - return self._lower_value_compare( - params, left, comp, lambda l, r: l < r, "<" - ) - case ast.LtE(): - return self._lower_value_compare( - params, left, comp, lambda l, r: l <= r, "<=" - ) - case ast.Gt(): - return self._lower_value_compare( - params, left, comp, lambda l, r: l > r, ">" - ) - case ast.GtE(): - return self._lower_value_compare( - params, left, comp, lambda l, r: l >= r, ">=" - ) - case _: + # `in` / `not in` are container membership, not a chainable + # value comparison, so only the un-chained single-operator form + # is lowered. `a < b in (1, 2)` is legal Python; refuse it + # explicitly rather than relying on `_convert_chunk` + # incidentally rejecting the ast.Tuple operand. + if any(isinstance(o, (ast.In, ast.NotIn)) for o in ops): + if len(ops) != 1: raise UnsupportedOperationException( - f"comparison operator {type(ops[0]).__name__} " - "is not supported by the transpiler" + "`in` / `not in` inside a chained comparison (e.g. " + "`a < b in (1, 2)`) is not supported by the " + "transpiler; falling back to interpreted Python" ) + return self._lower_in( + params, left, comps[0], negate=isinstance(ops[0], ast.NotIn) + ) + if len(ops) > _MAX_CHAIN_OPS: + raise UnsupportedOperationException( + f"chained comparison with {len(ops)} operators exceeds " + f"the {_MAX_CHAIN_OPS} the transpiler will lower; the " + "UDF falls back to interpreted Python" + ) + # Python evaluates `a OP1 b OP2 c` as `(a OP1 b) and (b OP2 c)` + # with `b` evaluated once and with SHORT-CIRCUIT: when + # `a OP1 b` is falsy, `b OP2 c` is never evaluated. That is not + # cosmetic here, because `_lower_value_compare` emits a + # `raise_error` for NULL operands to mirror Python's TypeError: + # `f(1, 0, None)` for `lambda a, b, c: a < b < c` must return + # False, not raise. `_lazy_and` explains why reproducing that + # needs nested CASE WHENs rather than `&`, and why every link + # being total (see `_lower_compare_pair`) removes the need for + # `coalesce`. + # + # Python's "evaluate `b` once" rule is about side effects. The + # *Python-level* operand grammar we accept (Constant / Name / + # BinOp / UnaryOp / IfExp / Compare / BoolOp) is deterministic + # and side-effect free -- ANSI raises aside, and those are + # idempotent -- so re-lowering an interior operand once per + # adjacent pair is safe. Not so for the CALL-SITE arguments: a + # parameter lowers to a `_udf_param_N` placeholder and the JVM's + # `resolveUDFParams` substitutes the caller's expression at EVERY + # occurrence, four of them for a two-link ordering chain. A + # NON-DETERMINISTIC argument therefore becomes that many + # independent expressions with independent state, and because the + # fold above evaluates later links only on the rows where earlier + # ones held, those copies desynchronize and stop describing the + # same value: `50 < x < 60` over `rand(7) * 100` yields ~30% true + # where Python yields ~10%. + # + # That is an OPEN BUG, not something this code can fix: the + # transpiler runs at UDF construction, long before a call site + # exists. The fix belongs in the analyzer, which is the first + # place a resolved argument can be tested for determinism, and is + # being handled as a companion change; until it lands the + # divergence is pinned by a deliberately failing test in + # test_udf_transpile_unit.py. + # + # Be precise about the blast radius, because it is easy to + # overstate as merely pre-existing. `and` / `or` were already + # wrong this way, since their operands are separate subtrees. A + # SINGLE comparison is NOT: its duplicate operand appears only + # inside an `IS NULL` test, whose outcome does not depend on the + # value drawn, so `50 < x` over `rand()` samples correctly. So + # lowering chained comparisons -- which previously fell back to + # interpreted Python, and were therefore correct -- newly exposes + # them to this bug. That argues for landing the analyzer guard + # before or with this change. + nodes = [left, *comps] + return _lazy_and( + [ + self._lower_compare_pair(params, op, left_node, right_node) + for op, left_node, right_node in zip(ops, nodes, nodes[1:]) + ] + ) case ast.BinOp(left=left, op=op, right=right): # Operator selection is driven by the operand *categories* under # the current input-type variant (see ``_category``): Python's diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala index f9295eb17b00c..dd5381ac70ab5 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.analysis import org.apache.spark.api.python.PythonEvalType import org.apache.spark.sql.catalyst.dsl.expressions._ -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Concat, Expression, Literal, PythonUDF, TranspiledPythonUDF} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Concat, Expression, GreaterThan, Literal, Multiply, PythonUDF, Rand, TranspiledPythonUDF} import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project} import org.apache.spark.sql.types.LongType @@ -92,4 +92,40 @@ class ResolveTranspiledPythonUDFOptionsSuite extends PlanTest { val pruned = prune(node, LocalRelation(a)) assert(pruned.transpiledOptions == List(onlyOpt)) } + + test("empties the options when an argument is non-deterministic") { + // EXPECTED TO FAIL until the companion fix to this rule lands. Deliberately + // red: it pins a real wrong-results bug rather than tolerating it silently. + // Do not ignore or cancel it. + // + // The builder substitutes the argument expression at EVERY `_udf_param_N` + // occurrence, so a non-deterministic argument becomes several independent + // expressions with independent state. Since `and` / `or` and chained + // comparisons evaluate later copies only on rows where earlier links held, + // the copies desynchronize and stop describing the same value. The rule must + // drop every option so ConvertToCatalyst falls back to the Python UDF, which + // evaluates each argument exactly once per row. + // + // This must be checked here rather than in the builder: at call-construction + // time `rand(7)` is still an `UnresolvedFunction` whose `deterministic` + // derives from its `Literal` seed and so reports true. + val a = $"a".double + val arg = Multiply(Rand(Literal(7L)), Literal(100.0)) + assert(!arg.deterministic, "the fixture argument must be non-deterministic") + val opt = GreaterThan(arg, Literal(50.0)) + val node = TranspiledPythonUDF("udf", pyUDF(Seq(arg)), List(opt), List(List("numeric"))) + val pruned = prune(node, LocalRelation(a)) + assert(pruned.transpiledOptions.isEmpty) + assert(pruned.optionInputCategories.isEmpty) + } + + test("keeps options when the argument is deterministic (control for the check above)") { + // Same shape as the non-deterministic case, but with a plain column argument, so + // the non-determinism guard must not fire and normal category matching applies. + val a = $"a".double + val opt = GreaterThan(a, Literal(50.0)) + val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)), List(opt), List(List("numeric"))) + val pruned = prune(node, LocalRelation(a)) + assert(pruned.transpiledOptions == List(opt)) + } } From 2e0482791fa7f9dd0db965f1d55d496b7f66cc8c Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Tue, 11 Aug 2026 21:43:16 +0000 Subject: [PATCH 02/18] [SPARK-58456][PYTHON][SQL][MINOR] Address review findings in the comparison lowering Review follow-up. No behavior change to the lowering itself. - Make the pushdown regression test assert something. It never checked that `raise_error` was in the plan at all, so it passed vacuously once the guard folded away, and its "the guard must stay ABOVE the join" comment did not match the plan: the guard lands inside the join CONDITION, on the Join's own line, so the `i > join_line` scan skipped it. Assert the guard is present, match the `Join` node at line start rather than the substring (which also occurs inside a printed join condition), require exactly one Join node, and pin the guard to the condition. Being in the condition is safe -- the predicate stays whole, so `and` still short-circuits -- and the comment now says that instead of claiming otherwise. - Restore the arity check `_seed_tuple_examples` lost when a two-name tuple unpack became `dict(zip(keys, values))`. `zip` stops at the shorter argument, so a mismatched seed tuple silently registered an `@example` for a different input and quietly shrank the seeded NULL-combination coverage. - Fold `_is_never_null_boolean` and `_is_definitely_boolean` onto one traversal parameterized by whether a NULL operand is acceptable. They differed only in the `ast.Constant` arm, so every new lowering had to be mirrored in both; missing one silently re-admits a NULL operand into `_lazy_and`. - Look the four ordering comparisons up in `_ORDERING_OPS` instead of four near-identical `case` arms with throwaway lambdas. - Label the parameterized comparison tests explicitly. `subTest(func=str(func))` and `f"{func}"` printed a lambda's memory address, so a CI failure named nothing: the real output read ` at 0xffff5b6d6b60`. Follows the `label` pattern the fallback test already used. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../tests/test_udf_transpile_hypothesis.py | 13 +- .../sql/tests/test_udf_transpile_unit.py | 120 +++++++++++------- python/pyspark/sql/transpile.py | 107 ++++++++-------- 3 files changed, 144 insertions(+), 96 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py index 38a3a22cb978b..726eb17fb862a 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -241,10 +241,21 @@ def wrapper(method): return wrapper def _seed_tuple_examples(tuples, keys=("x", "y", "z")): - """Stack one ``@example`` decorator per seed tuple, zipped onto ``keys``.""" + """Stack one ``@example`` decorator per seed tuple, zipped onto ``keys``. + + The arity check is load-bearing: ``zip`` stops at the shorter argument, + so a seed tuple whose length does not match ``keys`` would silently + register an ``@example`` for a different input than intended and quietly + shrink the seeded NULL-combination coverage. Fail loudly instead. + """ def wrapper(method): for values in reversed(tuples): + if len(values) != len(keys): + raise ValueError( + f"seed tuple {values!r} has {len(values)} value(s) but " + f"{len(keys)} key(s) {keys!r}" + ) method = example(**dict(zip(keys, values)))(method) return method diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index 69465973f260b..aef0e8552d8d3 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -24,6 +24,7 @@ property-based suite lives in ``test_udf_transpile_hypothesis.py``. """ +import re import unittest from pyspark.sql import Row @@ -581,8 +582,11 @@ def test_udf_transpile_boolop_short_circuits(self): # runtime symptom is scheduling-dependent: with the old `&` lowering the # pushed-down `raise_error` fires only on some runs (whether the failing # task gets far enough before the join yields no rows). Asserting on the - # optimized plan instead is deterministic -- the raise guard must stay - # ABOVE the join rather than being pushed beneath it. + # optimized plan instead is deterministic -- the guard must not be pushed + # BENEATH the join, onto one side, where it would be evaluated on rows + # the other operand excluded. It does end up inside the join CONDITION, + # which is safe for the same reason the lowering exists: the predicate + # stays whole, so the `and` still short-circuits. def both_positive(a, c): return a > 0 and c > 0 @@ -612,14 +616,27 @@ def either_positive(a, c): filtered = joined.filter(and_udf("a", "c")) plan = filtered._jdf.queryExecution().optimizedPlan().toString() - # An optimized plan prints root-first, depth-first, so anything - # pushed below the join appears on a LATER line than the Join node. lines = plan.split("\n") - join_line = next( - (i for i, line in enumerate(lines) if "Join" in line), - None, + # Assert the guard is actually there first: without this every check + # below passes vacuously the moment the guard folds away or the UDF + # stops transpiling. + self.assertTrue( + any("raise_error" in line for line in lines), + f"expected the NULL raise guard somewhere in the plan:\n{plan}", ) - self.assertIsNotNone(join_line, f"expected a Join in the plan:\n{plan}") + # An optimized plan prints root-first, depth-first, so a node pushed + # below the Join appears on a LATER line than the Join itself. Match + # the node name at the start of the line rather than the substring + # "Join", which also occurs inside a printed join condition. + join_lines = [i for i, line in enumerate(lines) if re.match(r"^[\s:+|-]*Join ", line)] + self.assertEqual(1, len(join_lines), f"expected exactly one Join node:\n{plan}") + join_line = join_lines[0] + # The hazard is the guard being pushed to ONE SIDE of the join, where + # it would see rows that never join. Landing in the join CONDITION is + # acceptable: the predicate stays whole, so `and` still + # short-circuits and the guard is only reached for candidate pairs. + # That is where it lands today -- hence `> join_line`, which excludes + # the Join's own line -- so also pin that it did not decompose. pushed = [ line for i, line in enumerate(lines) if i > join_line and "raise_error" in line ] @@ -629,6 +646,13 @@ def either_positive(a, c): "the NULL raise guard was pushed below the Join, so it can be " f"evaluated on rows the other operand excluded:\n{plan}", ) + self.assertIn( + "raise_error", + lines[join_line], + "expected the whole predicate to stay in the join condition; if " + "it moved above the Join that is also safe, but the lowering " + f"changed and this test should be re-read:\n{plan}", + ) # And it still returns Python's answer rather than raising. self.assertEqual([], filtered.select("k").collect()) @@ -680,14 +704,16 @@ def chained_eq(a, b, c): gt_gte = lambda a, b, c: a > b >= c # noqa: E731 lte_lte = lambda a, b, c: a <= b <= c # noqa: E731 four_ops = lambda a, b, c, d, e: a < b < c < d < e # noqa: E731 - for func, rows, expected in [ - (lt_eq, [(1, 2, 2), (1, 2, 3), (2, 1, 1)], [True, False, False]), - (eq_lt, [(1, 1, 2), (1, 1, 0), (1, 2, 3)], [True, False, False]), - (gt_gte, [(3, 2, 2), (3, 2, 5), (1, 2, 2)], [True, False, False]), - (lte_lte, [(1, 1, 2), (1, 2, 2), (2, 1, 3)], [True, True, False]), + for label, func, rows, expected in [ + ("a < b == c", lt_eq, [(1, 2, 2), (1, 2, 3), (2, 1, 1)], [True, False, False]), + ("a == b < c", eq_lt, [(1, 1, 2), (1, 1, 0), (1, 2, 3)], [True, False, False]), + ("a > b >= c", gt_gte, [(3, 2, 2), (3, 2, 5), (1, 2, 2)], [True, False, False]), + ("a <= b <= c", lte_lte, [(1, 1, 2), (1, 2, 2), (2, 1, 3)], [True, True, False]), ]: - with self.subTest(func=str(func)): - self.assertEqual(self._vals(func, B, "a long, b long, c long", rows), expected) + with self.subTest(func=label): + self.assertEqual( + self._vals(func, B, "a long, b long, c long", rows), expected, label + ) self.assertEqual( self._vals( four_ops, @@ -712,15 +738,15 @@ def test_udf_transpile_in_lowered(self): not_in_none = lambda x: x not in (1, None) # noqa: E731 in_plain = lambda x: x in (1, 2) # noqa: E731 with self.sql_conf(_TRANSPILE_ON): - for func, cases in [ - (in_plain, [(None, False), (1, True), (2, True), (3, False)]), - (in_none, [(None, True), (1, True), (2, False)]), - (not_in_none, [(None, False), (1, False), (2, True)]), + for label, func, cases in [ + ("x in (1, 2)", in_plain, [(None, False), (1, True), (2, True), (3, False)]), + ("x in (1, None)", in_none, [(None, True), (1, True), (2, False)]), + ("x not in (1, None)", not_in_none, [(None, False), (1, False), (2, True)]), ]: pudf = UserDefinedFunction(func, BooleanType()) - self.assertTrue(pudf.transpiled, f"{func} should transpile") + self.assertTrue(pudf.transpiled, f"{label} should transpile") for value, expected in cases: - with self.subTest(func=str(func), value=value): + with self.subTest(func=label, value=value): df = self.spark.createDataFrame([Row(a=value)], schema=schema) [row] = df.select(pudf("a")).collect() self.assertEqual(row[0], expected) @@ -740,10 +766,13 @@ def test_udf_transpile_in_empty_container(self): conf = dict(_TRANSPILE_ON) conf["spark.sql.legacy.nullInEmptyListBehavior"] = legacy with self.sql_conf(conf): - for func, expected in [(empty_in, False), (empty_not_in, True)]: - with self.subTest(legacy=legacy, func=str(func)): + for label, func, expected in [ + ("x in ()", empty_in, False), + ("x not in ()", empty_not_in, True), + ]: + with self.subTest(legacy=legacy, func=label): pudf = UserDefinedFunction(func, BooleanType()) - self.assertTrue(pudf.transpiled, f"{func} should transpile") + self.assertTrue(pudf.transpiled, f"{label} should transpile") df = self.spark.createDataFrame([Row(a=None)], schema=schema) [row] = df.select(pudf("a")).collect() self.assertEqual(row[0], expected) @@ -774,30 +803,30 @@ def test_udf_transpile_comparison_fallbacks(self): in_long_rt = lambda x: x in (1, 2) # noqa: E731 six_longs = "a long, b long, c long, d long, e long, f long" - # (func, schema, rows, expected). `expected` is what interpreted Python - # produces, so the query must still run and return it. + # (label, func, schema, rows, expected). `expected` is what interpreted + # Python produces, so the query must still run and return it. cases = [ - (str_in, "a string", [("b",), ("z",)], [True, False]), - (bytes_in, "a long", [(97,), (1,)], [True, False]), - (dict_in, "a long", [(1,), (2,)], [True, False]), - (runtime_in, "a long, b long", [(1, 1), (3, 5)], [True, False]), - (chained_in, "a long, b long", [(0, 1), (1, 0)], [True, False]), - (mixed_in, "a long", [(1,), (3,)], [True, False]), - (bool_in, "a long", [(1,), (5,)], [True, False]), - (five_ops, six_longs, [(1, 2, 3, 4, 5, 6)], [True]), + ("`in` on a str", str_in, "a string", [("b",), ("z",)], [True, False]), + ("`in` on bytes", bytes_in, "a long", [(97,), (1,)], [True, False]), + ("`in` on a dict", dict_in, "a long", [(1,), (2,)], [True, False]), + ("runtime `in`", runtime_in, "a long, b long", [(1, 1), (3, 5)], [True, False]), + ("chain + `in`", chained_in, "a long, b long", [(0, 1), (1, 0)], [True, False]), + ("mixed-category `in`", mixed_in, "a long", [(1,), (3,)], [True, False]), + ("`in` over bools", bool_in, "a long", [(1,), (5,)], [True, False]), + ("a 5-op chain (over the cap)", five_ops, six_longs, [(1, 2, 3, 4, 5, 6)], [True]), ] with self.sql_conf(_TRANSPILE_ON): - for func, schema, rows, expected in cases: - with self.subTest(func=str(func)): + for label, func, schema, rows, expected in cases: + with self.subTest(func=label): with _warnings.catch_warnings(record=True): _warnings.simplefilter("always") pudf = UserDefinedFunction(func, BooleanType()) - self.assertEqual([], pudf.transpiled, f"{func} must NOT transpile") + self.assertEqual([], pudf.transpiled, f"{label} must NOT transpile") df = self.spark.createDataFrame(rows, schema) self.assertEqual( [r[0] for r in df.select(pudf(*df.columns)).collect()], expected, - f"{func}: interpreted result diverged", + f"{label}: interpreted result diverged", ) # `0 < x < "z"` raises in interpreted Python too (cross-type # ordering), so assert only the fallback for it. @@ -2503,8 +2532,11 @@ def test_udf_transpile_plan_elision(self): # EvalPython node disappears for them too. bounded = lambda x: 0 < x < 10 # noqa: E731 member = lambda x: x in (1, 2, 3) # noqa: E731 - for func, expected in [(bounded, [1, 7]), (member, [1])]: - with self.subTest(func=str(func)): + for label, func, expected in [ + ("0 < x < 10", bounded, [1, 7]), + ("x in (1,2,3)", member, [1]), + ]: + with self.subTest(func=label): f = UserDefinedFunction(func, BooleanType()) self.assertTrue(f.transpiled) fdf = self.spark.createDataFrame([(7,), (11,), (1,)], "a long").filter(f("a")) @@ -2547,9 +2579,9 @@ def test_udf_transpile_nondeterministic_argument_falls_back(self): band = lambda x: 50 < x < 60 # noqa: E731 band_or = lambda x: x < 50 or x > 60 # noqa: E731 with self.sql_conf(_TRANSPILE_ON): - for func in (band, band_or): + for label, func in (("50 < x < 60", band), ("x < 50 or x > 60", band_or)): pudf = UserDefinedFunction(func, BooleanType()) - self.assertTrue(pudf.transpiled, f"{func} should transpile on its own") + self.assertTrue(pudf.transpiled, f"{label} should transpile on its own") base = self.spark.range(0, 8).select((col("id") * 10).alias("a")) # Deterministic argument -> transpiled and elided. @@ -2557,7 +2589,7 @@ def test_udf_transpile_nondeterministic_argument_falls_back(self): self.assertEqual( 0, self._eval_python_count(det), - f"{func}: a deterministic argument should still be elided", + f"{label}: a deterministic argument should still be elided", ) # Non-deterministic argument -> must fall back. @@ -2565,7 +2597,7 @@ def test_udf_transpile_nondeterministic_argument_falls_back(self): self.assertEqual( 1, self._eval_python_count(nondet), - f"{func}: a non-deterministic argument must NOT be transpiled", + f"{label}: a non-deterministic argument must NOT be transpiled", ) # And it still produces results. self.assertEqual(8, len(nondet.collect())) diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 89219e404a623..88cbc7e941f75 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -102,6 +102,7 @@ import contextlib import inspect import itertools +import operator import sys import textwrap import threading @@ -147,6 +148,16 @@ _MAX_CHAIN_OPS = 4 _MAX_IN_ELEMENTS = 64 +# The four ordering comparisons, which all lower through +# ``_lower_value_compare`` and differ only in the operator to apply and the +# symbol to name it in an error message. +_ORDERING_OPS: dict[type, Tuple[Callable[[Any, Any], Column], str]] = { + ast.Lt: (operator.lt, "<"), + ast.LtE: (operator.le, "<="), + ast.Gt: (operator.gt, ">"), + ast.GtE: (operator.ge, ">="), +} + class AbstractTranspiler(object): """Base class for transpilers. All experimental.""" @@ -201,42 +212,63 @@ def _is_definitely_basic_type(node: ast.AST) -> bool: return False -def _is_definitely_boolean(node: ast.AST) -> bool: - """Return True when ``node`` is statically guaranteed to produce a Python - ``bool`` (or ``None``, which round-trips through ``coalesce``). +def _is_boolean_expr(node: ast.AST, *, allow_null: bool) -> bool: + """Shared traversal behind ``_is_definitely_boolean`` (``allow_null=True``) + and ``_is_never_null_boolean`` (``allow_null=False``). - Used to gate ``if``/ternary lowering: we only allow the test expression - into Catalyst's ``when(coalesce(test, false), ...)`` form when it provably - produces a boolean. Everything else (bare Name, arithmetic, function calls, - subscript, ...) must force a fallback to interpreted Python instead of - silently diverging. + The two predicates differ only in whether a NULL-valued operand is + acceptable, so they share one match rather than two near-copies: with + separate copies every new lowering had to be reflected in both, and + forgetting one silently re-admits a NULL operand into ``_lazy_and``, where a + NULL ``CASE WHEN`` condition reads as not-taken and turns Python's ``None`` + into ``False``. """ match node: case ast.Constant(value=v): - return v is None or isinstance(v, bool) + # A literal None is "boolean" for coalesce purposes but IS NULL, so + # only the NULL-tolerant caller may accept it. + return isinstance(v, bool) or (allow_null and v is None) case ast.Compare(left=left, ops=ops, comparators=comparators): # Every Python comparison produces a bool, chained ones and - # ``in`` / ``not in`` included. Don't inspect an ``in`` right - # operand: it is a CONTAINER, which ``_is_definitely_basic_type`` - # rejects, yet the position still produces a bool. Being - # over-permissive is safe -- ``_convert_chunk`` fails closed on - # anything it cannot lower. + # ``in`` / ``not in`` included, and every comparison lowering is + # total (see ``_lower_compare_pair``) -- so a NULL-intolerant caller + # needs no extra restriction here, only the operand shapes below. + # Don't inspect an ``in`` right operand: it is a CONTAINER, which + # ``_is_definitely_basic_type`` rejects, yet the position still + # produces a bool. Being over-permissive is safe -- + # ``_convert_chunk`` fails closed on anything it cannot lower. return _is_definitely_basic_type(left) and all( isinstance(op, (ast.In, ast.NotIn)) or _is_definitely_basic_type(c) for op, c in zip(ops, comparators) ) - case ast.BoolOp(values=values): - return all(_is_definitely_boolean(v) for v in values) case ast.UnaryOp(op=ast.Not()): - # `not x` always produces bool. + # `not x` always produces bool, and lowers to + # ``coalesce(~operand, lit(True))``, which is total. return True + case ast.BoolOp(values=values): + return all(_is_boolean_expr(v, allow_null=allow_null) for v in values) case ast.IfExp(body=body, orelse=orelse): # Ternary is boolean only if both branches are. - return _is_definitely_boolean(body) and _is_definitely_boolean(orelse) + return _is_boolean_expr(body, allow_null=allow_null) and _is_boolean_expr( + orelse, allow_null=allow_null + ) case _: return False +def _is_definitely_boolean(node: ast.AST) -> bool: + """Return True when ``node`` is statically guaranteed to produce a Python + ``bool`` (or ``None``, which round-trips through ``coalesce``). + + Used to gate ``if``/ternary lowering: we only allow the test expression + into Catalyst's ``when(coalesce(test, false), ...)`` form when it provably + produces a boolean. Everything else (bare Name, arithmetic, function calls, + subscript, ...) must force a fallback to interpreted Python instead of + silently diverging. + """ + return _is_boolean_expr(node, allow_null=True) + + def _truthiness_col(cat: Optional[str], c: Column) -> Optional[Column]: """Return a boolean Column expressing Python's ``bool(c)`` for the given category. @@ -269,23 +301,7 @@ def _is_never_null_boolean(node: ast.AST) -> bool: ``not`` and if/ternary *tests* keep using ``_is_definitely_boolean``, since they wrap the operand in ``coalesce`` and so handle NULL explicitly. """ - match node: - case ast.Constant(value=v): - # A literal None is "boolean" for coalesce purposes but is NULL. - return isinstance(v, bool) - case ast.Compare(): - # Every comparison lowering is total (see ``_lower_compare_pair``), - # so only the operand shapes still need checking. - return _is_definitely_boolean(node) - case ast.UnaryOp(op=ast.Not()): - # Lowered as ``coalesce(~operand, lit(True))``, which is total. - return True - case ast.BoolOp(values=values): - return all(_is_never_null_boolean(v) for v in values) - case ast.IfExp(body=body, orelse=orelse): - return _is_never_null_boolean(body) and _is_never_null_boolean(orelse) - case _: - return False + return _is_boolean_expr(node, allow_null=False) def _lazy_and(cols: List[Column]) -> Column: @@ -672,22 +688,11 @@ def _lower_compare_pair( return self._lower_eq(params, left_node, right_node, equal=True) case ast.NotEq(): return self._lower_eq(params, left_node, right_node, equal=False) - case ast.Lt(): - return self._lower_value_compare( - params, left_node, right_node, lambda l, r: l < r, "<" - ) - case ast.LtE(): - return self._lower_value_compare( - params, left_node, right_node, lambda l, r: l <= r, "<=" - ) - case ast.Gt(): - return self._lower_value_compare( - params, left_node, right_node, lambda l, r: l > r, ">" - ) - case ast.GtE(): - return self._lower_value_compare( - params, left_node, right_node, lambda l, r: l >= r, ">=" - ) + case ast.Lt() | ast.LtE() | ast.Gt() | ast.GtE(): + # All four orderings share `_lower_value_compare` and differ only + # in the operator, so look it up rather than repeating the call. + op_fn, op_str = _ORDERING_OPS[type(op)] + return self._lower_value_compare(params, left_node, right_node, op_fn, op_str) case _: raise UnsupportedOperationException( f"comparison operator {type(op).__name__} is not supported by the transpiler" From 1dee547c9adaff2e792140ffabece434e8f4df3c Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Fri, 14 Aug 2026 23:14:21 +0000 Subject: [PATCH 03/18] [SPARK-58456][PYTHON][SQL] Narrow the transpiler change to chained comparisons Scopes this change down to chained comparisons and pins parameter categories in the new hypothesis tests. - Drop `in` / `not in` lowering, deferred to its own ticket. `_convert_chunk` now refuses them so the UDF falls back to interpreted Python, and `_lower_in` / `_const_container_elements` are removed with their tests. This also removes a silent wrong answer the lowering carried: `when(subject.isNull(), lit(has_none))` could not tell a Python `None` probe from arithmetic that produced NULL, so `(x + 1) in (1, 2)` on a NULL row returned False where Python raises TypeError. - Drop `_lazy_and` / `_lazy_or`. `and` / `or` return to the pre-existing `&` / `|` fold, so already-shipped UDFs keep the splittable `And` that predicate pushdown, partition pruning and InferFiltersFromConstraints need; chained comparisons fold the same way. `&` short-circuits at evaluation time, so the runtime semantics hold. - Replace the per-form `_MAX_CHAIN_OPS` / `_MAX_IN_ELEMENTS` caps with a single `_MAX_COMPARISON_TERMS`. A token count is the wrong bound and a cap per syntax form does not converge; a configurable size budget is the follow-up. - Annotate the new hypothesis UDFs' parameters. This emits one option variant instead of two and drops variants those tests can never bind. It does NOT speed the suite up -- measured 73s untyped vs 76s typed over the five tests at 15 examples, because the cost is the two Spark jobs per example against a fresh one-row DataFrame. Batching rows is the lever that would matter. - Point the NaN divergence comments at SPARK-58781, and note they pin current behavior rather than endorse it. Known incomplete, and the reason this is not submittable yet: - Two tests are deliberately red pending the analyzer determinism guard, one Python and one Scala. - Reverting `and` / `or` reopened a NULL divergence: a ternary evaluating to None reaches the fold and Spark's `NULL AND false` collapses to false where Python returns None. The `_is_never_null_boolean` gate that closed it needs restoring. - Ten `in` cases in `test_udf_transpile_lowers_operators`, plus `ternary_in` in the return-type test, still assert transpilation and now fail. - With the `&` fold, a chain under a filter over a join can have its `raise_error` NULL guard split out and pushed below the join, raising on rows Python never evaluates. This needs `RaiseError` to override `Expression.throwable`, which lands separately and is a prerequisite. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../tests/test_udf_transpile_hypothesis.py | 94 +--- .../sql/tests/test_udf_transpile_unit.py | 177 ++------ python/pyspark/sql/transpile.py | 409 +++++------------- 3 files changed, 164 insertions(+), 516 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py index 726eb17fb862a..bffa694fdbda3 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -384,16 +384,33 @@ def neq_pair(x, y): return x != y -def chained_bounds(x): +def chained_bounds(x: int) -> bool: # Single-parameter chained comparison: ``ast.Compare`` with two ops, which # Python evaluates as ``(0 < x) and (x < 10)`` with ``x`` evaluated once. # No None guard, so a NULL input reaches the first link's raise guard -- # matching Python's TypeError -- and ``_run``'s "both raised" equivalence # covers it. + # + # The parameters here and in the chains below are ANNOTATED deliberately. + # An untyped parameter is transpiled as both numeric and string, so the + # option matrix carries variants these tests can never use: they bind Long + # columns, so only the all-numeric variant can match. Pinning the category + # emits one option instead of two (measured: 1 vs 2 -- ``_param_category_ + # combos`` already collapses three untyped params to all-numeric plus + # all-string rather than 2**3). + # + # This costs no coverage, but do not expect it to speed the suite up: the + # per-example cost is dominated by the two Spark jobs ``_run`` executes + # against a fresh one-row DataFrame, not by transpilation. Measured over + # these five tests at 15 examples: 73s untyped vs 76s typed, i.e. no + # improvement. Batching rows per example is the lever that would matter. + # + # The strategies still generate ``None``; the annotation pins the category, + # not nullability. return 0 < x < 10 -def chained_lt(x, y, z): +def chained_lt(x: int, y: int, z: int) -> bool: # Three-column ordering chain. The interesting rows are the ones where an # earlier link is False and a later operand is NULL: Python short-circuits # and returns False without ever comparing against the NULL, so the @@ -401,13 +418,13 @@ def chained_lt(x, y, z): return x < y < z -def chained_ge(x, y, z): +def chained_ge(x: int, y: int, z: int) -> bool: # Reversed-direction chain, so the short-circuit fires on a different set # of rows than ``chained_lt``. return x >= y >= z -def chained_eq(x, y, z): +def chained_eq(x: int, y: int, z: int) -> bool: # ``x == y == z`` -- equality never raises in Python, so every one of the # eight NULL combinations has a defined boolean answer that the # four-branch ``_lower_eq`` per link plus the short-circuit fold must @@ -416,41 +433,12 @@ def chained_eq(x, y, z): return x == y == z -def mixed_chain(x, y, z): +def mixed_chain(x: int, y: int, z: int) -> bool: # Mixed operators in one chain (``ast.LtE`` then ``ast.Gt``), so the two # links go through different lowering helpers. return x <= y > z -def in_literal(x): - # ``ast.Compare(In)`` over a constant tuple. Python's ``in`` is - # equality-based and never raises on a None probe (``None in (1, 2, 3)`` is - # False), unlike Spark's ``IN``, which returns NULL for a NULL probe. - return x in (1, 2, 3) - - -def not_in_literal(x): - # Sister of ``in_literal`` for ast.Compare(NotIn); also exercises the - # negation, which is only exact because the ``in`` lowering never yields - # NULL. - return x not in (1, 2, 3) - - -def in_with_none(x): - # A None element in the container: ``None in (1, None)`` is True in Python, - # while Spark's ``IN`` would return NULL for any non-matching probe if the - # NULL element were left in the list. - return x in (1, None) - - -def in_and_compare(x, y): - # ``in`` as an ``and`` operand, so the container membership has to be - # recognized as a non-NULL boolean by the BoolOp gate. When the membership - # test is False, Python never evaluates ``y > 0``, so a NULL ``y`` must not - # raise. - return x in (1, 2) and y > 0 - - # A lambda captured at module scope so ``inspect.getsource`` can read # its definition. Exercises the ``ast.Lambda`` branch in # ``_get_function_from_ast``. @@ -846,44 +834,6 @@ def test_mixed_chain_matches_python(self, x, y, z): # ---- `in` / `not in` --------------------------------------------- - @_hyp_settings - @given(value=_long_strategy) - @_seed_examples(_LONG_EDGES) - def test_in_literal_matches_python(self, value): - df = self._single_arg_df(value, LongType()) - transpiled, interpreted = self._run(in_literal, BooleanType(), df, "a") - self.assertEqual(transpiled, interpreted, f"in_literal mismatch on {value!r}") - - @_hyp_settings - @given(value=_long_strategy) - @_seed_examples(_LONG_EDGES) - def test_not_in_literal_matches_python(self, value): - df = self._single_arg_df(value, LongType()) - transpiled, interpreted = self._run(not_in_literal, BooleanType(), df, "a") - self.assertEqual(transpiled, interpreted, f"not_in_literal mismatch on {value!r}") - - @_hyp_settings - @given(value=_long_strategy) - @_seed_examples(_LONG_EDGES) - def test_in_with_none_matches_python(self, value): - # ``None in (1, None)`` is True in Python. The lowering answers the - # NULL-probe case with a literal and drops the NULL element from - # the emitted ``IN`` list, so this must agree for a NULL probe as - # well as for matching and non-matching values. - df = self._single_arg_df(value, LongType()) - transpiled, interpreted = self._run(in_with_none, BooleanType(), df, "a") - self.assertEqual(transpiled, interpreted, f"in_with_none mismatch on {value!r}") - - @_hyp_settings - @given(x=_long_strategy, y=_long_strategy) - @_seed_pair_examples(_BOOLEAN_PAIR_EDGES) - def test_in_and_compare_matches_python(self, x, y): - df = self._two_long_arg_df(x, y) - transpiled, interpreted = self._run(in_and_compare, BooleanType(), df, "a", "b") - self.assertEqual( - transpiled, interpreted, f"in_and_compare mismatch on (x={x!r}, y={y!r})" - ) - class UDFTranspileHypothesisGatingTests(unittest.TestCase): """Smoke tests that always run, regardless of the env gate. diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index aef0e8552d8d3..54e1c6b50db3b 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -24,7 +24,6 @@ property-based suite lives in ``test_udf_transpile_hypothesis.py``. """ -import re import unittest from pyspark.sql import Row @@ -565,28 +564,21 @@ def chained(a, b, c): self.assertIn("cannot compare NULL", str(ctx.exception)) def test_udf_transpile_boolop_short_circuits(self): - # The `and` / `or` counterpart of the chain short-circuit, and the - # regression test for the pushdown hazard the CASE WHEN lowering closes. + # The `and` / `or` counterpart of the chain short-circuit. `&` short- + # circuits at evaluation time -- both `And.eval` and `And.doGenCode` skip + # the right operand when the left is false -- so a `raise_error` NULL + # guard in the second operand must not fire when the first already + # decided the result. # - # `x > 0 and y > 0` used to lower to `And(caseWhen1, caseWhen2)`. - # `splitConjunctivePredicates` decomposes `And` and nothing else, and - # `PushPredicateThroughJoin` treats a conjunct as pushable when - # `deterministic && !throwable` -- but `RaiseError` does not override - # `Expression.throwable`, so the NULL guard looks harmless and the - # second conjunct gets pushed to the far side of a join, where it is - # evaluated on rows the first conjunct excluded. That can raise where - # Python short-circuits to False. A single CASE WHEN is atomic to - # predicate splitting, so it cannot be separated that way. - # - # The guard below is a PLAN-SHAPE assertion, deliberately, because the - # runtime symptom is scheduling-dependent: with the old `&` lowering the - # pushed-down `raise_error` fires only on some runs (whether the failing - # task gets far enough before the join yields no rows). Asserting on the - # optimized plan instead is deterministic -- the guard must not be pushed - # BENEATH the join, onto one side, where it would be evaluated on rows - # the other operand excluded. It does end up inside the join CONDITION, - # which is safe for the same reason the lowering exists: the predicate - # stays whole, so the `and` still short-circuits. + # NOT asserted here: that the guard stays whole through the optimizer. + # `splitConjunctivePredicates` decomposes `And`, and + # `PushPredicateThroughJoin` will push a conjunct it believes harmless, + # which a `raise_error` currently is because `RaiseError` does not + # override `Expression.throwable`. Overriding it is the prerequisite fix; + # see DEFERED_AND_OR_CASE_WHEN_ISSUE. Until then the hazard is real but + # its runtime symptom is scheduling-dependent (the pushed-down guard + # fires only when the failing task gets far enough before the join yields + # no rows), so there is nothing deterministic to assert. def both_positive(a, c): return a > 0 and c > 0 @@ -607,53 +599,15 @@ def either_positive(a, c): self.assertEqual([True], [r[0] for r in df.select(or_udf("a", "c")).collect()]) # Same thing across a join, where the NULL operand lives entirely on - # the right side and would be a pushdown candidate if the predicate - # decomposed into conjuncts. + # the right side. The `and` must still short-circuit rather than + # raising on the row whose `c` is NULL. left = self.spark.createDataFrame([(1, -5)], "k long, a long") right = self.spark.createDataFrame([(1, None)], "k long, c long") joined = left.join(right, "k") self.assertEqual([False], [r[0] for r in joined.select(and_udf("a", "c")).collect()]) + # Filtering on it returns Python's answer rather than raising. filtered = joined.filter(and_udf("a", "c")) - plan = filtered._jdf.queryExecution().optimizedPlan().toString() - lines = plan.split("\n") - # Assert the guard is actually there first: without this every check - # below passes vacuously the moment the guard folds away or the UDF - # stops transpiling. - self.assertTrue( - any("raise_error" in line for line in lines), - f"expected the NULL raise guard somewhere in the plan:\n{plan}", - ) - # An optimized plan prints root-first, depth-first, so a node pushed - # below the Join appears on a LATER line than the Join itself. Match - # the node name at the start of the line rather than the substring - # "Join", which also occurs inside a printed join condition. - join_lines = [i for i, line in enumerate(lines) if re.match(r"^[\s:+|-]*Join ", line)] - self.assertEqual(1, len(join_lines), f"expected exactly one Join node:\n{plan}") - join_line = join_lines[0] - # The hazard is the guard being pushed to ONE SIDE of the join, where - # it would see rows that never join. Landing in the join CONDITION is - # acceptable: the predicate stays whole, so `and` still - # short-circuits and the guard is only reached for candidate pairs. - # That is where it lands today -- hence `> join_line`, which excludes - # the Join's own line -- so also pin that it did not decompose. - pushed = [ - line for i, line in enumerate(lines) if i > join_line and "raise_error" in line - ] - self.assertEqual( - [], - pushed, - "the NULL raise guard was pushed below the Join, so it can be " - f"evaluated on rows the other operand excluded:\n{plan}", - ) - self.assertIn( - "raise_error", - lines[join_line], - "expected the whole predicate to stay in the join condition; if " - "it moved above the Join that is also safe, but the lowering " - f"changed and this test should be re-read:\n{plan}", - ) - # And it still returns Python's answer rather than raising. self.assertEqual([], filtered.select("k").collect()) def test_udf_transpile_chain_operator_mix_and_eq_nulls(self): @@ -722,61 +676,9 @@ def chained_eq(a, b, c): [(1, 2, 3, 4, 5), (1, 2, 3, 5, 4)], ), [True, False], - "a chain of four operators is at the cap and must transpile", + "a chain of four operators is well under the cap and must transpile", ) - def test_udf_transpile_in_lowered(self): - # ``in`` / ``not in`` over a literal container is equality-based in - # Python and never raises on a None probe: ``None in (1, 2)`` is False - # and ``None in (1, None)`` is True. Spark's ``IN`` returns NULL for a - # NULL probe or a NULL list element, so the lowering answers the - # NULL-probe case directly and drops NULL elements from the list. - from pyspark.sql.types import StructField, StructType - - schema = StructType([StructField("a", LongType(), nullable=True)]) - in_none = lambda x: x in (1, None) # noqa: E731 - not_in_none = lambda x: x not in (1, None) # noqa: E731 - in_plain = lambda x: x in (1, 2) # noqa: E731 - with self.sql_conf(_TRANSPILE_ON): - for label, func, cases in [ - ("x in (1, 2)", in_plain, [(None, False), (1, True), (2, True), (3, False)]), - ("x in (1, None)", in_none, [(None, True), (1, True), (2, False)]), - ("x not in (1, None)", not_in_none, [(None, False), (1, False), (2, True)]), - ]: - pudf = UserDefinedFunction(func, BooleanType()) - self.assertTrue(pudf.transpiled, f"{label} should transpile") - for value, expected in cases: - with self.subTest(func=label, value=value): - df = self.spark.createDataFrame([Row(a=value)], schema=schema) - [row] = df.select(pudf("a")).collect() - self.assertEqual(row[0], expected) - - def test_udf_transpile_in_empty_container(self): - # ``x in ()`` is always False in Python, including for a None probe. - # The lowering spells that out with a literal rather than emitting - # ``IN ()``, whose NULL-probe result depends on - # ``spark.sql.legacy.nullInEmptyListBehavior``. Pin it under both - # settings so the answer cannot drift with that flag. - from pyspark.sql.types import StructField, StructType - - empty_in = lambda x: x in () # noqa: E731 - empty_not_in = lambda x: x not in () # noqa: E731 - schema = StructType([StructField("a", LongType(), nullable=True)]) - for legacy in (False, True): - conf = dict(_TRANSPILE_ON) - conf["spark.sql.legacy.nullInEmptyListBehavior"] = legacy - with self.sql_conf(conf): - for label, func, expected in [ - ("x in ()", empty_in, False), - ("x not in ()", empty_not_in, True), - ]: - with self.subTest(legacy=legacy, func=label): - pudf = UserDefinedFunction(func, BooleanType()) - self.assertTrue(pudf.transpiled, f"{label} should transpile") - df = self.spark.createDataFrame([Row(a=None)], schema=schema) - [row] = df.select(pudf("a")).collect() - self.assertEqual(row[0], expected) - def test_udf_transpile_comparison_fallbacks(self): # Comparison forms the transpiler must REFUSE, each for its own reason. # Every one asserts both that nothing transpiled and that a query using @@ -798,10 +700,8 @@ def test_udf_transpile_comparison_fallbacks(self): mixed_in = lambda x: x in (1, "a") # noqa: E731 bool_in = lambda x: x in (True, False) # noqa: E731 cross_chain = lambda x: 0 < x < "z" # noqa: E731 - five_ops = lambda a, b, c, d, e, f: a < b < c < d < e < f # noqa: E731 chain_long_rt = lambda x: 0 < x < 10 # noqa: E731 in_long_rt = lambda x: x in (1, 2) # noqa: E731 - six_longs = "a long, b long, c long, d long, e long, f long" # (label, func, schema, rows, expected). `expected` is what interpreted # Python produces, so the query must still run and return it. @@ -813,7 +713,6 @@ def test_udf_transpile_comparison_fallbacks(self): ("chain + `in`", chained_in, "a long, b long", [(0, 1), (1, 0)], [True, False]), ("mixed-category `in`", mixed_in, "a long", [(1,), (3,)], [True, False]), ("`in` over bools", bool_in, "a long", [(1,), (5,)], [True, False]), - ("a 5-op chain (over the cap)", five_ops, six_longs, [(1, 2, 3, 4, 5, 6)], [True]), ] with self.sql_conf(_TRANSPILE_ON): for label, func, schema, rows, expected in cases: @@ -908,20 +807,17 @@ def chain_of_compare(a, b, c): self.assertEqual([True], [r[0] for r in df.select(pudf("c", "x", "y")).collect()]) def test_udf_transpile_lowering_caps(self): - # The chain-operator and container-element caps bound plan growth: an - # interior chain operand's subtree is emitted once per adjacent pair, - # and every input-category variant carries its own copy for - # CheckAnalysis to walk. Exercise both boundaries directly rather than - # writing a 65-element tuple into the test source (a container built - # with `eval` has no retrievable source, so the UDF would fall back for - # the wrong reason and the test would pass vacuously). + # The chain cap bounds plan growth: an interior chain operand's subtree is + # emitted once per adjacent pair, and every input-category variant carries + # its own copy for CheckAnalysis to walk. Exercise the boundary directly + # rather than through a UDF, so the refusal cannot be confused with a + # fallback for some other reason. import ast as _ast from pyspark.errors import UnsupportedOperationException from pyspark.sql.transpile import ( CatalystTranspiler, - _MAX_CHAIN_OPS, - _MAX_IN_ELEMENTS, + _MAX_COMPARISON_TERMS, ) transpiler = CatalystTranspiler() @@ -930,28 +826,24 @@ def test_udf_transpile_lowering_caps(self): def compare_node(src): return _ast.parse(src, mode="eval").body - def in_src(n): - return "x in (%s)" % ", ".join(str(i) for i in range(n)) - def chain_src(n): # n operators over n + 1 operands: x < 0 < 1 < ... < n - 1 return " < ".join(["x"] + [str(i) for i in range(n)]) with self.sql_conf(_TRANSPILE_ON): - # At each cap the lowering succeeds. - node = compare_node(in_src(_MAX_IN_ELEMENTS)) - transpiler._lower_in(["x"], node.left, node.comparators[0], negate=False) - transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_CHAIN_OPS))) + # At the cap the lowering succeeds. + transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_TERMS))) # One past it, it refuses so the UDF falls back. with self.assertRaises(UnsupportedOperationException) as ctx: - node = compare_node(in_src(_MAX_IN_ELEMENTS + 1)) - transpiler._lower_in(["x"], node.left, node.comparators[0], negate=False) - self.assertIn("elements", str(ctx.exception)) - with self.assertRaises(UnsupportedOperationException) as ctx: - transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_CHAIN_OPS + 1))) + transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_TERMS + 1))) self.assertIn("operators", str(ctx.exception)) + # `in` / `not in` are refused outright, deferred to SPARK-56925. + with self.assertRaises(UnsupportedOperationException) as ctx: + transpiler._convert_chunk(["x"], compare_node("x in (1, 2)")) + self.assertIn("`in` / `not in`", str(ctx.exception)) + def test_udf_transpile_multi_row(self): # Every other transpile test uses a 1-row DataFrame; this one runs # the same arithmetic transpile on a multi-row input to catch any @@ -3004,6 +2896,11 @@ def test_udf_transpile_known_value_divergences(self): # second link whenever `a < b`, and Spark orders NaN above every value # (Python's `b < nan` is False); `a == b == NaN` likewise, since Spark # treats `NaN = NaN` as true where Python's `nan == nan` is False. + # + # These assertions pin the CURRENT (divergent) behavior so a change is + # visible, not an endorsement of it. Whether to special-case NaN, fall + # back to interpreted Python, or document it as a known limitation is + # tracked in SPARK-58781. nan_chain = lambda a, b, c: a < b < c # noqa: E731 nan_eq_chain = lambda a, b, c: a == b == c # noqa: E731 nan = float("nan") diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 88cbc7e941f75..b4a530116ee34 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -92,10 +92,11 @@ An argument the body never reads is not computed at all. Constructs that short-circuit in Python -- ``and`` / ``or`` and chained -comparisons such as ``0 <= x <= 100`` -- lower to nested ``CASE WHEN``s -rather than to ``&`` / ``|``, so that an operand Python never evaluates -cannot raise; see ``_lazy_and``. Their operands must also be provably -non-NULL booleans; see ``_is_never_null_boolean``. +comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``, which +short-circuit at evaluation time. Their operands must be provably boolean; +see ``_is_definitely_boolean``. Note the fold is splittable by +``splitConjunctivePredicates``, so a ``raise_error`` NULL guard can be pushed +across a join until ``RaiseError`` overrides ``Expression.throwable``. """ import ast @@ -139,14 +140,17 @@ from pyspark.sql._typing import DataTypeOrString -# Caps on how large a comparison we will lower, to bound plan growth: longer -# ones fall back to interpreted Python. Both are generous next to what shows up -# in real code (``0 <= x <= 100``, an allowlist of a few dozen values), and both -# matter because the lowered subtree is duplicated -- once per adjacent pair for -# a chain operand, once per input-category variant for every option -- and -# CheckAnalysis walks each copy. -_MAX_CHAIN_OPS = 4 -_MAX_IN_ELEMENTS = 64 +# Cap on how large a chained comparison we will lower, to bound plan growth: +# longer ones fall back to interpreted Python. It is generous next to what shows +# up in real code (``0 <= x <= 100``), and it matters because the lowered subtree +# is duplicated -- once per adjacent pair for a chain operand, once per +# input-category variant for every option -- and CheckAnalysis walks each copy. +# +# A token count is the wrong bound: it does not measure what actually matters, +# the size of the emitted expression tree, and a cap per syntax form does not +# converge as forms are added. A real size budget, configurable and counted +# across the whole lowering, is the follow-up that replaces this. +_MAX_COMPARISON_TERMS = 64 # The four ordering comparisons, which all lower through # ``_lower_value_compare`` and differ only in the operator to apply and the @@ -212,63 +216,35 @@ def _is_definitely_basic_type(node: ast.AST) -> bool: return False -def _is_boolean_expr(node: ast.AST, *, allow_null: bool) -> bool: - """Shared traversal behind ``_is_definitely_boolean`` (``allow_null=True``) - and ``_is_never_null_boolean`` (``allow_null=False``). +def _is_definitely_boolean(node: ast.AST) -> bool: + """Return True when ``node`` is statically guaranteed to produce a Python + ``bool`` (or ``None``, which round-trips through ``coalesce``). - The two predicates differ only in whether a NULL-valued operand is - acceptable, so they share one match rather than two near-copies: with - separate copies every new lowering had to be reflected in both, and - forgetting one silently re-admits a NULL operand into ``_lazy_and``, where a - NULL ``CASE WHEN`` condition reads as not-taken and turns Python's ``None`` - into ``False``. + Used to gate ``if``/ternary lowering and ``and`` / ``or``: we only allow the + expression into Catalyst's ``when(coalesce(test, false), ...)`` form, or onto + ``&`` / ``|``, when it provably produces a boolean. Everything else (bare + Name, arithmetic, function calls, subscript, ...) must force a fallback to + interpreted Python instead of silently diverging. """ match node: case ast.Constant(value=v): - # A literal None is "boolean" for coalesce purposes but IS NULL, so - # only the NULL-tolerant caller may accept it. - return isinstance(v, bool) or (allow_null and v is None) - case ast.Compare(left=left, ops=ops, comparators=comparators): - # Every Python comparison produces a bool, chained ones and - # ``in`` / ``not in`` included, and every comparison lowering is - # total (see ``_lower_compare_pair``) -- so a NULL-intolerant caller - # needs no extra restriction here, only the operand shapes below. - # Don't inspect an ``in`` right operand: it is a CONTAINER, which - # ``_is_definitely_basic_type`` rejects, yet the position still - # produces a bool. Being over-permissive is safe -- - # ``_convert_chunk`` fails closed on anything it cannot lower. - return _is_definitely_basic_type(left) and all( - isinstance(op, (ast.In, ast.NotIn)) or _is_definitely_basic_type(c) - for op, c in zip(ops, comparators) - ) + return v is None or isinstance(v, bool) + case ast.Compare(left=left, comparators=comparators): + # All comparison operators of simple types bool. Chained comparisons + # included: every operand still has to be a basic type. + return all(_is_definitely_basic_type(v) for v in comparators + [left]) + case ast.BoolOp(values=values): + return all(_is_definitely_boolean(v) for v in values) case ast.UnaryOp(op=ast.Not()): - # `not x` always produces bool, and lowers to - # ``coalesce(~operand, lit(True))``, which is total. + # `not x` always produces bool. return True - case ast.BoolOp(values=values): - return all(_is_boolean_expr(v, allow_null=allow_null) for v in values) case ast.IfExp(body=body, orelse=orelse): # Ternary is boolean only if both branches are. - return _is_boolean_expr(body, allow_null=allow_null) and _is_boolean_expr( - orelse, allow_null=allow_null - ) + return _is_definitely_boolean(body) and _is_definitely_boolean(orelse) case _: return False -def _is_definitely_boolean(node: ast.AST) -> bool: - """Return True when ``node`` is statically guaranteed to produce a Python - ``bool`` (or ``None``, which round-trips through ``coalesce``). - - Used to gate ``if``/ternary lowering: we only allow the test expression - into Catalyst's ``when(coalesce(test, false), ...)`` form when it provably - produces a boolean. Everything else (bare Name, arithmetic, function calls, - subscript, ...) must force a fallback to interpreted Python instead of - silently diverging. - """ - return _is_boolean_expr(node, allow_null=True) - - def _truthiness_col(cat: Optional[str], c: Column) -> Optional[Column]: """Return a boolean Column expressing Python's ``bool(c)`` for the given category. @@ -289,127 +265,6 @@ def _truthiness_col(cat: Optional[str], c: Column) -> Optional[Column]: return None -def _is_never_null_boolean(node: ast.AST) -> bool: - """Return True when ``node`` lowers to a boolean column that either raises - or yields True/False -- never NULL. - - Stricter than ``_is_definitely_boolean``, which also accepts operands that - may evaluate to NULL (a literal ``None``, a ternary with a ``None`` branch). - ``_lazy_and`` / ``_lazy_or`` need this: a CASE WHEN treats a NULL condition - as not-taken, so a NULL operand would silently become ``False`` where Python - returns ``None``. Bodies that fail this go back to interpreted Python. - ``not`` and if/ternary *tests* keep using ``_is_definitely_boolean``, since - they wrap the operand in ``coalesce`` and so handle NULL explicitly. - """ - return _is_boolean_expr(node, allow_null=False) - - -def _lazy_and(cols: List[Column]) -> Column: - """Fold ``cols`` into a boolean column with Python's ``and`` semantics, - evaluating each element only when every earlier one was True. - - Nests CASE WHENs instead of folding with ``&``. The reason is NOT that - ``And`` fails to short-circuit -- both ``And.eval`` and ``And.doGenCode`` - skip the right operand when the left is false. It is the optimizer: - ``splitConjunctivePredicates`` decomposes ``And`` and nothing else, and - ``PushPredicateThroughJoin`` gates pushdown on - ``cond.deterministic && !cond.throwable``. ``Expression.throwable`` defaults - to ``children.exists(_.throwable)`` and ``RaiseError`` does not override it - (its children are literals), so Catalyst believes the NULL guard that - ``_lower_value_compare`` emits cannot throw, and will push one conjunct to - the far side of a join where it sees rows another conjunct excluded. - - A CaseWhen resists that: it is atomic to predicate splitting, evaluates - branches in order (``CaseWhen.eval`` stops at the first TRUE condition), - stays non-foldable (``ConditionalExpression.foldable`` needs every child - foldable and ``RaiseError.foldable`` is false), and hides its lazy tail from - subexpression elimination -- ``EquivalentExpressions`` reaches the tail only - via ``branchGroups``, which intersects the branch values, and one of ours is - a ``Literal`` that ``updateExprTree`` skips, leaving that intersection - empty. - - SCOPE, precisely: this stops predicate splitting from tearing ONE lowered - predicate into independently-pushable conjuncts. It does NOT stop the whole - predicate being pushed as a single conjunct, since ``throwable`` is still - false for it -- ``left.join(right, "k").filter(f(col("c")))`` with ``c`` on - the right side still pushes below the join, so the guard fires on rows that - never join and raises where interpreted Python never saw them. That - divergence pre-dates chained comparisons (a single comparison already lowers - to one guarded CASE WHEN); only overriding ``throwable`` on ``RaiseError`` - in Catalyst closes it, which is a separate change because it shifts pushdown - behavior for every user of ``raise_error()`` / ``assert_true()``. - - Callers must have proven every element non-NULL (``_is_never_null_boolean``), - which is what lets this skip ``coalesce``. - - COST, and it is not free. When every operand column is non-nullable the - ``&`` fold used to simplify all the way down: ``NullPropagation`` folds - ``IsNull(e)`` to false for a non-nullable ``e``, so the guard disappears and - what remains is a plain ``And`` that ``splitConjunctivePredicates`` can tear - into conjuncts for datasource pushdown, partition pruning and - ``InferFiltersFromConstraints``. ``CaseWhen(Seq((c1, c2)), Some(false))`` has - no matching rule in ``SimplifyConditionals`` (which only collapses the - ``TrueLiteral`` / ``FalseLiteral`` branch shapes), so it survives as an - opaque predicate and none of that applies. Measured on a non-nullable - column: ``x > 20 and x < 80`` optimized to ``((id > 20) AND (id < 80))`` - under the old fold and stays ``CASE WHEN (id > 20) THEN (id < 80) ELSE - false END`` under this one. - - The trade is deliberate but narrow. The pushdown loss only bites when the - guard would have folded away -- i.e. non-nullable columns -- and that is - exactly the case where the ``And`` form was already SAFE, because with no - ``raise_error`` left there is nothing for the optimizer to relocate. When - the guard does survive (nullable columns) it blocks datasource pushdown in - either form, so nothing is lost. Overriding ``throwable`` on ``RaiseError`` - would let this go back to ``&`` / ``|`` and get both properties at once; - that is the right end state and is tracked separately. - """ - result = cols[-1] - for c in reversed(cols[:-1]): - result = when(c, result).otherwise(lit(False)) - return result - - -def _lazy_or(cols: List[Column]) -> Column: - """Fold ``cols`` with Python's ``or`` semantics, evaluating each element only - when every earlier one was False. The mirror of ``_lazy_and``; see there for - why this nests CASE WHENs rather than folding with ``|``. - """ - result = cols[-1] - for c in reversed(cols[:-1]): - result = when(c, lit(True)).otherwise(result) - return result - - -def _const_container_elements(node: ast.AST) -> Optional[List[ast.expr]]: - """Element nodes of a constant container literal, or ``None`` when ``node`` - is not one. - - Recognises ``ast.Tuple`` / ``ast.List`` / ``ast.Set`` displays whose every - element is an ``ast.Constant`` or a unary ``+`` / ``-`` on a non-bool numeric - ``ast.Constant`` (``x in (-1, 2)`` parses ``-1`` as a ``UnaryOp``). Anything - else -- a ``Name``, a ``Call``, a starred or nested element, a ``dict`` / - ``str`` / ``bytes`` right operand -- returns ``None`` so the caller falls - back; see ``CatalystTranspiler._lower_in`` for why element values must be - known at transpile time. - """ - if not isinstance(node, (ast.Tuple, ast.List, ast.Set)): - return None - for e in node.elts: - if isinstance(e, ast.Constant): - continue - if ( - isinstance(e, ast.UnaryOp) - and isinstance(e.op, (ast.USub, ast.UAdd)) - and isinstance(e.operand, ast.Constant) - and isinstance(e.operand.value, (int, float)) - and not isinstance(e.operand.value, bool) - ): - continue - return None - return list(node.elts) - - def _is_none_constant(node: ast.AST) -> bool: """Return True for the literal ``None``.""" return isinstance(node, ast.Constant) and node.value is None @@ -657,8 +512,8 @@ def _lower_compare_pair( widening as the only coercion Catalyst inserts, and under ANSI a failing cast raises instead of yielding NULL. - That invariant is what lets ``_lazy_and`` and ``_lower_in``'s negation - use plain ``when`` / ``~`` with no ``coalesce``. + That invariant is what lets the chain fold use plain ``&`` with no + ``coalesce``: every link is total, so none of them yields NULL. """ match op: case ast.Is() | ast.IsNot(): @@ -694,90 +549,14 @@ def _lower_compare_pair( op_fn, op_str = _ORDERING_OPS[type(op)] return self._lower_value_compare(params, left_node, right_node, op_fn, op_str) case _: + # Reached only if Python gains a new `cmpop`: the other eight are + # handled above and `in` / `not in` are refused in + # `_convert_chunk`. Kept so a new operator falls back rather than + # returning None from a `-> Column` function. raise UnsupportedOperationException( f"comparison operator {type(op).__name__} is not supported by the transpiler" ) - def _lower_in( - self, - params: List[str], - subject_node: ast.AST, - container_node: ast.AST, - negate: bool, - ) -> Column: - """Lower ``x in (c1, ..., cn)`` / ``x not in (...)`` over a literal - tuple / list / set of constants. - - Python's ``in`` over a container is EQUALITY-based and never raises on a - ``None`` probe: ``None in (1, 2)`` is False and ``None in (1, None)`` is - True. Spark's ``IN`` returns NULL when the probe is NULL *or* when any - list element is NULL, so we answer the NULL-probe case directly with a - literal and DROP the NULL elements from the ``isin`` list. Leaving one in - would make the whole expression NULL for a non-matching probe, breaking - the totality that ``not in``'s ``~`` relies on. - - The empty container is spelled out as ``lit(False)`` rather than - delegating to ``IN ()``: ``In.eval`` returns false for an empty list only - while ``spark.sql.legacy.nullInEmptyListBehavior`` is off; with it on, - ``OptimizeIn`` rewrites ``IN ()`` to ``If(IsNotNull(v), false, null)`` - and a NULL probe would yield NULL. - - Every element must share the SUBJECT's category. That guard is - load-bearing: ``InTypeCoercion`` finds no wider common type for e.g. - bigint-vs-string under ANSI, so ``In.checkInputDataTypes`` fails with - DATA_DIFF_TYPES -- and because the option rides along as a child of - ``TranspiledPythonUDF``, that fails the WHOLE query at CheckAnalysis - instead of falling back. Python's cross-type ``in`` is simply False, so - coercing would buy nothing anyway. - """ - elements = _const_container_elements(container_node) - if elements is None: - raise UnsupportedOperationException( - "`in` / `not in` is only supported over a literal tuple, list " - "or set of constants (e.g. `x in (1, 2)`): `in` on a str or " - "bytes is substring/subsequence containment rather than " - "equality (`'b' in 'abc'` is True while `'b' == 'abc'` is " - "False, and `b'ab'` iterates ints so `97 in b'ab'` is True), " - "dict membership tests keys, and a container holding runtime " - "values would need runtime NULL analysis; the transpiler falls " - "back to interpreted Python" - ) - if len(elements) > _MAX_IN_ELEMENTS: - raise UnsupportedOperationException( - f"`in` / `not in` container has {len(elements)} elements, more " - f"than the {_MAX_IN_ELEMENTS} the transpiler will lower; the " - "UDF falls back to interpreted Python to bound plan growth " - "across the input-type option matrix" - ) - # Raises for a literal-None or complex subject, dropping this variant. - subject_cat = self._category(params, subject_node) - has_none = False - value_nodes: List[ast.expr] = [] - for element in elements: - if _is_none_constant(element): - has_none = True - continue - # Classify a unary +/- element by its operand rather than leaning on - # `_category`'s catch-all, which happens to answer "numeric" for an - # `ast.UnaryOp` today. `_const_container_elements` has already - # restricted the operand to a non-bool numeric constant. - classify = element.operand if isinstance(element, ast.UnaryOp) else element - element_cat = self._category(params, classify) - if element_cat != subject_cat: - raise UnsupportedOperationException( - f"`in` / `not in` element category ({element_cat}) does not " - f"match the subject's ({subject_cat}); Python compares " - "across types as unequal while Spark would coerce or fail " - "analysis, so the transpiler falls back to interpreted " - "Python" - ) - value_nodes.append(element) - subject_col = self._convert_chunk(params, subject_node) - element_cols = [self._convert_chunk(params, e) for e in value_nodes] - membership = subject_col.isin(*element_cols) if element_cols else lit(False) - result = when(subject_col.isNull(), lit(has_none)).otherwise(membership) - return result.__invert__() if negate else result - def _category(self, params: List[str], node: ast.AST) -> str: """Infer ``"numeric"`` or ``"string"`` for ``node`` under the current ``self._param_categories`` assumption (set per input-type variant). @@ -957,27 +736,38 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # bitwise instead. Require statically known booleans so the # caller falls back rather than emitting a divergent plan. # - # The gate is `_is_never_null_boolean`, not - # `_is_definitely_boolean`, because `_lazy_and` / `_lazy_or` need - # non-NULL operands: a NULL CASE WHEN condition is not-taken, so - # `and` would yield False where Python returns None. That - # refuses a literal None operand (`None and (x > 0)` is None in - # Python) and a ternary with a None branch -- both of which - # diverge under Spark's three-valued `&` / `|` as well. - if not all(_is_never_null_boolean(v) for v in values): + if not all(_is_definitely_boolean(v) for v in values): + raise UnsupportedOperationException( + "`and` / `or` operand type is not statically known " + "to be boolean; Spark's `&` / `|` are bitwise, not " + "Python truthiness, so the transpiler refuses to " + "lower this and the UDF falls back to interpreted " + "Python" + ) + # A literal None operand short-circuits differently: Python's + # `None and (x > 0)` returns None regardless of x, but Spark's + # three-valued `null AND false` is false (and `null OR true` is + # true), so the lowered form diverges. `_is_definitely_boolean` + # accepts None for `not`/if-test contexts where coalesce handles + # it; here it must force a fallback instead. + if any(_is_none_constant(v) for v in values): raise UnsupportedOperationException( - "`and` / `or` operand is not statically known to be a " - "non-NULL boolean; Python's short-circuit returns the " - "operand itself (so `None and x` is None) and Spark's " - "`&` / `|` are bitwise rather than Python truthiness, " - "so the transpiler refuses to lower this and the UDF " + "literal None operand in `and` / `or` cannot be lowered: " + "Spark's three-valued logic diverges from Python's " + "short-circuit-return-operand semantics, so the UDF " "falls back to interpreted Python" ) cols = [self._convert_chunk(params, v) for v in values] if isinstance(op, ast.And): - return _lazy_and(cols) + result = cols[0] + for c in cols[1:]: + result = result & c + return result if isinstance(op, ast.Or): - return _lazy_or(cols) + result = cols[0] + for c in cols[1:]: + result = result | c + return result raise UnsupportedOperationException(f"BoolOp operator {op} is not supported") case ast.IfExp(test=test, body=body_expr, orelse=orelse_expr): # Ternary `body if test else orelse` -- shares the @@ -1002,37 +792,41 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: orelse[0] if orelse else None, ) case ast.Compare(left, ops, comps): - # `in` / `not in` are container membership, not a chainable - # value comparison, so only the un-chained single-operator form - # is lowered. `a < b in (1, 2)` is legal Python; refuse it - # explicitly rather than relying on `_convert_chunk` - # incidentally rejecting the ast.Tuple operand. + # `in` / `not in` are container membership rather than a chainable + # value comparison, and lowering them needs its own NULL analysis + # (Python's `in` is equality-based and never raises on a `None` + # probe, while an unguarded NULL subject is indistinguishable from + # one). Refuse them here so the UDF falls back to interpreted + # Python; support is deferred to SPARK-56925. if any(isinstance(o, (ast.In, ast.NotIn)) for o in ops): - if len(ops) != 1: - raise UnsupportedOperationException( - "`in` / `not in` inside a chained comparison (e.g. " - "`a < b in (1, 2)`) is not supported by the " - "transpiler; falling back to interpreted Python" - ) - return self._lower_in( - params, left, comps[0], negate=isinstance(ops[0], ast.NotIn) + raise UnsupportedOperationException( + "`in` / `not in` is not supported by the transpiler; the " + "UDF falls back to interpreted Python" ) - if len(ops) > _MAX_CHAIN_OPS: + if len(ops) > _MAX_COMPARISON_TERMS: raise UnsupportedOperationException( f"chained comparison with {len(ops)} operators exceeds " - f"the {_MAX_CHAIN_OPS} the transpiler will lower; the " + f"the {_MAX_COMPARISON_TERMS} the transpiler will lower; the " "UDF falls back to interpreted Python" ) # Python evaluates `a OP1 b OP2 c` as `(a OP1 b) and (b OP2 c)` # with `b` evaluated once and with SHORT-CIRCUIT: when - # `a OP1 b` is falsy, `b OP2 c` is never evaluated. That is not - # cosmetic here, because `_lower_value_compare` emits a - # `raise_error` for NULL operands to mirror Python's TypeError: - # `f(1, 0, None)` for `lambda a, b, c: a < b < c` must return - # False, not raise. `_lazy_and` explains why reproducing that - # needs nested CASE WHENs rather than `&`, and why every link - # being total (see `_lower_compare_pair`) removes the need for - # `coalesce`. + # `a OP1 b` is falsy, `b OP2 c` is never evaluated. That matters + # because `_lower_value_compare` emits a `raise_error` for NULL + # operands to mirror Python's TypeError: `f(1, 0, None)` for + # `lambda a, b, c: a < b < c` must return False, not raise. + # + # `&` short-circuits at evaluation time -- both `And.eval` and + # `And.doGenCode` skip the right operand when the left is false -- + # so the fold below reproduces that. What it does NOT survive is + # the optimizer: `splitConjunctivePredicates` decomposes `And`, + # and `PushPredicateThroughJoin` will push a conjunct it believes + # harmless, which a `raise_error` currently is because + # `RaiseError` does not override `Expression.throwable`. Fixing + # that override is the prerequisite change; see + # DEFERED_AND_OR_CASE_WHEN_ISSUE for the alternative (folding + # into nested CASE WHENs, which is atomic to predicate splitting + # but costs pushdown for every `and` / `or` UDF). # # Python's "evaluate `b` once" rule is about side effects. The # *Python-level* operand grammar we accept (Constant / Name / @@ -1068,13 +862,20 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # interpreted Python, and were therefore correct -- newly exposes # them to this bug. That argues for landing the analyzer guard # before or with this change. + # + # NaN is a separate known divergence: Spark orders NaN above every + # value and treats `NaN = NaN` as true, Python does neither, so a + # chain over doubles can disagree link by link. Tracked in + # SPARK-58781. nodes = [left, *comps] - return _lazy_and( - [ - self._lower_compare_pair(params, op, left_node, right_node) - for op, left_node, right_node in zip(ops, nodes, nodes[1:]) - ] - ) + links = [ + self._lower_compare_pair(params, op, left_node, right_node) + for op, left_node, right_node in zip(ops, nodes, nodes[1:]) + ] + result = links[0] + for link in links[1:]: + result = result & link + return result case ast.BinOp(left=left, op=op, right=right): # Operator selection is driven by the operand *categories* under # the current input-type variant (see ``_category``): Python's From 3f63fda570b0319f38262d9e7e2f413049833808 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Fri, 14 Aug 2026 23:36:13 +0000 Subject: [PATCH 04/18] [SPARK-58456][PYTHON][SQL] More closely match python behavior Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../sql/tests/test_udf_transpile_unit.py | 117 ++++++++---------- python/pyspark/sql/transpile.py | 108 ++++++++++------ 2 files changed, 118 insertions(+), 107 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index 54e1c6b50db3b..f35172b783fff 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -495,8 +495,8 @@ def gte_zero(x): def test_udf_transpile_chained_comparison_lowered(self): # ``0 < x < 10`` is a chained comparison: Python evaluates it as # ``(0 < x) and (x < 10)`` with ``x`` evaluated once. The transpiler - # folds the links into a nested CASE WHEN, so it transpiles and matches - # interpreted Python. A NULL input still raises through the ordering + # folds the links with `&`, so it transpiles and matches interpreted + # Python. A NULL input still raises through the ordering # guard on the first link (Python would raise TypeError), so the values # asserted here are the non-NULL ones. from pyspark.sql.types import StructField, StructType @@ -517,20 +517,23 @@ def chained(x): def test_udf_transpile_chain_short_circuits(self): # THE test for the chained lowering. Python's chain short-circuits: when # ``a < b`` is false, ``b < c`` is never evaluated, so a NULL ``c`` - # cannot raise. The transpiler must reproduce that -- which is why the - # links are folded into a nested CASE WHEN (whose branches Catalyst - # evaluates lazily and which predicate splitting cannot take apart) - # rather than `&`-ed together. + # cannot raise. The transpiler reproduces that by folding the links with + # `&`, which short-circuits at evaluation time. + # + # This covers evaluation only. `&` is splittable, so the guard's survival + # through the OPTIMIZER is a separate matter -- see + # `test_udf_transpile_boolop_short_circuits`, which pins that and depends + # on `RaiseError` overriding `Expression.throwable`. # # Note the literal form (`1 < 0 < None`) proves nothing here: a literal # None has no column category, so `_lower_value_compare` refuses the # whole UDF and it falls back. Nullable *columns* are required. # - # Run with whole-stage codegen both on and off: `CaseWhen` implements - # laziness twice over (`eval` walks branches until one is TRUE, and the - # single-branch codegen path defers to `If.doGenCode`, which emits each - # branch's code inside the generated `if`/`else`). A regression in - # either implementation must fail this test. + # Run with whole-stage codegen both on and off: `And` implements laziness + # twice over (`And.eval` skips the right operand when the left is false, + # and `And.doGenCode` emits the right operand's code inside the generated + # short-circuit). A regression in either implementation must fail this + # test. from pyspark.sql.types import StructField, StructType def chained(a, b, c): @@ -570,15 +573,19 @@ def test_udf_transpile_boolop_short_circuits(self): # guard in the second operand must not fire when the first already # decided the result. # - # NOT asserted here: that the guard stays whole through the optimizer. - # `splitConjunctivePredicates` decomposes `And`, and - # `PushPredicateThroughJoin` will push a conjunct it believes harmless, - # which a `raise_error` currently is because `RaiseError` does not - # override `Expression.throwable`. Overriding it is the prerequisite fix; - # see DEFERED_AND_OR_CASE_WHEN_ISSUE. Until then the hazard is real but - # its runtime symptom is scheduling-dependent (the pushed-down guard - # fires only when the failing task gets far enough before the join yields - # no rows), so there is nothing deterministic to assert. + # The join case below is a REGRESSION TEST FOR A PREREQUISITE and is RED + # until `RaiseError` overrides `Expression.throwable`. Without that + # override `splitConjunctivePredicates` decomposes the `And` and + # `PushPredicateThroughJoin` pushes a guard it believes harmless to one + # side of the join, where it fires on rows Python short-circuits past -- + # so the query raises USER_RAISED_EXCEPTION where Python returns False. + # + # AQE and broadcast joins are disabled for it deliberately. Under the + # default config whether the pushed-down guard fires is a race between + # the right-side filter task and AQE eliminating the empty join, so the + # same query intermittently returns [] instead of raising. A test that + # passes on some runs is worse than one that fails on all of them: pin + # the plan so this goes green exactly when the override lands. def both_positive(a, c): return a > 0 and c > 0 @@ -598,15 +605,21 @@ def either_positive(a, c): df = self.spark.createDataFrame([(5, None)], "a long, c long") self.assertEqual([True], [r[0] for r in df.select(or_udf("a", "c")).collect()]) - # Same thing across a join, where the NULL operand lives entirely on - # the right side. The `and` must still short-circuit rather than - # raising on the row whose `c` is NULL. + # Same thing across a join, where the NULL operand lives entirely on the + # right side. Deterministic plan: no AQE, no broadcast, so the pushed-down + # guard either exists or does not, rather than racing the join. + deterministic = dict(_TRANSPILE_ON) + deterministic["spark.sql.adaptive.enabled"] = False + deterministic["spark.sql.autoBroadcastJoinThreshold"] = -1 + with self.sql_conf(deterministic): + and_udf = UserDefinedFunction(both_positive, BooleanType()) left = self.spark.createDataFrame([(1, -5)], "k long, a long") right = self.spark.createDataFrame([(1, None)], "k long, c long") joined = left.join(right, "k") self.assertEqual([False], [r[0] for r in joined.select(and_udf("a", "c")).collect()]) - # Filtering on it returns Python's answer rather than raising. + # Filtering must return Python's answer rather than raising. RED + # until `RaiseError` overrides `Expression.throwable`; see above. filtered = joined.filter(and_udf("a", "c")) self.assertEqual([], filtered.select("k").collect()) @@ -763,10 +776,11 @@ def test_udf_transpile_boolean_body_requires_boolean_return_type(self): # `_category` now classifies every `ast.Compare` as "bool" directly. import warnings as _warnings - def ternary_in(c, x, y): - # `in` whose subject is a ternary -- reachable only since `in` is - # lowered at all, hence pinned here. - return (x if c > 0 else y) in (1, 2) + def ternary_compare(c, x, y): + # A comparison whose operand is itself a ternary, so + # `_is_definitely_boolean`'s basic-type requirement says False and + # the body relies on `_category` classifying `ast.Compare` as "bool". + return (x if c > 0 else y) == 1 def eq_true(a, b): # The same gap via `_lower_eq`; pre-dates `in` support. @@ -777,7 +791,7 @@ def chain_of_compare(a, b, c): return (a < b) <= (b < c) cases = [ - (ternary_in, "c long, x long, y long", [(1, 1, 9)]), + (ternary_compare, "c long, x long, y long", [(1, 1, 9)]), (eq_true, "a long, b long", [(1, 2)]), (chain_of_compare, "a long, b long, c long", [(1, 2, 3)]), ] @@ -802,8 +816,8 @@ def chain_of_compare(a, b, c): ) # Declared BooleanType, the same bodies transpile and are correct. df = self.spark.createDataFrame([(1, 1, 9)], "c long, x long, y long") - pudf = UserDefinedFunction(ternary_in, BooleanType()) - self.assertTrue(pudf.transpiled, "ternary_in should transpile for BooleanType") + pudf = UserDefinedFunction(ternary_compare, BooleanType()) + self.assertTrue(pudf.transpiled, "ternary_compare should transpile for BooleanType") self.assertEqual([True], [r[0] for r in df.select(pudf("c", "x", "y")).collect()]) def test_udf_transpile_lowering_caps(self): @@ -1345,13 +1359,11 @@ def wrapper(x): self.assertEqual(result[0], 60, "must run the wrapper, not the wrapped source") def test_udf_transpile_falls_back_for_none_in_boolop(self): - # Python's `None and x` short-circuits to None; neither Spark's - # three-valued `null AND false` (false) nor the nested CASE WHEN the - # transpiler emits (a NULL condition is not-taken, so also false) - # reproduces that. A literal None operand -- and, more generally, any - # operand `_is_never_null_boolean` cannot prove non-NULL, such as a - # ternary with a None branch -- must force a fallback rather than - # silently diverge. + # Python's `None and x` short-circuits to None, which Spark's + # three-valued `null AND false` (false) does not reproduce. A literal + # None operand -- and, more generally, any operand + # `_is_never_null_boolean` cannot prove non-NULL, such as a ternary with + # a None branch -- must force a fallback rather than silently diverge. import warnings as _warnings from pyspark.sql.types import StructField, StructType @@ -2289,19 +2301,6 @@ def test_udf_transpile_lowers_operators(self): assigned = lambda v: v + 1 # noqa: E731 bounded = lambda x: 0 <= x <= 100 # noqa: E731 str_chain = lambda x: "a" < x < "z" # noqa: E731 - in_tuple = lambda x: x in (1, 2, 3) # noqa: E731 - not_in_tuple = lambda x: x not in (1, 2) # noqa: E731 - in_negative = lambda x: x in (-1, 2) # noqa: E731 - in_list = lambda x: x in [1, 2] # noqa: E731 - in_set = lambda x: x in {1, 2} # noqa: E731 - in_dup = lambda x: x in (1, 1) # noqa: E731 - in_str = lambda x: x in ("a", "b") # noqa: E731 - # `not (x in ...)` is deliberately spelled with an explicit `not` rather - # than `not in`: it parses to UnaryOp(Not) over Compare(In), which is a - # different lowering path than Compare(NotIn) and is covered by - # `not_in_tuple` above. Hence the E713 suppression. - not_in_chain = lambda x: not (x in (1, 2)) # noqa: E713,E731 - in_and = lambda x, y: (x in (1, 2)) and (y > 0) # noqa: E731 chain_and = lambda a, b, c: (0 < a < 10) and (b == c) # noqa: E731 chain_ternary = lambda x: (0 < x < 10) if x is not None else None # noqa: E731 @@ -2313,12 +2312,6 @@ def if_elif_else(x): else: return 1 - def if_in(x): - if x in (1, 2): - return 10 - else: - return 20 - # (func, return_type, schema, rows, expected); arg columns come from the schema. cases = [ (modulo, L, "a long, b long", [(7, 3), (7, -3), (-7, 3), (-7, -3)], [1, -2, 2, -1]), @@ -2357,16 +2350,6 @@ def if_in(x): ), (chain_ternary, B, "a long", [(5,), (50,), (None,)], [True, False, None]), # `in` / `not in` over container literals. - (in_tuple, B, "a long", [(1,), (3,), (4,), (None,)], [True, True, False, False]), - (not_in_tuple, B, "a long", [(1,), (3,), (None,)], [False, True, True]), - (in_negative, B, "a long", [(-1,), (2,), (1,)], [True, True, False]), - (in_list, B, "a long", [(1,), (3,)], [True, False]), - (in_set, B, "a long", [(2,), (3,)], [True, False]), - (in_dup, B, "a long", [(1,), (2,)], [True, False]), - (in_str, B, "a string", [("a",), ("c",), (None,)], [True, False, False]), - (not_in_chain, B, "a long", [(1,), (3,), (None,)], [False, True, True]), - (in_and, B, "a long, b long", [(1, 5), (1, -5), (9, 5)], [True, False, False]), - (if_in, L, "a long", [(1,), (2,), (9,), (None,)], [10, 10, 20, 20]), ( nested, B, diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index b4a530116ee34..fdce74a02062d 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -245,6 +245,44 @@ def _is_definitely_boolean(node: ast.AST) -> bool: return False +def _is_never_null_boolean(node: ast.AST) -> bool: + """Return True when ``node`` lowers to a boolean column that is never NULL. + + Stricter than ``_is_definitely_boolean``, which also accepts operands that + may evaluate to NULL (a literal ``None``, a ternary with a ``None`` branch). + ``and`` / ``or`` need this even under the ``&`` / ``|`` fold, because Spark's + three-valued logic diverges from Python's short-circuit: ``NULL AND false`` + is ``false`` where Python's ``None and False`` is ``None``. Checking for a + bare ``None`` operand is not enough -- a ternary that *evaluates* to ``None`` + (``(True if c else None) and x``) reaches the fold the same way. + + ``not`` and if/ternary *tests* keep using ``_is_definitely_boolean``, since + they wrap the operand in ``coalesce`` and so handle NULL explicitly. + """ + match node: + case ast.Constant(value=v): + # A literal None is "boolean" for coalesce purposes but IS NULL. + return isinstance(v, bool) + case ast.Compare(): + # Every comparison lowering is total (see ``_lower_compare_pair``), + # so only the operand shapes still need checking. + return _is_definitely_boolean(node) + case ast.UnaryOp(op=ast.Not()): + # Lowered as ``coalesce(~operand, lit(True))``, which is total. + return True + case ast.BoolOp(values=values): + return all(_is_never_null_boolean(v) for v in values) + case ast.IfExp(body=body, orelse=orelse): + return _is_never_null_boolean(body) and _is_never_null_boolean(orelse) + case _: + return False + + +def _is_none_constant(node: ast.AST) -> bool: + """Return True for the literal ``None``.""" + return isinstance(node, ast.Constant) and node.value is None + + def _truthiness_col(cat: Optional[str], c: Column) -> Optional[Column]: """Return a boolean Column expressing Python's ``bool(c)`` for the given category. @@ -265,11 +303,6 @@ def _truthiness_col(cat: Optional[str], c: Column) -> Optional[Column]: return None -def _is_none_constant(node: ast.AST) -> bool: - """Return True for the literal ``None``.""" - return isinstance(node, ast.Constant) and node.value is None - - class CatalystTranspiler(AbstractTranspiler): """Transpiler that attempts to convert a Python UDF into native Spark SQL expressions.""" @@ -728,34 +761,26 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # Python `and` / `or` short-circuit and return one of the # operands rather than a strict boolean. For the booleans # produced by Compare / UnaryOp(Not) / nested BoolOps that maps - # cleanly onto a nested CASE WHEN. For non-boolean operands - # (including bare parameter names whose runtime type is unknown) - # the right semantics would need Python's truthiness rules - # (0 / "" / None / [] all falsy), which we can't reproduce - # without the input column types -- Spark's `&` / `|` would do - # bitwise instead. Require statically known booleans so the - # caller falls back rather than emitting a divergent plan. + # cleanly onto `&` / `|`. For non-boolean operands (including + # bare parameter names whose runtime type is unknown) the right + # semantics would need Python's truthiness rules (0 / "" / None / + # [] all falsy), which we can't reproduce without the input + # column types -- Spark's `&` / `|` would do bitwise instead. # - if not all(_is_definitely_boolean(v) for v in values): + # The gate is `_is_never_null_boolean`, not + # `_is_definitely_boolean`: Spark's three-valued `NULL AND false` + # is `false` where Python's `None and False` is `None`, so a + # NULL-valued operand diverges. Checking only for a bare `None` + # operand is not enough, because a ternary that evaluates to None + # (`(True if c else None) and x`) reaches the fold the same way. + if not all(_is_never_null_boolean(v) for v in values): raise UnsupportedOperationException( - "`and` / `or` operand type is not statically known " - "to be boolean; Spark's `&` / `|` are bitwise, not " - "Python truthiness, so the transpiler refuses to " - "lower this and the UDF falls back to interpreted " - "Python" - ) - # A literal None operand short-circuits differently: Python's - # `None and (x > 0)` returns None regardless of x, but Spark's - # three-valued `null AND false` is false (and `null OR true` is - # true), so the lowered form diverges. `_is_definitely_boolean` - # accepts None for `not`/if-test contexts where coalesce handles - # it; here it must force a fallback instead. - if any(_is_none_constant(v) for v in values): - raise UnsupportedOperationException( - "literal None operand in `and` / `or` cannot be lowered: " - "Spark's three-valued logic diverges from Python's " - "short-circuit-return-operand semantics, so the UDF " - "falls back to interpreted Python" + "`and` / `or` operand is not statically known to be a " + "non-NULL boolean; Python's short-circuit returns the " + "operand itself (so `None and x` is None) while Spark's " + "`&` / `|` are bitwise over three-valued logic, so the " + "transpiler refuses to lower this and the UDF falls back " + "to interpreted Python" ) cols = [self._convert_chunk(params, v) for v in values] if isinstance(op, ast.And): @@ -818,15 +843,18 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # # `&` short-circuits at evaluation time -- both `And.eval` and # `And.doGenCode` skip the right operand when the left is false -- - # so the fold below reproduces that. What it does NOT survive is - # the optimizer: `splitConjunctivePredicates` decomposes `And`, - # and `PushPredicateThroughJoin` will push a conjunct it believes - # harmless, which a `raise_error` currently is because - # `RaiseError` does not override `Expression.throwable`. Fixing - # that override is the prerequisite change; see - # DEFERED_AND_OR_CASE_WHEN_ISSUE for the alternative (folding - # into nested CASE WHENs, which is atomic to predicate splitting - # but costs pushdown for every `and` / `or` UDF). + # so the fold below reproduces that. What it does NOT survive on + # its own is the optimizer: `splitConjunctivePredicates` + # decomposes `And`, and `PushPredicateThroughJoin` pushes a + # conjunct it believes harmless, which a `raise_error` is until + # `RaiseError` overrides `Expression.throwable` (as `Sequence` + # already does). That override is a prerequisite for this + # lowering: without it, a chain under a filter over a join can + # have a link's guard pushed below the join and raise on rows + # Python never evaluates. The alternative -- folding into nested + # CASE WHENs, atomic to predicate splitting -- was rejected + # because it costs pushdown for every `and` / `or` UDF, including + # those over non-nullable columns that need no guard at all. # # Python's "evaluate `b` once" rule is about side effects. The # *Python-level* operand grammar we accept (Constant / Name / From 0d7c2dac1c2829a4a41523649516bb778a974516 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 15 Aug 2026 05:46:05 +0000 Subject: [PATCH 05/18] [SPARK-58456][PYTHON][SQL] Trim the comment prose and consolidate redundant tests Cuts the explanatory prose to what a reviewer actually needs and folds two redundant tests into ones that already cover them. Net -220 lines, no intended behavior change. - Fold `and` / `or` operands and chain links with `functools.reduce` rather than hand-rolled accumulate loops. - Rename `_MAX_COMPARISON_TERMS` to `_MAX_COMPARISON_OPS`, which is what it counts (operators, not terms). - Drop `test_udf_transpile_chained_comparison_lowered` and `test_chained_ge_matches_python`. Both were added by this branch and are subsumed: chain values and the `a == b == c` NULL patterns are covered by `test_udf_transpile_chain_operator_mix_and_eq_nulls`, `0 < x < 10` by `test_udf_transpile_plan_elision`, and `>=` in chain position by the `a > b >= c` case, so no operator loses coverage. Two fixes that are not prose: - Drop the stale `x in (1, 2, 3)` case from `test_udf_transpile_plan_elision`. Narrowing this branch to chained comparisons made `_convert_chunk` refuse `in` / `not in`, so the case asserted a UDF would transpile that no longer can, and it was failing in CI. - Exercise `And`'s two laziness implementations with `spark.sql.codegen.factoryMode` FALLBACK / NO_CODEGEN instead of toggling `spark.sql.codegen.wholeStage`. Only `NO_CODEGEN` skips codegen and takes the interpreted path (SQLConf.scala), so both legs were previously reaching `And.doGenCode` and `And.eval` was never covered. Still deliberately red, unchanged by this commit: the `RaiseError.throwable` prerequisite (SPARK-58627) and the analyzer determinism guard. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../tests/test_udf_transpile_hypothesis.py | 104 ++--- .../sql/tests/test_udf_transpile_unit.py | 359 +++++++----------- python/pyspark/sql/transpile.py | 246 +++++------- ...solveTranspiledPythonUDFOptionsSuite.scala | 21 +- 4 files changed, 255 insertions(+), 475 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py index bffa694fdbda3..ffaf2377c924a 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -206,13 +206,9 @@ def _env_opts_in(value: Optional[str]) -> bool: (1, -_LONG_BOUND), ) - # Triple edges for the chained-comparison tests. The eight NULL - # combinations are the point: a chain short-circuits, so whether a NULL - # operand's raise guard is reached depends on the earlier links. The - # ordering chains raise in pure Python on any NULL they reach (covered by - # ``_run``'s "both raised" equivalence), while ``x == y == z`` has a - # defined answer for every combination because Python's ``==`` never - # raises on None. + # Triple edges for the chained-comparison tests. The NULL patterns are the + # point: a chain short-circuits, so whether a NULL operand's raise guard is + # reached depends on the earlier links. _LONG_TRIPLE_EDGES = ( (None, None, None), (None, None, 0), @@ -243,19 +239,13 @@ def wrapper(method): def _seed_tuple_examples(tuples, keys=("x", "y", "z")): """Stack one ``@example`` decorator per seed tuple, zipped onto ``keys``. - The arity check is load-bearing: ``zip`` stops at the shorter argument, - so a seed tuple whose length does not match ``keys`` would silently - register an ``@example`` for a different input than intended and quietly - shrink the seeded NULL-combination coverage. Fail loudly instead. + The arity assert matters: ``zip`` stops at the shorter argument, so a + mismatched seed tuple would quietly seed a different input than intended. """ def wrapper(method): for values in reversed(tuples): - if len(values) != len(keys): - raise ValueError( - f"seed tuple {values!r} has {len(values)} value(s) but " - f"{len(keys)} key(s) {keys!r}" - ) + assert len(values) == len(keys), f"seed {values!r} does not fit {keys!r}" method = example(**dict(zip(keys, values)))(method) return method @@ -272,6 +262,12 @@ def _seed_pair_examples(pairs, keys=("x", "y")): # (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. +# +# Where a template annotates its parameters, that is to pin the input category: +# an untyped parameter emits both a numeric and a string variant, and these +# tests bind Long columns, so the string one is dead weight. It buys plan size, +# not wall clock. The strategies still generate ``None`` -- an annotation pins +# the category, not nullability. def plus_four(x): @@ -386,27 +382,9 @@ def neq_pair(x, y): def chained_bounds(x: int) -> bool: # Single-parameter chained comparison: ``ast.Compare`` with two ops, which - # Python evaluates as ``(0 < x) and (x < 10)`` with ``x`` evaluated once. - # No None guard, so a NULL input reaches the first link's raise guard -- - # matching Python's TypeError -- and ``_run``'s "both raised" equivalence - # covers it. - # - # The parameters here and in the chains below are ANNOTATED deliberately. - # An untyped parameter is transpiled as both numeric and string, so the - # option matrix carries variants these tests can never use: they bind Long - # columns, so only the all-numeric variant can match. Pinning the category - # emits one option instead of two (measured: 1 vs 2 -- ``_param_category_ - # combos`` already collapses three untyped params to all-numeric plus - # all-string rather than 2**3). - # - # This costs no coverage, but do not expect it to speed the suite up: the - # per-example cost is dominated by the two Spark jobs ``_run`` executes - # against a fresh one-row DataFrame, not by transpilation. Measured over - # these five tests at 15 examples: 73s untyped vs 76s typed, i.e. no - # improvement. Batching rows per example is the lever that would matter. - # - # The strategies still generate ``None``; the annotation pins the category, - # not nullability. + # Python evaluates as ``(0 < x) and (x < 10)`` with ``x`` evaluated once. No + # None guard, so a NULL input reaches the first link's raise guard, matching + # Python's TypeError; ``_run``'s "both raised" equivalence covers that. return 0 < x < 10 @@ -418,24 +396,16 @@ def chained_lt(x: int, y: int, z: int) -> bool: return x < y < z -def chained_ge(x: int, y: int, z: int) -> bool: - # Reversed-direction chain, so the short-circuit fires on a different set - # of rows than ``chained_lt``. - return x >= y >= z - - def chained_eq(x: int, y: int, z: int) -> bool: - # ``x == y == z`` -- equality never raises in Python, so every one of the - # eight NULL combinations has a defined boolean answer that the - # four-branch ``_lower_eq`` per link plus the short-circuit fold must - # reproduce exactly. This is the strongest NULL-semantics probe of the - # chained lowering. + # Equality never raises in Python, so every NULL pattern has a defined + # answer -- the strongest NULL probe of the chained lowering, with no "both + # raised" escape hatch to hide a divergence. return x == y == z def mixed_chain(x: int, y: int, z: int) -> bool: - # Mixed operators in one chain (``ast.LtE`` then ``ast.Gt``), so the two - # links go through different lowering helpers. + # Mixed operators in one chain (``ast.LtE`` then ``ast.Gt``), so the links + # disagree on direction and short-circuit on different rows again. return x <= y > z @@ -536,9 +506,9 @@ def _run(self, func, return_type, df, *udf_arg_columns, kwargs=None): 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. + # Asymmetric on purpose. If the transpiled path raises, the interpreted path + # must raise too -- that is the dangerous direction. The reverse is allowed: + # interpreted Python may raise where the transpiled form returns a value. if transpiled_error is not None: self.assertIsNotNone( interpreted_error, @@ -772,14 +742,11 @@ def test_lambda_plus_four_matches_python(self, value): # ---- Chained comparisons ---------------------------------------- # - # Python's ``a OP b OP c`` short-circuits: when the first link is falsy - # the second is never evaluated, so a NULL operand it would have - # touched must not raise. ``_run`` catches the dangerous direction of a - # regression here -- if the transpiled form raises where interpreted - # Python does not, ``_run`` asserts on it. (The reverse direction, - # Python raising where the transpiled form returns a value, is - # tolerated by ``_run``, so it is pinned by explicit assertions in - # test_udf_transpile_unit.py instead.) + # Python's ``a OP b OP c`` short-circuits, so a NULL operand the second + # link would have touched must not raise. ``_run`` catches the dangerous + # direction: the transpiled form raising where interpreted Python does + # not. It tolerates the reverse, which test_udf_transpile_unit.py pins + # explicitly instead. @_hyp_settings @given(value=_long_strategy) @@ -799,23 +766,10 @@ def test_chained_lt_matches_python(self, x, y, z): transpiled, interpreted, f"chained_lt mismatch on ({x!r}, {y!r}, {z!r})" ) - @_hyp_settings - @given(x=_long_strategy, y=_long_strategy, z=_long_strategy) - @_seed_tuple_examples(_LONG_TRIPLE_EDGES) - def test_chained_ge_matches_python(self, x, y, z): - df = self._three_long_arg_df(x, y, z) - transpiled, interpreted = self._run(chained_ge, BooleanType(), df, "a", "b", "c") - self.assertEqual( - transpiled, interpreted, f"chained_ge mismatch on ({x!r}, {y!r}, {z!r})" - ) - @_hyp_settings @given(x=_long_strategy, y=_long_strategy, z=_long_strategy) @_seed_tuple_examples(_LONG_TRIPLE_EDGES) def test_chained_eq_matches_python(self, x, y, z): - # Equality never raises in Python, so every NULL combination has a - # defined boolean answer on both sides -- no "both raised" - # escape hatch can mask a divergence here. df = self._three_long_arg_df(x, y, z) transpiled, interpreted = self._run(chained_eq, BooleanType(), df, "a", "b", "c") self.assertEqual( @@ -832,8 +786,6 @@ def test_mixed_chain_matches_python(self, x, y, z): transpiled, interpreted, f"mixed_chain mismatch on ({x!r}, {y!r}, {z!r})" ) - # ---- `in` / `not in` --------------------------------------------- - class UDFTranspileHypothesisGatingTests(unittest.TestCase): """Smoke tests that always run, regardless of the env gate. diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index f35172b783fff..b7b5855a6ca3a 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -270,9 +270,9 @@ def func_closure_capture(x): def test_udf_transpile_boolean_and_or_lowered(self): # When `and`/`or` operands are syntactically boolean (Compare - # results in this case), the transpiler should lower to a nested - # CASE WHEN that reproduces Python's short-circuit and produce - # results matching the interpreted UDF. + # results in this case), the transpiler should lower to `&`/`|`, + # which short-circuit like Python, and produce results matching + # the interpreted UDF. # Each UDF is a single top-level statement (the transpiler # doesn't support multi-statement bodies yet). from pyspark.sql.types import StructField, StructType @@ -492,67 +492,30 @@ def gte_zero(x): self.assertEqual(row_lte[0], lte_expected) self.assertEqual(row_gte[0], gte_expected) - def test_udf_transpile_chained_comparison_lowered(self): - # ``0 < x < 10`` is a chained comparison: Python evaluates it as - # ``(0 < x) and (x < 10)`` with ``x`` evaluated once. The transpiler - # folds the links with `&`, so it transpiles and matches interpreted - # Python. A NULL input still raises through the ordering - # guard on the first link (Python would raise TypeError), so the values - # asserted here are the non-NULL ones. - from pyspark.sql.types import StructField, StructType - - def chained(x): - return 0 < x < 10 - - schema = StructType([StructField("a", LongType(), nullable=False)]) - with self.sql_conf(_TRANSPILE_ON): - pudf = UserDefinedFunction(chained, BooleanType()) - self.assertTrue(pudf.transpiled, "chained comparison should transpile") - for value, expected in [(5, True), (0, False), (10, False), (-3, False), (9, True)]: - with self.subTest(value=value): - df = self.spark.createDataFrame([Row(a=value)], schema=schema) - [row] = df.select(pudf("a")).collect() - self.assertEqual(row[0], expected) - def test_udf_transpile_chain_short_circuits(self): # THE test for the chained lowering. Python's chain short-circuits: when - # ``a < b`` is false, ``b < c`` is never evaluated, so a NULL ``c`` - # cannot raise. The transpiler reproduces that by folding the links with - # `&`, which short-circuits at evaluation time. - # - # This covers evaluation only. `&` is splittable, so the guard's survival - # through the OPTIMIZER is a separate matter -- see - # `test_udf_transpile_boolop_short_circuits`, which pins that and depends - # on `RaiseError` overriding `Expression.throwable`. + # ``a < b`` is false, ``b < c`` is never evaluated, so a NULL ``c`` cannot + # raise. `&` short-circuits at evaluation time too, so the fold matches. + # Nullable *columns* are required: a literal `None` has no column + # category, so `1 < 0 < None` just falls back. # - # Note the literal form (`1 < 0 < None`) proves nothing here: a literal - # None has no column category, so `_lower_value_compare` refuses the - # whole UDF and it falls back. Nullable *columns* are required. - # - # Run with whole-stage codegen both on and off: `And` implements laziness - # twice over (`And.eval` skips the right operand when the left is false, - # and `And.doGenCode` emits the right operand's code inside the generated - # short-circuit). A regression in either implementation must fail this - # test. - from pyspark.sql.types import StructField, StructType - + # Both of `And`'s two laziness implementations are exercised: `doGenCode` + # under the default FALLBACK factory mode, and `eval` under NO_CODEGEN. + # `spark.sql.codegen.wholeStage=false` is NOT enough for the latter -- it + # disables whole-stage fusion, but expressions are still compiled through + # the codegen factory, so `doGenCode` runs either way. Survival through the + # OPTIMIZER is separate; see test_udf_transpile_boolop_short_circuits. def chained(a, b, c): return a < b < c - schema = StructType( - [ - StructField("a", LongType(), nullable=True), - StructField("b", LongType(), nullable=True), - StructField("c", LongType(), nullable=True), - ] - ) - for wholestage in (True, False): + schema = "a long, b long, c long" + for factory_mode in ("FALLBACK", "NO_CODEGEN"): conf = dict(_TRANSPILE_ON) - conf["spark.sql.codegen.wholeStage"] = wholestage - with self.subTest(wholestage=wholestage), self.sql_conf(conf): + conf["spark.sql.codegen.factoryMode"] = factory_mode + with self.subTest(factory_mode=factory_mode), self.sql_conf(conf): pudf = UserDefinedFunction(chained, BooleanType()) self.assertTrue(pudf.transpiled, "3-column chain should transpile") - df = self.spark.createDataFrame([Row(a=1, b=0, c=None)], schema=schema) + df = self.spark.createDataFrame([(1, 0, None)], schema) [row] = df.select(pudf("a", "b", "c")).collect() self.assertEqual( False, @@ -561,31 +524,26 @@ def chained(a, b, c): ) # ... but when the first link IS true the second link is reached # and its NULL guard fires, exactly as Python raises TypeError. - df = self.spark.createDataFrame([Row(a=0, b=1, c=None)], schema=schema) + df = self.spark.createDataFrame([(0, 1, None)], schema) with self.assertRaises(Exception) as ctx: df.select(pudf("a", "b", "c")).collect() self.assertIn("cannot compare NULL", str(ctx.exception)) def test_udf_transpile_boolop_short_circuits(self): - # The `and` / `or` counterpart of the chain short-circuit. `&` short- - # circuits at evaluation time -- both `And.eval` and `And.doGenCode` skip - # the right operand when the left is false -- so a `raise_error` NULL - # guard in the second operand must not fire when the first already - # decided the result. + # The `and` / `or` counterpart of the chain short-circuit: a NULL guard in + # the second operand must not fire when the first already decided the + # result. # - # The join case below is a REGRESSION TEST FOR A PREREQUISITE and is RED - # until `RaiseError` overrides `Expression.throwable`. Without that - # override `splitConjunctivePredicates` decomposes the `And` and - # `PushPredicateThroughJoin` pushes a guard it believes harmless to one - # side of the join, where it fires on rows Python short-circuits past -- - # so the query raises USER_RAISED_EXCEPTION where Python returns False. + # RED UNTIL `RaiseError` overrides `Expression.throwable` -- deliberately, + # to pin the prerequisite. Without the override, + # `splitConjunctivePredicates` decomposes the `And` and + # `PushPredicateThroughJoin` pushes the guard to one side of the join, + # where it fires on rows Python short-circuits past and raises + # USER_RAISED_EXCEPTION instead of returning False. # - # AQE and broadcast joins are disabled for it deliberately. Under the - # default config whether the pushed-down guard fires is a race between - # the right-side filter task and AQE eliminating the empty join, so the - # same query intermittently returns [] instead of raising. A test that - # passes on some runs is worse than one that fails on all of them: pin - # the plan so this goes green exactly when the override lands. + # The join case disables AQE and broadcast on purpose: by default whether + # the pushed-down guard fires races AQE eliminating the empty join, so the + # failure would be intermittent rather than a clean red. def both_positive(a, c): return a > 0 and c > 0 @@ -624,24 +582,13 @@ def either_positive(a, c): self.assertEqual([], filtered.select("k").collect()) def test_udf_transpile_chain_operator_mix_and_eq_nulls(self): - # Mixed operators in one chain, and `x == y == z` across every NULL - # combination. Equality never raises in Python (`None == None` is True, - # `None == 0` is False), so all eight combinations have a defined answer - # and the transpiled result must match it exactly. - from pyspark.sql.types import StructField, StructType - - schema = StructType( - [ - StructField("a", LongType(), nullable=True), - StructField("b", LongType(), nullable=True), - StructField("c", LongType(), nullable=True), - ] - ) - - def chained_eq(a, b, c): - return a == b == c + # Mixed operators in one chain, and `x == y == z` over every NULL pattern. + # Equality never raises in Python (`None == None` is True, `None == 0` is + # False), so every pattern has a defined answer to match exactly. + B = BooleanType() + chained_eq = lambda a, b, c: a == b == c # noqa: E731 - # (a, b, c, expected) -- the eight NULL combinations plus value cases. + # (a, b, c, expected) -- every NULL pattern plus value cases. eq_cases = [ (None, None, None, True), (None, None, 1, False), @@ -655,17 +602,13 @@ def chained_eq(a, b, c): (1, 1, 2, False), (1, 2, 3, False), ] - with self.sql_conf(_TRANSPILE_ON): - pudf = UserDefinedFunction(chained_eq, BooleanType()) - self.assertTrue(pudf.transpiled, "a == b == c should transpile") - for a, b, c, expected in eq_cases: - with self.subTest(a=a, b=b, c=c): - df = self.spark.createDataFrame([Row(a=a, b=b, c=c)], schema=schema) - [row] = df.select(pudf("a", "b", "c")).collect() - self.assertEqual(row[0], expected, f"a == b == c on ({a}, {b}, {c})") + self.assertEqual( + self._vals(chained_eq, B, "a long, b long, c long", [c[:3] for c in eq_cases]), + [c[3] for c in eq_cases], + "a == b == c over the NULL patterns", + ) # Operator mixes, on non-NULL rows so the ordering guards stay quiet. - B = BooleanType() lt_eq = lambda a, b, c: a < b == c # noqa: E731 eq_lt = lambda a, b, c: a == b < c # noqa: E731 gt_gte = lambda a, b, c: a > b >= c # noqa: E731 @@ -693,25 +636,22 @@ def chained_eq(a, b, c): ) def test_udf_transpile_comparison_fallbacks(self): - # Comparison forms the transpiler must REFUSE, each for its own reason. - # Every one asserts both that nothing transpiled and that a query using - # the UDF still runs and returns Python's answer -- the second half - # matters because a surviving option is type-checked by CheckAnalysis as - # a child of TranspiledPythonUDF, so a bad lowering fails the whole - # query instead of falling back (e.g. `x in (1, "a")` would raise - # DATA_DIFF_TYPES). + # Comparison forms the transpiler must REFUSE. Each case asserts both that + # nothing transpiled and that a query using the UDF still returns Python's + # answer. The second half matters because an option that survives is + # type-checked by CheckAnalysis as a child of TranspiledPythonUDF: a + # half-refused form would fail the whole query rather than fall back, so + # "did not transpile" alone is not enough to prove the fallback works. import warnings as _warnings - # `in` on a str/bytes is substring/subsequence containment, not - # equality; dict membership tests keys; a container of runtime values - # would need runtime NULL analysis; `in` cannot be chained. + # `in` / `not in` are refused wholesale (SPARK-56925), so these cases all + # take the same branch; they are kept because the reasons differ. `in` on a + # str is substring containment rather than equality, a container of runtime + # values would need runtime NULL analysis, and `in` can appear as a link in + # a chain (`a < b in (1, 2)` is one Compare), which the refusal also covers. str_in = lambda x: x in "abc" # noqa: E731 - bytes_in = lambda x: x in b"ab" # noqa: E731 - dict_in = lambda x: x in {1: "a"} # noqa: E731 runtime_in = lambda x, y: x in (y, 2) # noqa: E731 chained_in = lambda a, b: a < b in (1, 2) # noqa: E731 - mixed_in = lambda x: x in (1, "a") # noqa: E731 - bool_in = lambda x: x in (True, False) # noqa: E731 cross_chain = lambda x: 0 < x < "z" # noqa: E731 chain_long_rt = lambda x: 0 < x < 10 # noqa: E731 in_long_rt = lambda x: x in (1, 2) # noqa: E731 @@ -720,12 +660,8 @@ def test_udf_transpile_comparison_fallbacks(self): # Python produces, so the query must still run and return it. cases = [ ("`in` on a str", str_in, "a string", [("b",), ("z",)], [True, False]), - ("`in` on bytes", bytes_in, "a long", [(97,), (1,)], [True, False]), - ("`in` on a dict", dict_in, "a long", [(1,), (2,)], [True, False]), ("runtime `in`", runtime_in, "a long, b long", [(1, 1), (3, 5)], [True, False]), ("chain + `in`", chained_in, "a long, b long", [(0, 1), (1, 0)], [True, False]), - ("mixed-category `in`", mixed_in, "a long", [(1,), (3,)], [True, False]), - ("`in` over bools", bool_in, "a long", [(1,), (5,)], [True, False]), ] with self.sql_conf(_TRANSPILE_ON): for label, func, schema, rows, expected in cases: @@ -740,12 +676,10 @@ def test_udf_transpile_comparison_fallbacks(self): expected, f"{label}: interpreted result diverged", ) - # `0 < x < "z"` raises in interpreted Python too (cross-type - # ordering), so assert only the fallback for it. - # - # A chained / `in` body is boolean, so a non-BooleanType declared - # return type must drop every variant rather than emit a cast the - # interpreted path would never perform. + # `0 < x < "z"` raises in interpreted Python too (cross-type ordering), + # so assert only the fallback for it. And a chained / `in` body is + # boolean, so a non-BooleanType declared return type must drop every + # variant rather than emit a cast the interpreted path never performs. for func, rt, label in [ (cross_chain, BooleanType(), "a cross-category chain"), (chain_long_rt, LongType(), "a chain with a LongType return type"), @@ -761,33 +695,25 @@ def test_udf_transpile_comparison_fallbacks(self): ) def test_udf_transpile_boolean_body_requires_boolean_return_type(self): - # A comparison body is boolean no matter what its operands look like, so - # it must not satisfy a numeric declared return type. - # - # `_is_definitely_boolean`'s Compare arm additionally requires basic-type - # operands, so it answers False for a comparison whose operand is itself - # a ternary / comparison / boolean op. Those used to fall through - # `_category`'s numeric catch-all, which let the return-type match in - # `_transpile_from_ast` accept `LongType`: the lowered option returned - # `cast( as bigint)` -- 1 or 0 -- while the interpreted - # SQL_BATCHED_UDF path returns NULL for a Python bool declared LongType - # (EvaluatePython.makeFromJava nulls type-mismatched results). That is a - # silent divergence, exactly what the return-type match exists to stop. - # `_category` now classifies every `ast.Compare` as "bool" directly. + # A comparison body is boolean whatever its operands look like, so it must + # not satisfy a numeric declared return type. `_is_definitely_boolean` says + # False for a comparison over a ternary / comparison / boolean op (it also + # demands basic-type operands), and those used to fall through `_category`'s + # numeric catch-all: `_transpile_from_ast` then accepted `LongType` and the + # option returned `cast( as bigint)` -- 1 or 0 -- where the + # interpreted path returns NULL. `_category` now classifies every + # `ast.Compare` as "bool" directly. import warnings as _warnings def ternary_compare(c, x, y): - # A comparison whose operand is itself a ternary, so - # `_is_definitely_boolean`'s basic-type requirement says False and - # the body relies on `_category` classifying `ast.Compare` as "bool". return (x if c > 0 else y) == 1 def eq_true(a, b): - # The same gap via `_lower_eq`; pre-dates `in` support. + # The same gap via `_lower_eq`; pre-dates chained comparisons. return (a < b) == True # noqa: E712 def chain_of_compare(a, b, c): - # ... and via a chained comparison over a comparison operand. + # ... and via a chained comparison over comparison operands. return (a < b) <= (b < c) cases = [ @@ -821,17 +747,15 @@ def chain_of_compare(a, b, c): self.assertEqual([True], [r[0] for r in df.select(pudf("c", "x", "y")).collect()]) def test_udf_transpile_lowering_caps(self): - # The chain cap bounds plan growth: an interior chain operand's subtree is - # emitted once per adjacent pair, and every input-category variant carries - # its own copy for CheckAnalysis to walk. Exercise the boundary directly - # rather than through a UDF, so the refusal cannot be confused with a - # fallback for some other reason. + # The chain cap bounds plan growth. Exercise the boundary directly rather + # than through a UDF, so the refusal cannot be confused with a fallback for + # some other reason. import ast as _ast from pyspark.errors import UnsupportedOperationException from pyspark.sql.transpile import ( CatalystTranspiler, - _MAX_COMPARISON_TERMS, + _MAX_COMPARISON_OPS, ) transpiler = CatalystTranspiler() @@ -846,11 +770,11 @@ def chain_src(n): with self.sql_conf(_TRANSPILE_ON): # At the cap the lowering succeeds. - transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_TERMS))) + transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_OPS))) # One past it, it refuses so the UDF falls back. with self.assertRaises(UnsupportedOperationException) as ctx: - transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_TERMS + 1))) + transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_OPS + 1))) self.assertIn("operators", str(ctx.exception)) # `in` / `not in` are refused outright, deferred to SPARK-56925. @@ -1359,11 +1283,10 @@ def wrapper(x): self.assertEqual(result[0], 60, "must run the wrapper, not the wrapped source") def test_udf_transpile_falls_back_for_none_in_boolop(self): - # Python's `None and x` short-circuits to None, which Spark's - # three-valued `null AND false` (false) does not reproduce. A literal - # None operand -- and, more generally, any operand - # `_is_never_null_boolean` cannot prove non-NULL, such as a ternary with - # a None branch -- must force a fallback rather than silently diverge. + # Python's `None and x` short-circuits to None, which Spark's three-valued + # `null AND false` (false) does not reproduce. Any operand + # `_is_never_null_boolean` cannot prove non-NULL -- a literal None, or a + # ternary with a None branch -- must fall back rather than diverge. import warnings as _warnings from pyspark.sql.types import StructField, StructType @@ -2338,9 +2261,9 @@ def if_elif_else(x): [True, True, True, False, False], ), (str_chain, B, "a string", [("m",), ("a",), ("z",)], [True, False, False]), - # A chain as a boolean operand and inside a ternary: the fold's - # result is a non-NULL boolean, so it is accepted by the `and`/`or` - # gate and unifies with a None ternary branch. + # A chain as a boolean operand and inside a ternary: the fold's result + # is a non-NULL boolean, so the `and`/`or` gate accepts it and it + # unifies with a None ternary branch. ( chain_and, B, @@ -2349,7 +2272,6 @@ def if_elif_else(x): [True, False, False], ), (chain_ternary, B, "a long", [(5,), (50,), (None,)], [True, False, None]), - # `in` / `not in` over container literals. ( nested, B, @@ -2403,79 +2325,65 @@ def test_udf_transpile_plan_elision(self): self.assertEqual(chained.first()[0], 7.0) # ((10 + 3) + 1) / 2 self.assertEqual(2, self._eval_python_count(chained)) - # A chained comparison and an `in` are both fully lowered, so the - # EvalPython node disappears for them too. + # A chained comparison is fully lowered, so the EvalPython node + # disappears for it too. bounded = lambda x: 0 < x < 10 # noqa: E731 - member = lambda x: x in (1, 2, 3) # noqa: E731 - for label, func, expected in [ - ("0 < x < 10", bounded, [1, 7]), - ("x in (1,2,3)", member, [1]), - ]: - with self.subTest(func=label): - f = UserDefinedFunction(func, BooleanType()) - self.assertTrue(f.transpiled) - fdf = self.spark.createDataFrame([(7,), (11,), (1,)], "a long").filter(f("a")) - self.assertEqual(sorted(r[0] for r in fdf.collect()), expected) - self.assertEqual(0, self._eval_python_count(fdf)) + f = UserDefinedFunction(bounded, BooleanType()) + self.assertTrue(f.transpiled) + fdf = self.spark.createDataFrame([(7,), (11,), (1,)], "a long").filter(f("a")) + self.assertEqual(sorted(r[0] for r in fdf.collect()), [1, 7]) + self.assertEqual(0, self._eval_python_count(fdf)) def test_udf_transpile_nondeterministic_argument_falls_back(self): - # EXPECTED TO FAIL until the companion analyzer fix lands. This test is - # deliberately red: it pins a real wrong-results bug so the gap stays - # visible rather than being silently tolerated. Do not skip or xfail it. + # RED UNTIL the companion analyzer fix lands -- deliberately, to pin a real + # wrong-results bug rather than tolerate it silently. Do not skip or xfail. # - # A non-deterministic ARGUMENT expression must suppress transpilation at - # the call site, even though the UDF itself transpiled fine. (Distinct - # from test_udf_transpile_skips_nondeterministic, which covers a UDF - # marked non-deterministic.) + # A non-deterministic ARGUMENT must suppress transpilation at the call site + # even though the UDF itself transpiled fine. (Distinct from + # test_udf_transpile_skips_nondeterministic, which covers a UDF marked + # non-deterministic.) `resolveUDFParams` substitutes the argument at every + # `_udf_param_N` occurrence -- four of them for a two-link chain -- so a + # non-deterministic one becomes several independent expressions with + # independent state. Since `and` / `or` and chains evaluate later copies + # only on the rows where earlier links held, the copies desynchronize: + # `50 < x < 60` over `rand(7) * 100` is ~30% true instead of ~10%. # - # `resolveUDFParams` substitutes the argument expression at every - # `_udf_param_N` occurrence, and the lowerings reference an operand two - # to three times -- four for a two-link chain. For a deterministic - # argument that is only plan bloat, but a non-deterministic one becomes - # several independent expressions with independent internal state. Since - # `and` / `or` and chained comparisons evaluate later copies only on the - # rows where earlier links held, the copies advance at different rates - # and stop describing the same value: `50 < x < 60` over `rand(7) * 100` - # returns ~30% true instead of the correct ~10%. - # - # The bug pre-dates chained comparisons (`and` / `or` duplicate operands - # too) and cannot be fixed in Python: transpilation runs at UDF - # construction, before any call site exists. It also cannot be fixed in - # the JVM builder, where an argument is still unresolved (`rand(7)` - # reports deterministic=true as an UnresolvedFunction). The fix belongs - # in ResolveTranspiledPythonUDFOptions, which runs post-resolution. - # - # Assert on the plan rather than on observed proportions: a statistical - # assertion would be flaky, whereas leaving EvalPython in place means the - # interpreted UDF runs, which evaluates each argument exactly once per - # row and so is correct by construction. + # The fix belongs in ResolveTranspiledPythonUDFOptions. It cannot go in + # Python (no call site exists at UDF construction) nor in the JVM builder + # (`rand(7)` is still an UnresolvedFunction there and reports + # deterministic=true). Assert on the plan, not on observed proportions: a + # statistical assertion would be flaky, and keeping EvalPython means the + # interpreted UDF runs, which is correct by construction. from pyspark.sql.functions import col, rand band = lambda x: 50 < x < 60 # noqa: E731 band_or = lambda x: x < 50 or x > 60 # noqa: E731 with self.sql_conf(_TRANSPILE_ON): for label, func in (("50 < x < 60", band), ("x < 50 or x > 60", band_or)): - pudf = UserDefinedFunction(func, BooleanType()) - self.assertTrue(pudf.transpiled, f"{label} should transpile on its own") - base = self.spark.range(0, 8).select((col("id") * 10).alias("a")) + # subTest so the chain failing does not hide the `or` case, which + # has the same bug by the same mechanism. + with self.subTest(func=label): + pudf = UserDefinedFunction(func, BooleanType()) + self.assertTrue(pudf.transpiled, f"{label} should transpile on its own") + base = self.spark.range(0, 8).select((col("id") * 10).alias("a")) - # Deterministic argument -> transpiled and elided. - det = base.select(pudf("a")) - self.assertEqual( - 0, - self._eval_python_count(det), - f"{label}: a deterministic argument should still be elided", - ) + # Deterministic argument -> transpiled and elided. + det = base.select(pudf("a")) + self.assertEqual( + 0, + self._eval_python_count(det), + f"{label}: a deterministic argument should still be elided", + ) - # Non-deterministic argument -> must fall back. - nondet = base.select(pudf(rand(7) * 100)) - self.assertEqual( - 1, - self._eval_python_count(nondet), - f"{label}: a non-deterministic argument must NOT be transpiled", - ) - # And it still produces results. - self.assertEqual(8, len(nondet.collect())) + # Non-deterministic argument -> must fall back. + nondet = base.select(pudf(rand(7) * 100)) + self.assertEqual( + 1, + self._eval_python_count(nondet), + f"{label}: a non-deterministic argument must NOT be transpiled", + ) + # And it still produces results. + self.assertEqual(8, len(nondet.collect())) def test_udf_transpile_config_toggle_no_stale_nodes(self): # Built with the flags on, executed with them off -> clean fallback to @@ -2874,16 +2782,11 @@ def test_udf_transpile_known_value_divergences(self): self._vals(eq_strlit, BooleanType(), "a long", [(5,), (3,)], require_lowered=False), [False, False], ) - # Chained comparisons inherit the same NaN divergence, one link at a - # time -- no new mechanism, but easier to hit. `a < b < NaN` reaches the - # second link whenever `a < b`, and Spark orders NaN above every value - # (Python's `b < nan` is False); `a == b == NaN` likewise, since Spark - # treats `NaN = NaN` as true where Python's `nan == nan` is False. - # - # These assertions pin the CURRENT (divergent) behavior so a change is - # visible, not an endorsement of it. Whether to special-case NaN, fall - # back to interpreted Python, or document it as a known limitation is - # tracked in SPARK-58781. + # Chained comparisons inherit the same NaN divergence one link at a time -- + # no new mechanism, just easier to hit. Spark orders NaN above every value + # (Python's `b < nan` is False) and treats `NaN = NaN` as true (Python's + # `nan == nan` is False). These assertions pin the CURRENT divergent + # behavior so a change is visible; SPARK-58781 decides what to do about it. nan_chain = lambda a, b, c: a < b < c # noqa: E731 nan_eq_chain = lambda a, b, c: a == b == c # noqa: E731 nan = float("nan") diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index fdce74a02062d..533d464a6aeb2 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -92,15 +92,24 @@ An argument the body never reads is not computed at all. Constructs that short-circuit in Python -- ``and`` / ``or`` and chained -comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``, which -short-circuit at evaluation time. Their operands must be provably boolean; -see ``_is_definitely_boolean``. Note the fold is splittable by -``splitConjunctivePredicates``, so a ``raise_error`` NULL guard can be pushed -across a join until ``RaiseError`` overrides ``Expression.throwable``. +comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``, which also +short-circuit at evaluation time: ``And`` skips its right operand when the left +is false and ``Or`` when the left is true, in both ``eval`` and ``doGenCode``. +That is what keeps the ``raise_error`` NULL guard in a later operand from firing +on rows Python never reaches. ``and`` / ``or`` operands must be provably boolean +and non-NULL (see ``_is_never_null_boolean``); chain links need no such gate, +because every comparison lowering is total. + +Two known divergences in the fold: it is splittable by +``splitConjunctivePredicates``, so a NULL guard can be pushed across a join +until ``RaiseError`` overrides ``Expression.throwable``; and Spark orders NaN +above every value and treats ``NaN = NaN`` as true where Python does neither +(SPARK-58781). """ import ast import contextlib +import functools import inspect import itertools import operator @@ -140,21 +149,15 @@ from pyspark.sql._typing import DataTypeOrString -# Cap on how large a chained comparison we will lower, to bound plan growth: -# longer ones fall back to interpreted Python. It is generous next to what shows -# up in real code (``0 <= x <= 100``), and it matters because the lowered subtree -# is duplicated -- once per adjacent pair for a chain operand, once per -# input-category variant for every option -- and CheckAnalysis walks each copy. -# -# A token count is the wrong bound: it does not measure what actually matters, -# the size of the emitted expression tree, and a cap per syntax form does not -# converge as forms are added. A real size budget, configurable and counted -# across the whole lowering, is the follow-up that replaces this. -_MAX_COMPARISON_TERMS = 64 - -# The four ordering comparisons, which all lower through -# ``_lower_value_compare`` and differ only in the operator to apply and the -# symbol to name it in an error message. +# The most operators we will lower in one chained comparison (64 operators over +# 65 operands); longer chains fall back to interpreted Python. Each interior +# operand's subtree is emitted once per adjacent pair, and again per +# input-category variant, so this bounds plan growth -- generously, next to real +# code like ``0 <= x <= 100``. +_MAX_COMPARISON_OPS = 64 + +# The four ordering comparisons: same lowering, differing only in the operator +# and the symbol used to name it in an error message. _ORDERING_OPS: dict[type, Tuple[Callable[[Any, Any], Column], str]] = { ast.Lt: (operator.lt, "<"), ast.LtE: (operator.le, "<="), @@ -249,15 +252,16 @@ def _is_never_null_boolean(node: ast.AST) -> bool: """Return True when ``node`` lowers to a boolean column that is never NULL. Stricter than ``_is_definitely_boolean``, which also accepts operands that - may evaluate to NULL (a literal ``None``, a ternary with a ``None`` branch). - ``and`` / ``or`` need this even under the ``&`` / ``|`` fold, because Spark's - three-valued logic diverges from Python's short-circuit: ``NULL AND false`` - is ``false`` where Python's ``None and False`` is ``None``. Checking for a - bare ``None`` operand is not enough -- a ternary that *evaluates* to ``None`` - (``(True if c else None) and x``) reaches the fold the same way. - - ``not`` and if/ternary *tests* keep using ``_is_definitely_boolean``, since - they wrap the operand in ``coalesce`` and so handle NULL explicitly. + may evaluate to NULL. ``and`` / ``or`` need that: Spark's ``NULL AND false`` + is ``false`` where Python's ``None and False`` is ``None``, and rejecting a + bare ``None`` operand is not enough because a ternary can evaluate to one + (``(True if c else None) and x``). + + ``not`` and if/ternary *tests* keep using ``_is_definitely_boolean``: they + wrap the operand in ``coalesce``, so they handle NULL explicitly. + + This answers "would it be NULL?", not "can it be lowered at all" -- it says + True for shapes ``_convert_chunk`` goes on to refuse, such as ``x in "abc"``. """ match node: case ast.Constant(value=v): @@ -530,23 +534,14 @@ def _lower_compare_pair( ) -> Column: """Lower ONE adjacent comparison of a (possibly chained) ``ast.Compare``. - TOTALITY INVARIANT, which callers depend on: when the returned column - returns a value at all, that value is never NULL. - - * ``is`` / ``is not`` emit ``isNull`` / ``isNotNull``, non-nullable by - construction. - * ``_lower_eq`` answers every NULL combination with a boolean literal, - reaching its ``otherwise`` only when both operands are non-NULL. - * ``_lower_value_compare`` raises rather than returning NULL. - - Equal operand categories plus the JVM-side category contract in - ``ResolveTranspiledPythonUDFOptions`` (non-decimal ``NumericType`` for - "numeric", UTF8_BINARY ``StringType`` for "string") leave numeric + Callers rely on the result never being NULL: ``is`` / ``is not`` emit + ``isNull`` / ``isNotNull``, ``_lower_eq`` answers every NULL combination + with a boolean literal, and ``_lower_value_compare`` raises rather than + returning NULL. Its non-NULL branch cannot yield NULL either, because + equal operand categories plus the JVM-side category contract leave numeric widening as the only coercion Catalyst inserts, and under ANSI a failing - cast raises instead of yielding NULL. - - That invariant is what lets the chain fold use plain ``&`` with no - ``coalesce``: every link is total, so none of them yields NULL. + cast raises. That is what lets the chain fold below use a plain ``&`` with + no ``coalesce``. """ match op: case ast.Is() | ast.IsNot(): @@ -568,24 +563,17 @@ def _lower_compare_pair( ) subject_node = right_node if is_none_left else left_node subject_col = self._convert_chunk(params, subject_node) - if isinstance(op, ast.Is): - return subject_col.isNull() - else: - return subject_col.isNotNull() + return subject_col.isNull() if isinstance(op, ast.Is) else subject_col.isNotNull() case ast.Eq(): return self._lower_eq(params, left_node, right_node, equal=True) case ast.NotEq(): return self._lower_eq(params, left_node, right_node, equal=False) case ast.Lt() | ast.LtE() | ast.Gt() | ast.GtE(): - # All four orderings share `_lower_value_compare` and differ only - # in the operator, so look it up rather than repeating the call. op_fn, op_str = _ORDERING_OPS[type(op)] return self._lower_value_compare(params, left_node, right_node, op_fn, op_str) case _: # Reached only if Python gains a new `cmpop`: the other eight are - # handled above and `in` / `not in` are refused in - # `_convert_chunk`. Kept so a new operator falls back rather than - # returning None from a `-> Column` function. + # handled above and `in` / `not in` are refused by the caller. raise UnsupportedOperationException( f"comparison operator {type(op).__name__} is not supported by the transpiler" ) @@ -687,18 +675,12 @@ def branch_category(b: ast.AST) -> Optional[str]: ) return result_cat case ast.Compare(): - # EVERY Python comparison evaluates to a bool -- `in` / `not in` - # and chained forms included -- whatever its operands look like. - # `_is_definitely_boolean` can't answer this because its Compare - # arm also requires basic-type operands, so it says False for - # e.g. `(x if c else y) in (1, 2)` or `(a < b) == True`. Those - # then fell through to the numeric catch-all below, which let the - # return-type match in `_transpile_from_ast` accept a numeric - # declared return type for a boolean body: the lowered option - # returned `cast( as bigint)` (1/0) where the - # interpreted SQL_BATCHED_UDF path returns NULL, a silent - # divergence. Classify on the node kind instead, which is a fact - # about Python rather than an operand analysis. + # Checked by node kind, not through `_is_definitely_boolean`: that + # also demands basic-type operands, so it misses `(a < b) == True`, + # which then reached the numeric catch-all below and let a boolean + # body satisfy a numeric declared return type -- emitting + # `cast( as bigint)` where the interpreted path returns + # NULL. EVERY Python comparison evaluates to a bool. return "bool" case _ if _is_definitely_boolean(node): # `not` and boolean ops produce a boolean column. Labeling them @@ -767,12 +749,9 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # [] all falsy), which we can't reproduce without the input # column types -- Spark's `&` / `|` would do bitwise instead. # - # The gate is `_is_never_null_boolean`, not - # `_is_definitely_boolean`: Spark's three-valued `NULL AND false` - # is `false` where Python's `None and False` is `None`, so a - # NULL-valued operand diverges. Checking only for a bare `None` - # operand is not enough, because a ternary that evaluates to None - # (`(True if c else None) and x`) reaches the fold the same way. + # The gate is `_is_never_null_boolean` rather than + # `_is_definitely_boolean` because Spark's `NULL AND false` is + # `false` where Python's `None and False` is `None`. if not all(_is_never_null_boolean(v) for v in values): raise UnsupportedOperationException( "`and` / `or` operand is not statically known to be a " @@ -783,17 +762,15 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: "to interpreted Python" ) cols = [self._convert_chunk(params, v) for v in values] - if isinstance(op, ast.And): - result = cols[0] - for c in cols[1:]: - result = result & c - return result - if isinstance(op, ast.Or): - result = cols[0] - for c in cols[1:]: - result = result | c - return result - raise UnsupportedOperationException(f"BoolOp operator {op} is not supported") + match op: + case ast.And(): + return functools.reduce(operator.and_, cols) + case ast.Or(): + return functools.reduce(operator.or_, cols) + case _: + raise UnsupportedOperationException( + f"BoolOp operator {op} is not supported" + ) case ast.IfExp(test=test, body=body_expr, orelse=orelse_expr): # Ternary `body if test else orelse` -- shares the # NULL-as-falsy lowering with the if-statement case. @@ -817,93 +794,44 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: orelse[0] if orelse else None, ) case ast.Compare(left, ops, comps): - # `in` / `not in` are container membership rather than a chainable - # value comparison, and lowering them needs its own NULL analysis - # (Python's `in` is equality-based and never raises on a `None` - # probe, while an unguarded NULL subject is indistinguishable from - # one). Refuse them here so the UDF falls back to interpreted - # Python; support is deferred to SPARK-56925. + # `in` / `not in` are container membership, not a chainable value + # comparison, and need their own NULL analysis: Python's `in` is + # equality-based and never raises on a `None` probe. Deferred to + # SPARK-56925, so refuse them and fall back. if any(isinstance(o, (ast.In, ast.NotIn)) for o in ops): raise UnsupportedOperationException( "`in` / `not in` is not supported by the transpiler; the " "UDF falls back to interpreted Python" ) - if len(ops) > _MAX_COMPARISON_TERMS: + if len(ops) > _MAX_COMPARISON_OPS: raise UnsupportedOperationException( - f"chained comparison with {len(ops)} operators exceeds " - f"the {_MAX_COMPARISON_TERMS} the transpiler will lower; the " + f"chained comparison has {len(ops)} operators; the " + f"transpiler lowers at most {_MAX_COMPARISON_OPS}, so the " "UDF falls back to interpreted Python" ) # Python evaluates `a OP1 b OP2 c` as `(a OP1 b) and (b OP2 c)` - # with `b` evaluated once and with SHORT-CIRCUIT: when - # `a OP1 b` is falsy, `b OP2 c` is never evaluated. That matters - # because `_lower_value_compare` emits a `raise_error` for NULL - # operands to mirror Python's TypeError: `f(1, 0, None)` for - # `lambda a, b, c: a < b < c` must return False, not raise. - # - # `&` short-circuits at evaluation time -- both `And.eval` and - # `And.doGenCode` skip the right operand when the left is false -- - # so the fold below reproduces that. What it does NOT survive on - # its own is the optimizer: `splitConjunctivePredicates` - # decomposes `And`, and `PushPredicateThroughJoin` pushes a - # conjunct it believes harmless, which a `raise_error` is until - # `RaiseError` overrides `Expression.throwable` (as `Sequence` - # already does). That override is a prerequisite for this - # lowering: without it, a chain under a filter over a join can - # have a link's guard pushed below the join and raise on rows - # Python never evaluates. The alternative -- folding into nested - # CASE WHENs, atomic to predicate splitting -- was rejected - # because it costs pushdown for every `and` / `or` UDF, including - # those over non-nullable columns that need no guard at all. - # - # Python's "evaluate `b` once" rule is about side effects. The - # *Python-level* operand grammar we accept (Constant / Name / - # BinOp / UnaryOp / IfExp / Compare / BoolOp) is deterministic - # and side-effect free -- ANSI raises aside, and those are - # idempotent -- so re-lowering an interior operand once per - # adjacent pair is safe. Not so for the CALL-SITE arguments: a - # parameter lowers to a `_udf_param_N` placeholder and the JVM's - # `resolveUDFParams` substitutes the caller's expression at EVERY - # occurrence, four of them for a two-link ordering chain. A - # NON-DETERMINISTIC argument therefore becomes that many - # independent expressions with independent state, and because the - # fold above evaluates later links only on the rows where earlier - # ones held, those copies desynchronize and stop describing the - # same value: `50 < x < 60` over `rand(7) * 100` yields ~30% true - # where Python yields ~10%. - # - # That is an OPEN BUG, not something this code can fix: the - # transpiler runs at UDF construction, long before a call site - # exists. The fix belongs in the analyzer, which is the first - # place a resolved argument can be tested for determinism, and is - # being handled as a companion change; until it lands the - # divergence is pinned by a deliberately failing test in - # test_udf_transpile_unit.py. + # with `b` evaluated once, so the fold below re-lowers each + # interior operand once per adjacent pair. Safe for the operand + # grammar we accept, all of it deterministic and side-effect free + # -- but NOT for the call-site arguments the JVM's + # `resolveUDFParams` substitutes into `_udf_param_N`, where a + # non-deterministic argument becomes several independent + # expressions that desynchronize (`50 < x < 60` over + # `rand(7) * 100` is ~30% true where Python is ~10%). That has to + # be fixed in the analyzer, which is the first place a resolved + # argument exists; `and` / `or` already had it and chains are newly + # exposed to it. # - # Be precise about the blast radius, because it is easy to - # overstate as merely pre-existing. `and` / `or` were already - # wrong this way, since their operands are separate subtrees. A - # SINGLE comparison is NOT: its duplicate operand appears only - # inside an `IS NULL` test, whose outcome does not depend on the - # value drawn, so `50 < x` over `rand()` samples correctly. So - # lowering chained comparisons -- which previously fell back to - # interpreted Python, and were therefore correct -- newly exposes - # them to this bug. That argues for landing the analyzer guard - # before or with this change. - # - # NaN is a separate known divergence: Spark orders NaN above every - # value and treats `NaN = NaN` as true, Python does neither, so a - # chain over doubles can disagree link by link. Tracked in - # SPARK-58781. + # See the module docstring for the short-circuit and NULL-guard + # semantics this fold relies on. nodes = [left, *comps] - links = [ - self._lower_compare_pair(params, op, left_node, right_node) - for op, left_node, right_node in zip(ops, nodes, nodes[1:]) - ] - result = links[0] - for link in links[1:]: - result = result & link - return result + return functools.reduce( + operator.and_, + [ + self._lower_compare_pair(params, op, left_node, right_node) + for op, left_node, right_node in zip(ops, nodes, nodes[1:]) + ], + ) case ast.BinOp(left=left, op=op, right=right): # Operator selection is driven by the operand *categories* under # the current input-type variant (see ``_category``): Python's diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala index dd5381ac70ab5..7a6f6a5ae7ba5 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala @@ -94,21 +94,18 @@ class ResolveTranspiledPythonUDFOptionsSuite extends PlanTest { } test("empties the options when an argument is non-deterministic") { - // EXPECTED TO FAIL until the companion fix to this rule lands. Deliberately - // red: it pins a real wrong-results bug rather than tolerating it silently. - // Do not ignore or cancel it. + // RED UNTIL the companion fix to this rule lands -- deliberately, to pin a + // real wrong-results bug rather than tolerate it silently. Do not ignore it. // // The builder substitutes the argument expression at EVERY `_udf_param_N` // occurrence, so a non-deterministic argument becomes several independent // expressions with independent state. Since `and` / `or` and chained // comparisons evaluate later copies only on rows where earlier links held, - // the copies desynchronize and stop describing the same value. The rule must - // drop every option so ConvertToCatalyst falls back to the Python UDF, which - // evaluates each argument exactly once per row. - // - // This must be checked here rather than in the builder: at call-construction - // time `rand(7)` is still an `UnresolvedFunction` whose `deterministic` - // derives from its `Literal` seed and so reports true. + // the copies desynchronize. The rule must drop every option so + // ConvertToCatalyst falls back to the Python UDF, which evaluates each + // argument exactly once per row. It has to be checked here rather than in the + // builder: at call-construction time `rand(7)` is still an + // `UnresolvedFunction` and reports deterministic = true. val a = $"a".double val arg = Multiply(Rand(Literal(7L)), Literal(100.0)) assert(!arg.deterministic, "the fixture argument must be non-deterministic") @@ -120,8 +117,8 @@ class ResolveTranspiledPythonUDFOptionsSuite extends PlanTest { } test("keeps options when the argument is deterministic (control for the check above)") { - // Same shape as the non-deterministic case, but with a plain column argument, so - // the non-determinism guard must not fire and normal category matching applies. + // Same shape as the non-deterministic case but with a plain column argument, so + // the guard must not fire and normal category matching applies. val a = $"a".double val opt = GreaterThan(a, Literal(50.0)) val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)), List(opt), List(List("numeric"))) From 62e880d1b6b4281b2dbacc1e12f5aa4952d8fa9d Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 15 Aug 2026 05:49:27 +0000 Subject: [PATCH 06/18] [SPARK-58456][PYTHON][SQL][MINOR] Put back the rationale the trim cut too far Review of the previous commit flagged four places where the trim took the "why" along with the prose. Restored condensed, comments only. - `_MAX_COMPARISON_OPS` reads as a considered design without the note that it is a stopgap. An operator count does not measure the emitted tree size, and a cap per syntax form does not converge: `and` / `or` admits the same unbounded repetition and is capped nowhere, which is the standing proof. Says so again, and names the replacement (one budget counted across the whole lowering). No cap is added to `and` / `or` here; that belongs with the shared budget. - `_lower_compare_pair` referred to "the JVM-side category contract" with no pointer to it. Names `ResolveTranspiledPythonUDFOptions` and the contract again, since the two have to change in lockstep. - The chain fold lost the rejected alternative (nested CASE WHENs, atomic to predicate splitting, rejected because it costs pushdown for every `and` / `or` UDF including those needing no guard) and the blast-radius reasoning: a SINGLE comparison is not affected by the non-deterministic argument bug, because its duplicate operand sits only inside an `IS NULL` test whose outcome does not depend on the value drawn. The conclusion had survived, the reasoning that lets a reader check it had not. - The hypothesis templates asserted annotations "buy plan size, not wall clock" with the measurements removed. Numbers are back (1 option instead of 2; 73s untyped vs 76s typed at 15 examples). Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../tests/test_udf_transpile_hypothesis.py | 10 ++++-- python/pyspark/sql/transpile.py | 32 +++++++++++++++---- 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py index ffaf2377c924a..24b56f227a5cd 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -265,9 +265,13 @@ def _seed_pair_examples(pairs, keys=("x", "y")): # # Where a template annotates its parameters, that is to pin the input category: # an untyped parameter emits both a numeric and a string variant, and these -# tests bind Long columns, so the string one is dead weight. It buys plan size, -# not wall clock. The strategies still generate ``None`` -- an annotation pins -# the category, not nullability. +# tests bind Long columns, so the string one is dead weight (measured: 1 option +# emitted instead of 2). It buys plan size, not wall clock -- measured 73s +# untyped vs 76s typed over these five tests at 15 examples, because the cost is +# the two Spark jobs ``_run`` executes per example against a fresh one-row +# DataFrame. Batching rows per example is the lever that would matter. The +# strategies still generate ``None`` -- an annotation pins the category, not +# nullability. def plus_four(x): diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 533d464a6aeb2..3a7da57b36cce 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -154,6 +154,12 @@ # operand's subtree is emitted once per adjacent pair, and again per # input-category variant, so this bounds plan growth -- generously, next to real # code like ``0 <= x <= 100``. +# +# This is a stopgap, not the intended design. An operator count does not measure +# what actually matters, the size of the emitted tree, and a cap per syntax form +# does not converge as forms are added -- ``and`` / ``or`` admits the same +# unbounded repetition and is not capped at all. The replacement is one size +# budget, configurable and counted across the whole lowering. _MAX_COMPARISON_OPS = 64 # The four ordering comparisons: same lowering, differing only in the operator @@ -538,10 +544,12 @@ def _lower_compare_pair( ``isNull`` / ``isNotNull``, ``_lower_eq`` answers every NULL combination with a boolean literal, and ``_lower_value_compare`` raises rather than returning NULL. Its non-NULL branch cannot yield NULL either, because - equal operand categories plus the JVM-side category contract leave numeric - widening as the only coercion Catalyst inserts, and under ANSI a failing - cast raises. That is what lets the chain fold below use a plain ``&`` with - no ``coalesce``. + equal operand categories plus the category contract that + ``ResolveTranspiledPythonUDFOptions`` enforces on the JVM side + (non-decimal ``NumericType`` for "numeric", UTF8_BINARY ``StringType`` + for "string") leave numeric widening as the only coercion Catalyst + inserts, and under ANSI a failing cast raises. That is what lets the + chain fold below use a plain ``&`` with no ``coalesce``. """ match op: case ast.Is() | ast.IsNot(): @@ -819,8 +827,20 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # expressions that desynchronize (`50 < x < 60` over # `rand(7) * 100` is ~30% true where Python is ~10%). That has to # be fixed in the analyzer, which is the first place a resolved - # argument exists; `and` / `or` already had it and chains are newly - # exposed to it. + # argument exists. + # + # Be precise about the blast radius. `and` / `or` were already + # wrong this way, their operands being separate subtrees. A SINGLE + # comparison is not: its duplicate operand appears only inside an + # `IS NULL` test, whose outcome does not depend on the value drawn, + # so `50 < x` over `rand()` samples correctly. Chains previously + # fell back to interpreted Python, so lowering them is what newly + # exposes them to this. + # + # Folding into nested CASE WHENs instead -- atomic to predicate + # splitting -- was rejected: it costs pushdown for every + # `and` / `or` UDF, including those over non-nullable columns that + # need no guard at all. # # See the module docstring for the short-circuit and NULL-guard # semantics this fold relies on. From 918908ab14394e849da43e0529aa3abeafdab961 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 15 Aug 2026 05:51:25 +0000 Subject: [PATCH 07/18] [SPARK-58456][PYTHON][SQL][MINOR] Replace the cap's hand-waving with measurements The comment claimed the cap bounds plan growth "generously" and implied the emitted tree multiplies per input-category variant. Measured it instead, at the cap over three untyped params: ops=2 variants=2 expr_chars=564 raise_error=2 ops=8 variants=2 expr_chars=2196 raise_error=8 ops=64 variants=2 expr_chars=17428 raise_error=64 So growth is linear in the operator count, not combinatorial, and only two options ever survive: every adjacent pair of a chain must share a category, which leaves all-numeric and all-string as the only variants that lower at all. Three untyped params would otherwise give 2**3 = 8 combos, and the other six are refused rather than emitted. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- python/pyspark/sql/transpile.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 3a7da57b36cce..655ac773c5db8 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -151,9 +151,12 @@ # The most operators we will lower in one chained comparison (64 operators over # 65 operands); longer chains fall back to interpreted Python. Each interior -# operand's subtree is emitted once per adjacent pair, and again per -# input-category variant, so this bounds plan growth -- generously, next to real -# code like ``0 <= x <= 100``. +# operand is lowered once per adjacent pair, so growth is linear in the operator +# count, and only two options can survive: every adjacent pair must share a +# category, which leaves the all-numeric and all-string variants as the only ones +# that lower at all. Measured at the cap over three untyped params: 2 options, +# one `raise_error` guard per operator, ~17 KB of expression per option. Generous +# next to real code like ``0 <= x <= 100``. # # This is a stopgap, not the intended design. An operator count does not measure # what actually matters, the size of the emitted tree, and a cap per syntax form From fc480294127f7cbfe796b5bd6e4c18618b734d6b Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Tue, 18 Aug 2026 22:34:28 +0000 Subject: [PATCH 08/18] [SPARK-58456][PYTHON][SQL] Bound the whole lowering, not one syntax form Review found `_MAX_COMPARISON_OPS` does not bound plan growth at all, and the comment claiming it does was wrong. A comparison's operand may itself be a comparison -- both have category "bool", so `(a < b) < (c < d)` type-checks -- and a chain re-lowers each interior operand once per adjacent pair. Nesting the two doubles the work per level while every node stays far under the cap. Measured, with `L(d) = "((a 0.57s transpiled; depths 8, 11, 15, 20 -> under 1s, fall back `_MAX_COMPARISON_OPS` stays. The budget subsumes it above 128 operators, but it still bounds a single chain's width, which is what sets the ~17 KB-per-option figure, and removing it would silently start transpiling 65-128 operator chains that currently fall back. Also from the review: - The module docstring presented the `&` / `|` short-circuit as making the fold safe, listed two divergences, and named `Expression.throwable` as the remedy. Measured three, and the remedy fixes one. A PURE chain with no `and` / `or` (`lambda a, b, c: a < b < c`) over a join returns False correctly under `select` but raises USER_RAISED_EXCEPTION under `filter`, so SPARK-58627 is a prerequisite. Separately, on a plain `select` with no join and no filter, `(a < b < c) or (d > 0 and b < c)` on `(5, 1, None, -1)` raises because `spark.sql.subexpressionElimination.skipForShortcutExpr` defaults false and hoists a conjunct out of a short-circuited position; setting it true alone fixes it, and `throwable` does not help since that rule never consults it. Also corrected "skips when the left is false" to "exactly false" -- a NULL left operand still evaluates the right. - Restored the fallback-warning assertion the earlier consolidation dropped: both sites used `catch_warnings(record=True)` with no binding, so they only suppressed. Factored the filter the suite repeats in six places into `_fallback_warnings`. - Dropped the no-op `sql_conf(_TRANSPILE_ON)` from the caps test; nothing on that path reads SQLConf. Extended it to cover the budget, asserting every node is under the per-chain cap so it is provably the budget that fires. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../sql/tests/test_udf_transpile_unit.py | 95 ++++++++++++++---- python/pyspark/sql/transpile.py | 98 ++++++++++++++----- 2 files changed, 152 insertions(+), 41 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index b7b5855a6ca3a..ed608a480a132 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -53,6 +53,20 @@ } +def _fallback_warnings(caught): + """The warnings that mean "this UDF fell back to interpreted Python". + + Asserting on these is what distinguishes a deliberate refusal from a + transpile that silently produced nothing: ``transpiled == []`` alone is also + what you get if an exception is swallowed before the reason is recorded. + """ + return [ + w + for w in caught + if "Unable to transpile" in str(w.message) or "Errors encountered" in str(w.message) + ] + + @unittest.skipIf( is_remote_only(), "UDF transpilation is only supported in regular (non-Connect) Spark.", @@ -666,10 +680,14 @@ def test_udf_transpile_comparison_fallbacks(self): with self.sql_conf(_TRANSPILE_ON): for label, func, schema, rows, expected in cases: with self.subTest(func=label): - with _warnings.catch_warnings(record=True): + with _warnings.catch_warnings(record=True) as caught: _warnings.simplefilter("always") pudf = UserDefinedFunction(func, BooleanType()) self.assertEqual([], pudf.transpiled, f"{label} must NOT transpile") + self.assertTrue( + _fallback_warnings(caught), + f"{label}: expected a fallback warning", + ) df = self.spark.createDataFrame(rows, schema) self.assertEqual( [r[0] for r in df.select(pudf(*df.columns)).collect()], @@ -686,13 +704,17 @@ def test_udf_transpile_comparison_fallbacks(self): (in_long_rt, LongType(), "an `in` with a LongType return type"), ]: with self.subTest(label=label): - with _warnings.catch_warnings(record=True): + with _warnings.catch_warnings(record=True) as caught: _warnings.simplefilter("always") self.assertEqual( [], UserDefinedFunction(func, rt).transpiled, f"{label} must NOT transpile", ) + self.assertTrue( + _fallback_warnings(caught), + f"{label}: expected a fallback warning", + ) def test_udf_transpile_boolean_body_requires_boolean_return_type(self): # A comparison body is boolean whatever its operands look like, so it must @@ -747,15 +769,18 @@ def chain_of_compare(a, b, c): self.assertEqual([True], [r[0] for r in df.select(pudf("c", "x", "y")).collect()]) def test_udf_transpile_lowering_caps(self): - # The chain cap bounds plan growth. Exercise the boundary directly rather - # than through a UDF, so the refusal cannot be confused with a fallback for - # some other reason. + # Both bounds on comparison lowering. Exercised directly rather than + # through a UDF so a refusal cannot be confused with a fallback for some + # other reason. No `sql_conf` here on purpose: nothing on this path reads + # SQLConf (only `_get_transpilers` does), so wrapping it would imply the + # caps are conf-gated when they are module constants. import ast as _ast from pyspark.errors import UnsupportedOperationException from pyspark.sql.transpile import ( CatalystTranspiler, _MAX_COMPARISON_OPS, + _MAX_LOWERED_COMPARISONS, ) transpiler = CatalystTranspiler() @@ -764,23 +789,57 @@ def test_udf_transpile_lowering_caps(self): def compare_node(src): return _ast.parse(src, mode="eval").body + def lower(src): + # `_transpile_from_ast` is what normally zeroes the budget; reset it + # here so each case is charged on its own. + transpiler._lowered_comparisons = 0 + return transpiler._convert_chunk(["x"], compare_node(src)) + def chain_src(n): # n operators over n + 1 operands: x < 0 < 1 < ... < n - 1 return " < ".join(["x"] + [str(i) for i in range(n)]) - with self.sql_conf(_TRANSPILE_ON): - # At the cap the lowering succeeds. - transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_OPS))) - - # One past it, it refuses so the UDF falls back. - with self.assertRaises(UnsupportedOperationException) as ctx: - transpiler._convert_chunk(["x"], compare_node(chain_src(_MAX_COMPARISON_OPS + 1))) - self.assertIn("operators", str(ctx.exception)) - - # `in` / `not in` are refused outright, deferred to SPARK-56925. - with self.assertRaises(UnsupportedOperationException) as ctx: - transpiler._convert_chunk(["x"], compare_node("x in (1, 2)")) - self.assertIn("`in` / `not in`", str(ctx.exception)) + # ---- Per-chain operator cap ---- + lower(chain_src(_MAX_COMPARISON_OPS)) # at the cap, succeeds + with self.assertRaises(UnsupportedOperationException) as ctx: + lower(chain_src(_MAX_COMPARISON_OPS + 1)) + self.assertIn("operators", str(ctx.exception)) + + # ---- Whole-body budget ---- + # A comparison may be an operand of another comparison (both category + # "bool"), and a chain re-lowers each interior operand once per adjacent + # pair, so nesting doubles the link count per level while every node stays + # far under the per-chain cap. Before the budget existed this reached hours + # of driver-side py4j work from a few hundred characters of source. + def nested_src(depth): + src = "(x < 0)" + for _ in range(depth): + src = "((x < 0) < %s < (x < 0))" % src + return src + + # Shallow nesting stays inside the budget and still lowers. + lower(nested_src(2)) + + # Deep nesting is refused by the budget, not by the per-chain cap: every + # `ast.Compare` in here has exactly two operators. + deep = nested_src(10) + node = compare_node(deep) + self.assertTrue( + all( + len(n.ops) <= _MAX_COMPARISON_OPS + for n in _ast.walk(node) + if isinstance(n, _ast.Compare) + ), + "every node must be under the per-chain cap, so the budget is what fires", + ) + with self.assertRaises(UnsupportedOperationException) as ctx: + lower(deep) + self.assertIn(str(_MAX_LOWERED_COMPARISONS), str(ctx.exception)) + + # `in` / `not in` are refused outright, deferred to SPARK-56925. + with self.assertRaises(UnsupportedOperationException) as ctx: + lower("x in (1, 2)") + self.assertIn("`in` / `not in`", str(ctx.exception)) def test_udf_transpile_multi_row(self): # Every other transpile test uses a 1-row DataFrame; this one runs diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 655ac773c5db8..52a3ddaf83cdb 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -92,19 +92,39 @@ An argument the body never reads is not computed at all. Constructs that short-circuit in Python -- ``and`` / ``or`` and chained -comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``, which also -short-circuit at evaluation time: ``And`` skips its right operand when the left -is false and ``Or`` when the left is true, in both ``eval`` and ``doGenCode``. -That is what keeps the ``raise_error`` NULL guard in a later operand from firing -on rows Python never reaches. ``and`` / ``or`` operands must be provably boolean +comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``. That fold is +what is meant to keep the ``raise_error`` NULL guard in a later operand from +firing on rows Python never reaches, since ``And`` skips its right operand when +the left is exactly ``false`` (and ``Or`` when the left is exactly ``true``) in +both ``eval`` and ``doGenCode``. Note the "exactly": a NULL left operand still +evaluates the right. ``and`` / ``or`` operands must therefore be provably boolean and non-NULL (see ``_is_never_null_boolean``); chain links need no such gate, because every comparison lowering is total. -Two known divergences in the fold: it is splittable by -``splitConjunctivePredicates``, so a NULL guard can be pushed across a join -until ``RaiseError`` overrides ``Expression.throwable``; and Spark orders NaN -above every value and treats ``NaN = NaN`` as true where Python does neither -(SPARK-58781). +That short-circuit is an EVALUATION-time property only, and it does not survive +the optimizer. Three divergences, all measured rather than reasoned about: + +* Predicate pushdown. The fold is splittable by ``splitConjunctivePredicates``, + and ``Optimizer`` partitions on ``cond.deterministic && !cond.throwable``, + which ``RaiseError`` fails to set. So a guard can be pushed below a join and + run on rows the join drops. Measured on a PURE chain with no ``and`` / ``or``: + ``lambda a, b, c: a < b < c`` over a join, ``select`` returns False correctly + but ``filter`` raises USER_RAISED_EXCEPTION where Python returns False. Fixed + by ``RaiseError`` overriding ``Expression.throwable`` (SPARK-58627), which is + a prerequisite for this lowering rather than a nice-to-have. +* Subexpression elimination. ``spark.sql.subexpressionElimination.enabled`` + defaults true while ``...skipForShortcutExpr`` defaults FALSE, so a conjunct + living only in a short-circuited position is hoisted and evaluated + unconditionally. Measured on a plain ``select``, no join and no filter: + ``(a < b < c) or (d > 0 and b < c)`` on ``(5, 1, None, -1)`` raises where + Python returns False, and setting ``skipForShortcutExpr=true`` alone fixes it. + ``Expression.throwable`` does NOT help here -- this rule never consults it. +* NaN. Spark orders NaN above every value and treats ``NaN = NaN`` as true where + Python does neither (SPARK-58781). + +A chain duplicates every interior operand by construction, so it manufactures +the common subexpressions the second bullet needs. Treat both of the first two +as open, not theoretical. """ import ast @@ -150,21 +170,34 @@ # The most operators we will lower in one chained comparison (64 operators over -# 65 operands); longer chains fall back to interpreted Python. Each interior -# operand is lowered once per adjacent pair, so growth is linear in the operator -# count, and only two options can survive: every adjacent pair must share a -# category, which leaves the all-numeric and all-string variants as the only ones -# that lower at all. Measured at the cap over three untyped params: 2 options, -# one `raise_error` guard per operator, ~17 KB of expression per option. Generous -# next to real code like ``0 <= x <= 100``. -# -# This is a stopgap, not the intended design. An operator count does not measure -# what actually matters, the size of the emitted tree, and a cap per syntax form -# does not converge as forms are added -- ``and`` / ``or`` admits the same -# unbounded repetition and is not capped at all. The replacement is one size -# budget, configurable and counted across the whole lowering. +# 65 operands); longer chains fall back to interpreted Python. For a FLAT chain +# this is close to linear: each interior operand is lowered once per adjacent +# pair, and only two options survive because every adjacent pair must share a +# category, leaving all-numeric and all-string. Measured over three untyped +# params: 2 options, one `raise_error` guard per operator, ~17 KB of expression +# per option. Generous next to real code like ``0 <= x <= 100``. _MAX_COMPARISON_OPS = 64 +# Total comparison links we will lower for one function body, across every +# nesting level. This is the bound that actually matters, and a per-form operator +# count cannot express it: a comparison's operand may itself be a comparison +# (both have category "bool", so `(a < b) < (c < d)` type-checks), and a chain +# re-lowers each interior operand once per adjacent pair. Nesting those two facts +# doubles the work per level while every individual node stays far under +# `_MAX_COMPARISON_OPS`. +# +# Measured before this bound existed, with `L(d) = "((a None: self._param_categories: dict[int, str] = {} self._category_cache: dict[int, str] = {} @@ -553,7 +591,17 @@ def _lower_compare_pair( for "string") leave numeric widening as the only coercion Catalyst inserts, and under ANSI a failing cast raises. That is what lets the chain fold below use a plain ``&`` with no ``coalesce``. + + Every comparison link routes through here, which is what makes this the + place to charge ``_MAX_LOWERED_COMPARISONS``. """ + self._lowered_comparisons += 1 + if self._lowered_comparisons > _MAX_LOWERED_COMPARISONS: + raise UnsupportedOperationException( + f"lowering this body needs more than {_MAX_LOWERED_COMPARISONS} " + "comparisons, so the UDF falls back to interpreted Python; " + "nested comparisons double the count per level" + ) match op: case ast.Is() | ast.IsNot(): # Only lower `x is None` / `None is x` (and their @@ -987,6 +1035,10 @@ def _transpile_from_ast( # Category inference depends on the per-variant assumptions above. Cache # each AST node only for this lowering so recursive conversion stays linear. self._category_cache = {} + # Charged by ``_lower_compare_pair``. Reset per variant, not per process: + # the budget bounds one body's lowering, and each variant lowers the body + # again from scratch. + self._lowered_comparisons = 0 function_body = function_ast.body if len(function_body) != 1: raise UnsupportedOperationException( From ba4b14eee62d5095a2d1218ba985ed7794099bb8 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Fri, 21 Aug 2026 23:22:34 +0000 Subject: [PATCH 09/18] [SPARK-58456][PYTHON][SQL] SPARK-58627 landed; retire the stale RED UNTIL notes Rebasing onto current master picked up 7f10fc41071 [SPARK-58627][SQL] Mark raise error Throwable & fix sequence throwable, so `RaiseError` now overrides `Expression.throwable` (misc.scala:87). That was the prerequisite this branch was pinning, and it closes the pushdown divergence: - `test_udf_transpile_boolop_short_circuits` passes. It was documented "RED UNTIL `RaiseError` overrides `Expression.throwable`", which is now a marker inviting someone to fix an already-correct test. Reframed as the regression test for SPARK-58627, with why it should stay: the override is what makes the `&` / `|` fold survive the optimizer. - Re-ran the pure-chain repro that used to fail. `lambda a, b, c: a < b < c` over a join, no `and` / `or` anywhere: `filter` returns 0 rows, matching Python's False, where before it raised USER_RAISED_EXCEPTION. `select` was and is correct. The module docstring now marks each of the three divergences CLOSED or OPEN rather than implying all are live, and says explicitly that SPARK-58627 does NOT cover the subexpression-elimination one, because that rule never consults `throwable`. Still open: subexpression elimination hoisting a conjunct out of a short-circuited position, and NaN ordering (SPARK-58781). Local, on the rebased tree: 59 unit tests with 2 failures, both the analyzer determinism red; catalyst 6 passed / 1 failed, the same red. That red stays -- the analyzer fix has no JIRA or branch yet, so its RED UNTIL notes are accurate and untouched. Co-Authored-By: Claude Opus 5 (1M context) Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../sql/tests/test_udf_transpile_unit.py | 23 ++++++------ python/pyspark/sql/transpile.py | 35 ++++++++++--------- 2 files changed, 30 insertions(+), 28 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index ed608a480a132..6406f7119dfd0 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -548,16 +548,17 @@ def test_udf_transpile_boolop_short_circuits(self): # the second operand must not fire when the first already decided the # result. # - # RED UNTIL `RaiseError` overrides `Expression.throwable` -- deliberately, - # to pin the prerequisite. Without the override, - # `splitConjunctivePredicates` decomposes the `And` and - # `PushPredicateThroughJoin` pushes the guard to one side of the join, - # where it fires on rows Python short-circuits past and raises - # USER_RAISED_EXCEPTION instead of returning False. + # This is the regression test for SPARK-58627. Before `RaiseError` + # overrode `Expression.throwable`, `splitConjunctivePredicates` decomposed + # the `And` and `PushPredicateThroughJoin` pushed the guard to one side of + # the join, where it fired on rows Python short-circuits past and raised + # USER_RAISED_EXCEPTION instead of returning False. Keep it: the override + # is what makes the `&` / `|` fold survive the optimizer, so this is the + # test that notices if it regresses. # - # The join case disables AQE and broadcast on purpose: by default whether - # the pushed-down guard fires races AQE eliminating the empty join, so the - # failure would be intermittent rather than a clean red. + # The join case disables AQE and broadcast on purpose: otherwise whether + # the pushed-down guard fires races AQE eliminating the empty join, and the + # test would pass for the wrong reason. def both_positive(a, c): return a > 0 and c > 0 @@ -590,8 +591,8 @@ def either_positive(a, c): joined = left.join(right, "k") self.assertEqual([False], [r[0] for r in joined.select(and_udf("a", "c")).collect()]) - # Filtering must return Python's answer rather than raising. RED - # until `RaiseError` overrides `Expression.throwable`; see above. + # Filtering must return Python's answer rather than raising. This is + # the assertion SPARK-58627 fixed; see above. filtered = joined.filter(and_udf("a", "c")) self.assertEqual([], filtered.select("k").collect()) diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 52a3ddaf83cdb..8e908280e2564 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -101,30 +101,31 @@ and non-NULL (see ``_is_never_null_boolean``); chain links need no such gate, because every comparison lowering is total. -That short-circuit is an EVALUATION-time property only, and it does not survive -the optimizer. Three divergences, all measured rather than reasoned about: - -* Predicate pushdown. The fold is splittable by ``splitConjunctivePredicates``, - and ``Optimizer`` partitions on ``cond.deterministic && !cond.throwable``, - which ``RaiseError`` fails to set. So a guard can be pushed below a join and - run on rows the join drops. Measured on a PURE chain with no ``and`` / ``or``: - ``lambda a, b, c: a < b < c`` over a join, ``select`` returns False correctly - but ``filter`` raises USER_RAISED_EXCEPTION where Python returns False. Fixed - by ``RaiseError`` overriding ``Expression.throwable`` (SPARK-58627), which is - a prerequisite for this lowering rather than a nice-to-have. -* Subexpression elimination. ``spark.sql.subexpressionElimination.enabled`` +That short-circuit is an EVALUATION-time property only, and it does not by itself +survive the optimizer. Three divergences, all measured rather than reasoned about: + +* Predicate pushdown -- CLOSED by SPARK-58627. ``Optimizer`` partitions on + ``cond.deterministic && !cond.throwable``, and the fold is splittable by + ``splitConjunctivePredicates``, so while ``RaiseError`` still reported + non-throwable a guard could be pushed below a join and run on rows the join + drops. That hit a PURE chain with no ``and`` / ``or``: + ``lambda a, b, c: a < b < c`` over a join returned False correctly under + ``select`` but raised USER_RAISED_EXCEPTION under ``filter``. ``RaiseError`` + now overrides ``throwable``, and the same repro returns Python's answer. + ``test_udf_transpile_boolop_short_circuits`` is the regression test. +* Subexpression elimination -- OPEN. ``spark.sql.subexpressionElimination.enabled`` defaults true while ``...skipForShortcutExpr`` defaults FALSE, so a conjunct living only in a short-circuited position is hoisted and evaluated unconditionally. Measured on a plain ``select``, no join and no filter: ``(a < b < c) or (d > 0 and b < c)`` on ``(5, 1, None, -1)`` raises where Python returns False, and setting ``skipForShortcutExpr=true`` alone fixes it. - ``Expression.throwable`` does NOT help here -- this rule never consults it. -* NaN. Spark orders NaN above every value and treats ``NaN = NaN`` as true where - Python does neither (SPARK-58781). + ``Expression.throwable`` does NOT help -- this rule never consults it, so + SPARK-58627 does not cover this one. +* NaN -- OPEN. Spark orders NaN above every value and treats ``NaN = NaN`` as + true where Python does neither (SPARK-58781). A chain duplicates every interior operand by construction, so it manufactures -the common subexpressions the second bullet needs. Treat both of the first two -as open, not theoretical. +the common subexpressions the second bullet needs. Treat that one as live. """ import ast From 92678505fc7d951eea0961dd40023562db530be9 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 26 Sep 2026 04:25:51 +0000 Subject: [PATCH 10/18] [SPARK-58456][PYTHON][SQL][MINOR] Clear the category cache between direct lowerings in the caps test `test_udf_transpile_lowering_caps` drives `_convert_chunk` directly, reusing one `CatalystTranspiler` across separately parsed sources. SPARK-59028's `_category_cache` keys on `id(node)`, and CPython reuses a collected node's id for a node of a different type (measured: 16 such reuses over 200 parses), so a chain's int `Constant` could inherit an earlier `Compare`'s cached "bool" and the lowering would refuse. Reset the cache with the budget. Production is unaffected: `_transpile_from_ast` clears the cache on entry next to the `_param_categories` write it guards, `_get_transpilers` builds a transpiler per UDF, and the caller holds the AST alive for the whole lowering. Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- python/pyspark/sql/tests/test_udf_transpile_unit.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index 6406f7119dfd0..a40239883e265 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -791,9 +791,13 @@ def compare_node(src): return _ast.parse(src, mode="eval").body def lower(src): - # `_transpile_from_ast` is what normally zeroes the budget; reset it - # here so each case is charged on its own. + # `_transpile_from_ast` is what normally zeroes the per-lowering state; + # reset it here so each case is charged on its own. The category cache + # has to go too: it keys on `id(node)`, and CPython reuses the id of a + # collected node for one of a different type, so a cache kept across + # separately parsed sources hands back a stale category. transpiler._lowered_comparisons = 0 + transpiler._category_cache = {} return transpiler._convert_chunk(["x"], compare_node(src)) def chain_src(n): From a334662c9b11aed96bc865cd82f5ffaf87819999 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 26 Sep 2026 05:31:11 +0000 Subject: [PATCH 11/18] [SPARK-58456][PYTHON][SQL] Drop the non-determinism RED tests SPARK-58626 made obsolete The branch carried two deliberately-red tests pinning a real wrong-results bug: a non-deterministic ARGUMENT was copied into every `_udf_param_N` occurrence, so a chain or an `and` / `or`, which evaluates later operands only on rows where earlier ones held, saw the copies desynchronize -- `50 < x < 60` over `rand(7) * 100` came out ~30% true where Python is ~10%. SPARK-58626 fixed that bug, in a different place than the tests demanded: `_udf_param_N` is now a reference, and `Optimizer.mustPreEvaluate` / `cheapToRepeat` give an argument the body reads more than once one column per row, or decline to lower the call. So both tests now assert the wrong remedy -- a fallback that is no longer the right answer. Measured before deleting rather than replacing with a retargeted version. The chain over a non-deterministic argument is one draw shared through one column, structurally identical to the shapes `test_udf_transpile_evaluates_inputs_once` and `test_udf_transpile_shares_a_nondeterministic_argument_in_a_predicate` already cover: Filter ((50.0 < _udf_param_5) AND (_udf_param_5 < 60.0)) +- Project [id, (rand(7) * 100.0) AS _udf_param_5] 195/2000 rows, i.e. 9.75% against Python's 10%. The chain's `And` looked like it might reach a path a single comparison cannot, being splittable by `splitConjunctivePredicates` -- but the conjuncts stay put: the non-determinism that makes the argument need a column is exactly what stops `PushPredicateThroughNonJoin` pushing a Filter through the Project holding it. The two conditions are mutually exclusive, so there is nothing left to assert that would not duplicate existing coverage. The Scala control goes with it: the suite's first test already asserts a strict superset on a richer fixture (two options, numeric kept and string dropped, plus `optionInputCategories.isEmpty`), so the control only existed as the negative half of the red pair. The file is now identical to master's. Also fixes two comments the merge left stale: the `in` / `not in` refusal cited SPARK-56925, which shipped as bare `if x` support and is unrelated, and the comparison fold still described the pre-SPARK-58626 substitution. Co-authored-by: Holden Karau Co-authored-by: Holden Karau --- .../sql/tests/test_udf_transpile_unit.py | 51 ------------------- python/pyspark/sql/transpile.py | 29 ++++------- ...solveTranspiledPythonUDFOptionsSuite.scala | 35 +------------ 3 files changed, 12 insertions(+), 103 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index a40239883e265..1c04562ff9f2d 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -2398,57 +2398,6 @@ def test_udf_transpile_plan_elision(self): self.assertEqual(sorted(r[0] for r in fdf.collect()), [1, 7]) self.assertEqual(0, self._eval_python_count(fdf)) - def test_udf_transpile_nondeterministic_argument_falls_back(self): - # RED UNTIL the companion analyzer fix lands -- deliberately, to pin a real - # wrong-results bug rather than tolerate it silently. Do not skip or xfail. - # - # A non-deterministic ARGUMENT must suppress transpilation at the call site - # even though the UDF itself transpiled fine. (Distinct from - # test_udf_transpile_skips_nondeterministic, which covers a UDF marked - # non-deterministic.) `resolveUDFParams` substitutes the argument at every - # `_udf_param_N` occurrence -- four of them for a two-link chain -- so a - # non-deterministic one becomes several independent expressions with - # independent state. Since `and` / `or` and chains evaluate later copies - # only on the rows where earlier links held, the copies desynchronize: - # `50 < x < 60` over `rand(7) * 100` is ~30% true instead of ~10%. - # - # The fix belongs in ResolveTranspiledPythonUDFOptions. It cannot go in - # Python (no call site exists at UDF construction) nor in the JVM builder - # (`rand(7)` is still an UnresolvedFunction there and reports - # deterministic=true). Assert on the plan, not on observed proportions: a - # statistical assertion would be flaky, and keeping EvalPython means the - # interpreted UDF runs, which is correct by construction. - from pyspark.sql.functions import col, rand - - band = lambda x: 50 < x < 60 # noqa: E731 - band_or = lambda x: x < 50 or x > 60 # noqa: E731 - with self.sql_conf(_TRANSPILE_ON): - for label, func in (("50 < x < 60", band), ("x < 50 or x > 60", band_or)): - # subTest so the chain failing does not hide the `or` case, which - # has the same bug by the same mechanism. - with self.subTest(func=label): - pudf = UserDefinedFunction(func, BooleanType()) - self.assertTrue(pudf.transpiled, f"{label} should transpile on its own") - base = self.spark.range(0, 8).select((col("id") * 10).alias("a")) - - # Deterministic argument -> transpiled and elided. - det = base.select(pudf("a")) - self.assertEqual( - 0, - self._eval_python_count(det), - f"{label}: a deterministic argument should still be elided", - ) - - # Non-deterministic argument -> must fall back. - nondet = base.select(pudf(rand(7) * 100)) - self.assertEqual( - 1, - self._eval_python_count(nondet), - f"{label}: a non-deterministic argument must NOT be transpiled", - ) - # And it still produces results. - self.assertEqual(8, len(nondet.collect())) - def test_udf_transpile_config_toggle_no_stale_nodes(self): # Built with the flags on, executed with them off -> clean fallback to # interpreted Python (the optimizer drops the transpiled node), no error. diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 8e908280e2564..337c0707082f3 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -856,8 +856,8 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: case ast.Compare(left, ops, comps): # `in` / `not in` are container membership, not a chainable value # comparison, and need their own NULL analysis: Python's `in` is - # equality-based and never raises on a `None` probe. Deferred to - # SPARK-56925, so refuse them and fall back. + # equality-based and never raises on a `None` probe. Not lowered + # yet, so refuse them and fall back. if any(isinstance(o, (ast.In, ast.NotIn)) for o in ops): raise UnsupportedOperationException( "`in` / `not in` is not supported by the transpiler; the " @@ -872,22 +872,15 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # Python evaluates `a OP1 b OP2 c` as `(a OP1 b) and (b OP2 c)` # with `b` evaluated once, so the fold below re-lowers each # interior operand once per adjacent pair. Safe for the operand - # grammar we accept, all of it deterministic and side-effect free - # -- but NOT for the call-site arguments the JVM's - # `resolveUDFParams` substitutes into `_udf_param_N`, where a - # non-deterministic argument becomes several independent - # expressions that desynchronize (`50 < x < 60` over - # `rand(7) * 100` is ~30% true where Python is ~10%). That has to - # be fixed in the analyzer, which is the first place a resolved - # argument exists. - # - # Be precise about the blast radius. `and` / `or` were already - # wrong this way, their operands being separate subtrees. A SINGLE - # comparison is not: its duplicate operand appears only inside an - # `IS NULL` test, whose outcome does not depend on the value drawn, - # so `50 < x` over `rand()` samples correctly. Chains previously - # fell back to interpreted Python, so lowering them is what newly - # exposes them to this. + # grammar we accept, all of it deterministic and side-effect free. + # It does mean the body READS a parameter once per adjacent pair, + # which SPARK-58626 is what makes safe: `_udf_param_N` is a + # reference, and `ConvertToCatalyst` gives an argument read more + # than once one column per row -- or declines to lower the call + # where no column fits -- so a non-deterministic argument is drawn + # once and every read agrees. While each read still carried its own + # copy of the argument, `50 < x < 60` over `rand(7) * 100` was ~30% + # true where Python is ~10%. # # Folding into nested CASE WHENs instead -- atomic to predicate # splitting -- was rejected: it costs pushdown for every diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala index 7a6f6a5ae7ba5..f9295eb17b00c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.analysis import org.apache.spark.api.python.PythonEvalType import org.apache.spark.sql.catalyst.dsl.expressions._ -import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Concat, Expression, GreaterThan, Literal, Multiply, PythonUDF, Rand, TranspiledPythonUDF} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Concat, Expression, Literal, PythonUDF, TranspiledPythonUDF} import org.apache.spark.sql.catalyst.plans.PlanTest import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project} import org.apache.spark.sql.types.LongType @@ -92,37 +92,4 @@ class ResolveTranspiledPythonUDFOptionsSuite extends PlanTest { val pruned = prune(node, LocalRelation(a)) assert(pruned.transpiledOptions == List(onlyOpt)) } - - test("empties the options when an argument is non-deterministic") { - // RED UNTIL the companion fix to this rule lands -- deliberately, to pin a - // real wrong-results bug rather than tolerate it silently. Do not ignore it. - // - // The builder substitutes the argument expression at EVERY `_udf_param_N` - // occurrence, so a non-deterministic argument becomes several independent - // expressions with independent state. Since `and` / `or` and chained - // comparisons evaluate later copies only on rows where earlier links held, - // the copies desynchronize. The rule must drop every option so - // ConvertToCatalyst falls back to the Python UDF, which evaluates each - // argument exactly once per row. It has to be checked here rather than in the - // builder: at call-construction time `rand(7)` is still an - // `UnresolvedFunction` and reports deterministic = true. - val a = $"a".double - val arg = Multiply(Rand(Literal(7L)), Literal(100.0)) - assert(!arg.deterministic, "the fixture argument must be non-deterministic") - val opt = GreaterThan(arg, Literal(50.0)) - val node = TranspiledPythonUDF("udf", pyUDF(Seq(arg)), List(opt), List(List("numeric"))) - val pruned = prune(node, LocalRelation(a)) - assert(pruned.transpiledOptions.isEmpty) - assert(pruned.optionInputCategories.isEmpty) - } - - test("keeps options when the argument is deterministic (control for the check above)") { - // Same shape as the non-deterministic case but with a plain column argument, so - // the guard must not fire and normal category matching applies. - val a = $"a".double - val opt = GreaterThan(a, Literal(50.0)) - val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)), List(opt), List(List("numeric"))) - val pruned = prune(node, LocalRelation(a)) - assert(pruned.transpiledOptions == List(opt)) - } } From 1d3555d7965d7494ed92c8af2baa3f87b4e809fd Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 26 Sep 2026 07:22:36 +0000 Subject: [PATCH 12/18] [SPARK-58456][PYTHON][SQL] Correct the optimizer-divergence notes against measurement Comments only, no behavior change. Two of the module docstring's claims about which optimizer rules break the `&` / `|` fold were measurably wrong, and a code review of the merge turned up two more divergences than the list had. * "Predicate pushdown -- CLOSED by SPARK-58627" is true for joins and false one operator over. Only the three join sites test `!cond.throwable`. The Aggregate arm says so outright -- "We can push down deterministic predicate through Aggregate, including throwable predicate" -- and Window / Union / EventTimeWatermark partition on `_.deterministic` alone. Measured: a `groupBy(...).agg(...).filter(sa < b < c)` splits the fold, pushes the `b < c` guard below the `Aggregate` and raises. That arm's justification reasons about ROWS, which is sound, and is silent about conjunct ORDER, which is the only thing a guard in a fold relies on. * "setting `skipForShortcutExpr=true` alone fixes it" is false. There is a second independently sufficient cause: `BooleanSimplification`, which consults neither `deterministic` nor `throwable`. On `(a < b < c) or (d < b < c)` over `(5, 1, None, 7)`, each of `skipForShortcutExpr=true`, `subexpressionElimination.enabled=false` and excluding `BooleanSimplification` still raises alone; only excluding the rule AND disabling a subexpression-elimination lever returns Python's False. * New, and the worst of the set because it is a wrong answer rather than a spurious error: `BooleanSimplification`'s absorption case `(a && b) || b => b` DROPS a guard instead of moving it. `(a < b < c) or (b < c)` on `(None, 0, 1)` returns True where Python raises TypeError, with one of the two `raise_error` guards gone from the plan. Found by checking the direction our differential tests tolerate -- see below. None of this is specific to chained comparisons, which the docstring now says: the `and` form of each repro predates the chain lowering and diverges identically, so a chain is new syntax reaching an existing hole. Recorded rather than fixed; fixing belongs in the optimizer and wants its own ticket. Also: the `__init__` seed comment justified itself by a test that no longer needs it (the caps test resets the counter itself) and the field lacked a type annotation. Retargeted at the real hazard -- a subclass overriding `_transpile_from_ast` must reset every per-lowering field, since a stale budget refuses part-way through the variant list and silently emits fewer options. And `_is_never_null_boolean`'s docstring named only `_is_definitely_boolean` for if/ternary tests, which describes a gate that never fires for the bare `if x` SPARK-56925 shipped; the NULL rationale holds, the mechanism needed the `_truthiness_col` path added. Co-authored-by: Holden Karau Co-authored-by: Holden Karau Co-authored-by: Cursor --- .../sql/tests/test_udf_transpile_unit.py | 2 +- python/pyspark/sql/transpile.py | 121 +++++++++++------- 2 files changed, 79 insertions(+), 44 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index 1c04562ff9f2d..3947ee4df8079 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -779,9 +779,9 @@ def test_udf_transpile_lowering_caps(self): from pyspark.errors import UnsupportedOperationException from pyspark.sql.transpile import ( - CatalystTranspiler, _MAX_COMPARISON_OPS, _MAX_LOWERED_COMPARISONS, + CatalystTranspiler, ) transpiler = CatalystTranspiler() diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 337c0707082f3..e37be90c1a144 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -35,6 +35,70 @@ functions with more than three untyped parameters only emit the all-numeric and all-string variants. +Constructs that short-circuit in Python -- ``and`` / ``or`` and chained +comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``. That fold is +what is meant to keep the ``raise_error`` NULL guard in a later operand from +firing on rows Python never reaches, since ``And`` skips its right operand when +the left is exactly ``false`` (and ``Or`` when the left is exactly ``true``) in +both ``eval`` and ``doGenCode``. Note the "exactly": a NULL left operand still +evaluates the right. ``and`` / ``or`` operands must therefore be provably boolean +and non-NULL (see ``_is_never_null_boolean``); chain links need no such gate, +because every comparison lowering is total. + +That short-circuit is an EVALUATION-time property only, and it does not by itself +survive the optimizer. Four divergences, all measured rather than reasoned about. +None of them is specific to chained comparisons: the ``and`` form of each repro +(``a < b and b < c`` for ``a < b < c``) predates this lowering and diverges +identically, so a chain is new syntax reaching an existing hole. + +* Predicate pushdown through a JOIN -- CLOSED by SPARK-58627. ``Optimizer`` + partitions on ``cond.deterministic && !cond.throwable``, and the fold is + splittable by ``splitConjunctivePredicates``, so while ``RaiseError`` still + reported non-throwable a guard could be pushed below a join and run on rows the + join drops. That hit a PURE chain with no ``and`` / ``or``: + ``lambda a, b, c: a < b < c`` over a join returned False correctly under + ``select`` but raised USER_RAISED_EXCEPTION under ``filter``. ``RaiseError`` + now overrides ``throwable``, and the same repro returns Python's answer. + ``test_udf_transpile_boolop_short_circuits`` is the regression test. +* Predicate pushdown through an AGGREGATE or WINDOW -- OPEN, and the reason the + bullet above says "through a JOIN". Only the three join sites test + ``!cond.throwable``; the Aggregate arm is explicit that it does not + ("We can push down deterministic predicate through Aggregate, including + throwable predicate"), and the Window / Union / EventTimeWatermark arms + partition on ``_.deterministic`` alone. Measured: + ``groupBy("b", "c").agg(sum("a").alias("sa")).filter(lambda: sa < b < c)`` + splits the fold, pushing the ``b < c`` guard BELOW the ``Aggregate`` while + ``sa < b`` stays above, and raises. The Aggregate arm's own reasoning is about + ROWS -- grouping keys cannot gain distinct values -- and says nothing about + conjunct ORDER, which is the only thing a guard in a fold relies on. +* BooleanSimplification -- OPEN, and it breaks the fold two different ways. It + consults neither ``deterministic`` nor ``throwable``. Common-factor elimination + for disjunction (``(a && b) || (a && c) => a && (b || c)``) hoists a shared + guard to the FRONT of a conjunction, where it evaluates unconditionally: + ``(a < b < c) or (d < b < c)`` on ``(5, 1, None, 7)`` raises where Python + returns False. Its absorption case (``(a && b) || b => b``) is worse, because + it drops a guard instead of moving it: ``(a < b < c) or (b < c)`` on + ``(None, 0, 1)`` returns True where Python raises TypeError, with only one of + the two ``raise_error`` guards surviving. A silently wrong answer, not a + spurious error. +* Subexpression elimination -- OPEN. ``spark.sql.subexpressionElimination.enabled`` + defaults true while ``...skipForShortcutExpr`` defaults FALSE, so a conjunct + living only in a short-circuited position is hoisted and evaluated + unconditionally. ``Expression.throwable`` does NOT help -- this rule never + consults it, so SPARK-58627 does not cover this one. +* NaN -- OPEN. Spark orders NaN above every value and treats ``NaN = NaN`` as + true where Python does neither (SPARK-58781). + +Subexpression elimination and BooleanSimplification are INDEPENDENTLY sufficient, +which is worth knowing before reaching for a conf to prove a theory. On the +``(a < b < c) or (d < b < c)`` repro, each of ``skipForShortcutExpr=true``, +``subexpressionElimination.enabled=false`` and excluding ``BooleanSimplification`` +still raises on its own; only excluding ``BooleanSimplification`` AND disabling +one of the subexpression-elimination levers returns Python's answer. + +A chain duplicates every interior operand by construction, so it manufactures the +common subexpressions and the common factors these rules need. Treat them as live. + ``len(s)`` lowers to Catalyst ``length`` for a string operand (SPARK-55214). Other ``len`` arguments (numbers, lists, ...) stay interpreted Python. ``len(None)`` raises (via ``raise_error`` on NULL), matching CPython. @@ -90,42 +154,6 @@ expression. A draw is never inlined, because those rules check determinism. An argument the body never reads is not computed at all. - -Constructs that short-circuit in Python -- ``and`` / ``or`` and chained -comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``. That fold is -what is meant to keep the ``raise_error`` NULL guard in a later operand from -firing on rows Python never reaches, since ``And`` skips its right operand when -the left is exactly ``false`` (and ``Or`` when the left is exactly ``true``) in -both ``eval`` and ``doGenCode``. Note the "exactly": a NULL left operand still -evaluates the right. ``and`` / ``or`` operands must therefore be provably boolean -and non-NULL (see ``_is_never_null_boolean``); chain links need no such gate, -because every comparison lowering is total. - -That short-circuit is an EVALUATION-time property only, and it does not by itself -survive the optimizer. Three divergences, all measured rather than reasoned about: - -* Predicate pushdown -- CLOSED by SPARK-58627. ``Optimizer`` partitions on - ``cond.deterministic && !cond.throwable``, and the fold is splittable by - ``splitConjunctivePredicates``, so while ``RaiseError`` still reported - non-throwable a guard could be pushed below a join and run on rows the join - drops. That hit a PURE chain with no ``and`` / ``or``: - ``lambda a, b, c: a < b < c`` over a join returned False correctly under - ``select`` but raised USER_RAISED_EXCEPTION under ``filter``. ``RaiseError`` - now overrides ``throwable``, and the same repro returns Python's answer. - ``test_udf_transpile_boolop_short_circuits`` is the regression test. -* Subexpression elimination -- OPEN. ``spark.sql.subexpressionElimination.enabled`` - defaults true while ``...skipForShortcutExpr`` defaults FALSE, so a conjunct - living only in a short-circuited position is hoisted and evaluated - unconditionally. Measured on a plain ``select``, no join and no filter: - ``(a < b < c) or (d > 0 and b < c)`` on ``(5, 1, None, -1)`` raises where - Python returns False, and setting ``skipForShortcutExpr=true`` alone fixes it. - ``Expression.throwable`` does NOT help -- this rule never consults it, so - SPARK-58627 does not cover this one. -* NaN -- OPEN. Spark orders NaN above every value and treats ``NaN = NaN`` as - true where Python does neither (SPARK-58781). - -A chain duplicates every interior operand by construction, so it manufactures -the common subexpressions the second bullet needs. Treat that one as live. """ import ast @@ -300,8 +328,11 @@ def _is_never_null_boolean(node: ast.AST) -> bool: bare ``None`` operand is not enough because a ternary can evaluate to one (``(True if c else None) and x``). - ``not`` and if/ternary *tests* keep using ``_is_definitely_boolean``: they - wrap the operand in ``coalesce``, so they handle NULL explicitly. + ``not`` and if/ternary *tests* do not need this: both wrap the operand in + ``coalesce``, so both handle NULL explicitly. ``not`` gates on + ``_is_definitely_boolean``; an if/ternary test that fails that gate falls to + ``_truthiness_col`` (SPARK-56925), which coalesces in every category, so a + bare ``if x`` is total too. This answers "would it be NULL?", not "can it be lowered at all" -- it says True for shapes ``_convert_chunk`` goes on to refuse, such as ``x in "abc"``. @@ -355,14 +386,18 @@ class CatalystTranspiler(AbstractTranspiler): variety = "catalyst" - # Class-level default so the lowering helpers can be driven directly (tests - # call ``_convert_chunk`` without going through ``_transpile_from_ast``, which - # is what normally resets this). - _lowered_comparisons = 0 - def __init__(self) -> None: self._param_categories: dict[int, str] = {} self._category_cache: dict[int, str] = {} + # Every field above is per-lowering state that ``_transpile_from_ast`` + # resets on entry; they are seeded here so an instance is usable before + # the first lowering. A subclass overriding that seam -- it is the + # documented override point for + # ``spark.sql.experimental.optimizer.pyTranspilers`` -- has to reset ALL + # of them, not just the cache: a stale counter is a monotonic budget that + # refuses part-way through the variant list and silently emits fewer + # options. + self._lowered_comparisons: int = 0 # TODO (SPARK-55218): handle implicit-None return bodies like # ``def f(x): x + x`` -- no return statement means return None; From 5dcfb25b05439a38ff1f01f528fd53fd1d1847bb Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Fri, 2 Oct 2026 15:17:01 +0000 Subject: [PATCH 13/18] [SPARK-58456][PYTHON][SQL][MINOR] Fix the divergence count in the module docstring The intro says "Four divergences" but the list has five bullets, and the fifth (NaN) is a value-level difference, not an optimizer rule. Say so. Co-authored-by: Holden Karau Co-authored-by: Holden Karau Co-authored-by: Cursor --- python/pyspark/sql/transpile.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index e37be90c1a144..bcbc50b841ada 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -46,8 +46,9 @@ because every comparison lowering is total. That short-circuit is an EVALUATION-time property only, and it does not by itself -survive the optimizer. Four divergences, all measured rather than reasoned about. -None of them is specific to chained comparisons: the ``and`` form of each repro +survive the optimizer. Four optimizer rules diverge, all measured rather than +reasoned about, and the last bullet adds the NaN value-level difference. None of +them is specific to chained comparisons: the ``and`` form of each repro (``a < b and b < c`` for ``a < b < c``) predates this lowering and diverges identically, so a chain is new syntax reaching an existing hole. From aa42c9b2388c2ba81ff0911e0fd030985b51ebfe Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Fri, 2 Oct 2026 21:13:26 +0000 Subject: [PATCH 14/18] [SPARK-58456][PYTHON][SQL][MINOR] Trim the chained-comparison hypothesis seed edges Reduce _LONG_TRIPLE_EDGES from 14 hand-listed tuples to 9 pulled from the single-value edges, dropping the redundant 2-of-3-None-position cases (covered by the single-None cases plus all-None) and the all-equal tuples (covered by random generation plus the monotonic orderings). Trim the _seed_tuple_examples docstring and drop the over-defensive arity assert; every seed matches its keys today, and a too-short seed fails loudly at decoration time on its own. Co-authored-by: Holden Karau Co-authored-by: Holden Karau Co-authored-by: Cursor --- .../tests/test_udf_transpile_hypothesis.py | 24 +++++++------------ 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py index 24b56f227a5cd..1cce17d5e79dc 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -206,22 +206,21 @@ def _env_opts_in(value: Optional[str]) -> bool: (1, -_LONG_BOUND), ) - # Triple edges for the chained-comparison tests. The NULL patterns are the - # point: a chain short-circuits, so whether a NULL operand's raise guard is - # reached depends on the earlier links. + # Triple edges for the chained-comparison tests, pulled from the + # single-value edges above. A chain short-circuits, so the NULL patterns + # -- which position holds the None -- are the point; the boundary orderings + # probe sign/overflow at the type's edges. _LONG_TRIPLE_EDGES = ( + # NULL position: which link's guard fires -- or is short-circuited past + # -- depends on where the None sits. (None, None, None), - (None, None, 0), - (None, 0, None), - (0, None, None), (None, 0, 1), (0, None, 1), (0, 1, None), - (1, 0, None), - (0, 0, 0), + (1, 0, None), # first link false -> None never reached + # Boundary orderings. (0, 1, 2), (2, 1, 0), - (1, 1, 1), (_INT32_MIN, 0, _INT32_MAX), (-_LONG_BOUND, 0, _LONG_BOUND), ) @@ -237,15 +236,10 @@ def wrapper(method): return wrapper def _seed_tuple_examples(tuples, keys=("x", "y", "z")): - """Stack one ``@example`` decorator per seed tuple, zipped onto ``keys``. - - The arity assert matters: ``zip`` stops at the shorter argument, so a - mismatched seed tuple would quietly seed a different input than intended. - """ + """Stack one ``@example`` decorator per seed tuple, zipped onto ``keys``.""" def wrapper(method): for values in reversed(tuples): - assert len(values) == len(keys), f"seed {values!r} does not fit {keys!r}" method = example(**dict(zip(keys, values)))(method) return method From dcba3ad1952e9e5840f1924e2c7578f7a33f4fc5 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Thu, 22 Oct 2026 21:26:04 +0000 Subject: [PATCH 15/18] [SPARK-58456][PYTHON][SQL] Fold chains to CASE WHEN and bound materialized size Address review on #59216: - Fold a chained comparison into a sequential CASE WHEN (CASE WHEN NOT link1 THEN false ... ELSE linkN END) instead of `&`, so splitConjunctivePredicates / BooleanSimplification / subexpression elimination cannot move a later link's NULL guard ahead of the links that gate it. Chains had no lowering before this PR, so the atomic fold loses no pushdown versus the merge target; `and` / `or` keep `&` / `|`. - Charge the comparison budget by materialized (tree) size, not link count: a right-nested body is ~2d+1 links but ~2^d nodes (3^d with `==`), since each operand Column is referenced more than once and the JVM converter does not memoize. Refuse once a body's materialized weight passes a bound, so pathological nesting falls back at UDF construction. - Assert zero EvalPython in the SPARK-58627 join test so a JVM discard of the transpiled options cannot hide a regression behind a matching result. - Fix the `in` / `not in` refusal comment (Python `in` semantics depend on the container; `None in "abc"` raises) and drop the stale SPARK-56925 references (it is the bare-`if` truthiness work, not membership). - Refresh the `_is_definitely_boolean` docstring: `and` / `or` now gate on `_is_never_null_boolean`, not this helper. - Add test_udf_transpile_chain_survives_optimizer covering the Aggregate pushdown and BooleanSimplification common-factor / absorption shapes. Co-authored-by: Holden Karau Co-authored-by: Holden Karau Co-authored-by: Cursor --- .../sql/tests/test_udf_transpile_unit.py | 150 +++++++++- python/pyspark/sql/transpile.py | 267 ++++++++++++++---- 2 files changed, 351 insertions(+), 66 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index 3947ee4df8079..173d47ad33d5d 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -543,6 +543,76 @@ def chained(a, b, c): df.select(pudf("a", "b", "c")).collect() self.assertIn("cannot compare NULL", str(ctx.exception)) + def test_udf_transpile_chain_survives_optimizer(self): + # A chain folds to a sequential CASE WHEN, which is atomic to + # `splitConjunctivePredicates` and `BooleanSimplification`. Before this + # fix the chain folded to `&`, which those rules could split or reorder, + # moving a later link's NULL guard ahead of the links that gate it and + # diverging from Python. Each repro below is an optimizer shape that broke + # the `&` fold; all must return Python's answer (or raise like Python) + # with the chain fully lowered (no EvalPython in the plan). + + # PushDownPredicates / Aggregate: a later link over only grouping keys + # used to be pushed below the Aggregate, where it ran on groups Python + # short-circuits past. `sa < b < c` over groupBy(b,c).agg(sum(a)): the + # group (b=1, c=NULL, sa=5) has `sa < b` false, so Python returns False + # without touching `b < c` (which would raise on the NULL c). + from pyspark.sql.functions import sum as sum_ + + def chain_over_agg(sa, b, c): + return sa < b < c + + with self.sql_conf(_TRANSPILE_ON): + udf = UserDefinedFunction(chain_over_agg, BooleanType()) + self.assertTrue(udf.transpiled, "chain over agg should transpile") + agg = ( + self.spark.createDataFrame( + [(1, 1, None), (5, 1, None), (2, 2, 3)], "a long, b long, c long" + ) + .groupBy("b", "c") + .agg(sum_("a").alias("sa")) + ) + filtered = agg.filter(udf("sa", "b", "c")) + # Python: row (sa=5, b=1, c=None) -> 5 < 1 is False -> False (no raise). + # row (sa=2, b=2, c=3) -> 2 < 2 is False -> False. + # No row passes, and crucially no raise leaks out of the NULL c. + self.assertEqual([], filtered.select("b", "c").collect()) + self.assertEqual(0, self._eval_python_count(filtered), "chain must stay lowered") + + # BooleanSimplification common-factor: `(a < b < c) or (d < b < c)` on + # (5, 1, None, 7). The `&` fold let common-factor elimination hoist the + # shared `b < c` guard to the front, where the NULL c raised where Python + # returns False. The CASE WHEN fold is atomic, so the shared link cannot + # be hoisted out of either chain. + def or_shared_chain(a, b, c, d): + return (a < b < c) or (d < b < c) + + with self.sql_conf(_TRANSPILE_ON): + udf = UserDefinedFunction(or_shared_chain, BooleanType()) + self.assertTrue(udf.transpiled, "or of shared chains should transpile") + df = self.spark.createDataFrame([(5, 1, None, 7)], "a long, b long, c long, d long") + # Python: (5 < 1) False -> first chain False; (7 < 1) False -> second + # chain False; or -> False. The NULL c is never reached. + self.assertEqual([False], [r[0] for r in df.select(udf("a", "b", "c", "d")).collect()]) + projected = df.select(udf("a", "b", "c", "d")) + self.assertEqual(0, self._eval_python_count(projected), "or of chains stays lowered") + + # BooleanSimplification absorption: `(a < b < c) or (b < c)` on + # (None, 0, 1). The `&` fold's absorption (`(a && b) || b => b`) dropped + # the first link's guard, returning True where Python raises TypeError on + # `None < 0`. The CASE WHEN fold keeps both links, so the first chain + # reaches its NULL guard and raises like Python. + def or_absorb_chain(a, b, c): + return (a < b < c) or (b < c) + + with self.sql_conf(_TRANSPILE_ON): + udf = UserDefinedFunction(or_absorb_chain, BooleanType()) + self.assertTrue(udf.transpiled, "or with absorption shape should transpile") + df = self.spark.createDataFrame([(None, 0, 1)], "a long, b long, c long") + with self.assertRaises(Exception) as ctx: + df.select(udf("a", "b", "c")).collect() + self.assertIn("cannot compare NULL", str(ctx.exception)) + def test_udf_transpile_boolop_short_circuits(self): # The `and` / `or` counterpart of the chain short-circuit: a NULL guard in # the second operand must not fire when the first already decided the @@ -589,12 +659,18 @@ def either_positive(a, c): left = self.spark.createDataFrame([(1, -5)], "k long, a long") right = self.spark.createDataFrame([(1, None)], "k long, c long") joined = left.join(right, "k") - self.assertEqual([False], [r[0] for r in joined.select(and_udf("a", "c")).collect()]) + joined_select = joined.select(and_udf("a", "c")) + self.assertEqual([False], [r[0] for r in joined_select.collect()]) + # Results alone do not prove the lowering ran -- interpreted Python + # returns the same `[False]`. Assert zero EvalPython nodes so a JVM + # discard of the transpiled options cannot hide a regression. + self.assertEqual(0, self._eval_python_count(joined_select), "joined select lowered") # Filtering must return Python's answer rather than raising. This is # the assertion SPARK-58627 fixed; see above. filtered = joined.filter(and_udf("a", "c")) self.assertEqual([], filtered.select("k").collect()) + self.assertEqual(0, self._eval_python_count(filtered), "joined filter lowered") def test_udf_transpile_chain_operator_mix_and_eq_nulls(self): # Mixed operators in one chain, and `x == y == z` over every NULL pattern. @@ -659,11 +735,13 @@ def test_udf_transpile_comparison_fallbacks(self): # "did not transpile" alone is not enough to prove the fallback works. import warnings as _warnings - # `in` / `not in` are refused wholesale (SPARK-56925), so these cases all - # take the same branch; they are kept because the reasons differ. `in` on a - # str is substring containment rather than equality, a container of runtime - # values would need runtime NULL analysis, and `in` can appear as a link in - # a chain (`a < b in (1, 2)` is one Compare), which the refusal also covers. + # `in` / `not in` are refused wholesale, so these cases all take the same + # branch; they are kept because the reasons differ. `in` on a str is + # substring containment rather than equality, a container of runtime + # values would need runtime NULL analysis, and `in` can appear as a link + # in a chain (`a < b in (1, 2)` is one Compare), which the refusal also + # covers. Membership needs its own NULL analysis and is deferred to its + # own ticket (SPARK-56925 is the unrelated bare-`if` truthiness work). str_in = lambda x: x in "abc" # noqa: E731 runtime_in = lambda x, y: x in (y, 2) # noqa: E731 chained_in = lambda a, b: a < b in (1, 2) # noqa: E731 @@ -780,6 +858,7 @@ def test_udf_transpile_lowering_caps(self): from pyspark.errors import UnsupportedOperationException from pyspark.sql.transpile import ( _MAX_COMPARISON_OPS, + _MAX_LOWERED_COMPARISON_WEIGHT, _MAX_LOWERED_COMPARISONS, CatalystTranspiler, ) @@ -795,9 +874,13 @@ def lower(src): # reset it here so each case is charged on its own. The category cache # has to go too: it keys on `id(node)`, and CPython reuses the id of a # collected node for one of a different type, so a cache kept across - # separately parsed sources hands back a stale category. + # separately parsed sources hands back a stale category. The weight + # cache keys on `id(node)` for the same reason. transpiler._lowered_comparisons = 0 + transpiler._lowered_comparison_weight = 0 + transpiler._compare_depth = 0 transpiler._category_cache = {} + transpiler._comparison_weight_cache = {} return transpiler._convert_chunk(["x"], compare_node(src)) def chain_src(n): @@ -839,9 +922,58 @@ def nested_src(depth): ) with self.assertRaises(UnsupportedOperationException) as ctx: lower(deep) - self.assertIn(str(_MAX_LOWERED_COMPARISONS), str(ctx.exception)) + # This shape's LINK count also grows with nesting (a chain re-lowers each + # interior operand once per adjacent pair), so it can trip the link budget + # too; either budget refusing is correct. Assert the materialized-weight + # bound fires, since that is the bound this shape is now judged by. + self.assertIn(str(_MAX_LOWERED_COMPARISON_WEIGHT), str(ctx.exception)) + + # ---- Right-nested body: linear links, exponential materialized size ---- + # `(x < 0) < ((x < 0) < (...))` nests in ONE operand of a SINGLE + # comparison per level, so the link count grows linearly (~d + 1) while + # the materialized size doubles per level (each operand Column is + # referenced twice and the JVM converter does not memoize). The link + # budget above does NOT fire here -- depth 20 is 21 links, well under + # `_MAX_LOWERED_COMPARISONS` -- so the materialized-weight budget is what + # has to refuse it. + def right_nested_src(depth): + src = "(x < 0)" + for _ in range(depth): + src = "(x < 0) < (%s)" % src + return src + + # Shallow right-nesting stays inside the weight budget and still lowers. + lower(right_nested_src(2)) + + deep_right = right_nested_src(20) + right_node = compare_node(deep_right) + # The link count is linear and well under the link budget, proving the + # weight budget -- not the link budget -- is what refuses this shape. + self.assertLess( + sum(len(n.ops) for n in _ast.walk(right_node) if isinstance(n, _ast.Compare)), + _MAX_LOWERED_COMPARISONS, + "right-nest link count must stay under the link budget, so the weight " + "budget is what fires", + ) + with self.assertRaises(UnsupportedOperationException) as ctx: + lower(deep_right) + self.assertIn(str(_MAX_LOWERED_COMPARISON_WEIGHT), str(ctx.exception)) + + # `==` references each operand three times (vs twice for `<`), so an + # equality right-nest compounds faster; the same depth already refuses. + def eq_right_nested_src(depth): + src = "(x == 1)" + for _ in range(depth): + src = "(x == 1) == (%s)" % src + return src + + with self.assertRaises(UnsupportedOperationException) as ctx: + lower(eq_right_nested_src(20)) + self.assertIn(str(_MAX_LOWERED_COMPARISON_WEIGHT), str(ctx.exception)) - # `in` / `not in` are refused outright, deferred to SPARK-56925. + # `in` / `not in` are refused outright; membership needs its own NULL + # analysis and is deferred to its own ticket (SPARK-56925 is the + # unrelated bare-`if` truthiness work). with self.assertRaises(UnsupportedOperationException) as ctx: lower("x in (1, 2)") self.assertIn("`in` / `not in`", str(ctx.exception)) diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index bcbc50b841ada..51c6a5831d61c 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -36,52 +36,60 @@ all-numeric and all-string variants. Constructs that short-circuit in Python -- ``and`` / ``or`` and chained -comparisons such as ``0 <= x <= 100`` -- fold with ``&`` / ``|``. That fold is -what is meant to keep the ``raise_error`` NULL guard in a later operand from -firing on rows Python never reaches, since ``And`` skips its right operand when -the left is exactly ``false`` (and ``Or`` when the left is exactly ``true``) in -both ``eval`` and ``doGenCode``. Note the "exactly": a NULL left operand still -evaluates the right. ``and`` / ``or`` operands must therefore be provably boolean -and non-NULL (see ``_is_never_null_boolean``); chain links need no such gate, -because every comparison lowering is total. - -That short-circuit is an EVALUATION-time property only, and it does not by itself -survive the optimizer. Four optimizer rules diverge, all measured rather than -reasoned about, and the last bullet adds the NaN value-level difference. None of -them is specific to chained comparisons: the ``and`` form of each repro -(``a < b and b < c`` for ``a < b < c``) predates this lowering and diverges -identically, so a chain is new syntax reaching an existing hole. +comparisons such as ``0 <= x <= 100`` -- keep the ``raise_error`` NULL guard in +a later operand from firing on rows Python never reaches. ``and`` / ``or`` fold +with ``&`` / ``|``: ``And`` skips its right operand when the left is exactly +``false`` (and ``Or`` when the left is exactly ``true``) in both ``eval`` and +``doGenCode``. Note the "exactly": a NULL left operand still evaluates the right, +so ``and`` / ``or`` operands must be provably boolean and non-NULL (see +``_is_never_null_boolean``). A chain folds instead into a sequential CASE WHEN +(``CASE WHEN NOT link1 THEN false ... WHEN NOT link(n-1) THEN false ELSE linkN +END``): each later link sits in a branch that runs only after every earlier link +is true, so it short-circuits the same way. Chain links need no non-NULL gate, +because every comparison lowering is total, so ``NOT link`` is never NULL. + +That short-circuit is an EVALUATION-time property only, and it does not by +itself survive the optimizer. The ``&`` / ``|`` fold that ``and`` / ``or`` use +is splittable, and four optimizer rules diverge on it, all measured rather than +reasoned about (the last bullet adds the NaN value-level difference). A chain's +CASE WHEN is atomic -- ``splitConjunctivePredicates`` matches only ``And``, +``BooleanSimplification``'s patterns match ``And`` / ``Or``, and subexpression +elimination recurses only into a CASE WHEN's first condition -- so no rule can +move a later guard ahead of the links that gate it, and a chain predicate moves +only as a whole when pushdown applies. The divergences below are therefore the +``and`` / ``or`` fold's; the ``and`` form of each repro (``a < b and b < c`` for +``a < b < c``) predates this lowering and diverges identically. * Predicate pushdown through a JOIN -- CLOSED by SPARK-58627. ``Optimizer`` - partitions on ``cond.deterministic && !cond.throwable``, and the fold is + partitions on ``cond.deterministic && !cond.throwable``, and the ``&`` fold is splittable by ``splitConjunctivePredicates``, so while ``RaiseError`` still reported non-throwable a guard could be pushed below a join and run on rows the - join drops. That hit a PURE chain with no ``and`` / ``or``: - ``lambda a, b, c: a < b < c`` over a join returned False correctly under - ``select`` but raised USER_RAISED_EXCEPTION under ``filter``. ``RaiseError`` - now overrides ``throwable``, and the same repro returns Python's answer. - ``test_udf_transpile_boolop_short_circuits`` is the regression test. + join drops. ``RaiseError`` now overrides ``throwable``, and the same repro + returns Python's answer. ``test_udf_transpile_boolop_short_circuits`` is the + regression test. (A pure chain ``lambda a, b, c: a < b < c`` over a join was + the original repro; it no longer reaches this path, since the chain now folds + to an atomic CASE WHEN that cannot be split.) * Predicate pushdown through an AGGREGATE or WINDOW -- OPEN, and the reason the bullet above says "through a JOIN". Only the three join sites test ``!cond.throwable``; the Aggregate arm is explicit that it does not ("We can push down deterministic predicate through Aggregate, including throwable predicate"), and the Window / Union / EventTimeWatermark arms partition on ``_.deterministic`` alone. Measured: - ``groupBy("b", "c").agg(sum("a").alias("sa")).filter(lambda: sa < b < c)`` + ``groupBy("b", "c").agg(sum("a").alias("sa")).filter(lambda: sa < b and b < c)`` splits the fold, pushing the ``b < c`` guard BELOW the ``Aggregate`` while ``sa < b`` stays above, and raises. The Aggregate arm's own reasoning is about ROWS -- grouping keys cannot gain distinct values -- and says nothing about conjunct ORDER, which is the only thing a guard in a fold relies on. -* BooleanSimplification -- OPEN, and it breaks the fold two different ways. It - consults neither ``deterministic`` nor ``throwable``. Common-factor elimination - for disjunction (``(a && b) || (a && c) => a && (b || c)``) hoists a shared - guard to the FRONT of a conjunction, where it evaluates unconditionally: - ``(a < b < c) or (d < b < c)`` on ``(5, 1, None, 7)`` raises where Python - returns False. Its absorption case (``(a && b) || b => b``) is worse, because - it drops a guard instead of moving it: ``(a < b < c) or (b < c)`` on - ``(None, 0, 1)`` returns True where Python raises TypeError, with only one of - the two ``raise_error`` guards surviving. A silently wrong answer, not a - spurious error. +* BooleanSimplification -- OPEN, and it breaks the ``&`` / ``|`` fold two + different ways. It consults neither ``deterministic`` nor ``throwable``. + Common-factor elimination for disjunction (``(a && b) || (a && c) => a && (b + || c)``) hoists a shared guard to the FRONT of a conjunction, where it + evaluates unconditionally: ``(a < b and b < c) or (d < b and b < c)`` on + ``(5, 1, None, 7)`` raises where Python returns False. Its absorption case + (``(a && b) || b => b``) is worse, because it drops a guard instead of moving + it: ``(a < b and b < c) or (b < c)`` on ``(None, 0, 1)`` returns True where + Python raises TypeError, with only one of the two ``raise_error`` guards + surviving. A silently wrong answer, not a spurious error. * Subexpression elimination -- OPEN. ``spark.sql.subexpressionElimination.enabled`` defaults true while ``...skipForShortcutExpr`` defaults FALSE, so a conjunct living only in a short-circuited position is hoisted and evaluated @@ -92,13 +100,15 @@ Subexpression elimination and BooleanSimplification are INDEPENDENTLY sufficient, which is worth knowing before reaching for a conf to prove a theory. On the -``(a < b < c) or (d < b < c)`` repro, each of ``skipForShortcutExpr=true``, +``(a < b and b < c) or (d < b and b < c)`` repro, each of ``skipForShortcutExpr=true``, ``subexpressionElimination.enabled=false`` and excluding ``BooleanSimplification`` still raises on its own; only excluding ``BooleanSimplification`` AND disabling one of the subexpression-elimination levers returns Python's answer. -A chain duplicates every interior operand by construction, so it manufactures the -common subexpressions and the common factors these rules need. Treat them as live. +A chain duplicates every interior operand by construction, so where it folds to +``&`` / ``|`` (it does not -- it folds to a CASE WHEN) it would manufacture the +common subexpressions and common factors these rules need. The ``and`` / ``or`` +form does, so treat them as live there. ``len(s)`` lowers to Catalyst ``length`` for a string operand (SPARK-55214). Other ``len`` arguments (numbers, lists, ...) stay interpreted Python. @@ -228,6 +238,32 @@ # nested case to well under a second. _MAX_LOWERED_COMPARISONS = 128 +# Bound on the materialized Catalyst expression size a body's comparisons may +# emit. Each comparison lowering references its operand Columns more than once -- +# `_lower_value_compare` uses each operand in the NULL guard and the comparison, +# `_lower_eq` uses each one three times -- and `ColumnNodeToExpressionConverter` +# recurses into every child with no memoization, so a shared Column DAG expands +# into a tree. A right-nested body (`(x < 0) < ((x < 0) < ...)`) is charged about +# `2d + 1` links but materializes on the order of `2^d` nodes (`3^d` with `==`), +# so the link budget above fires far too late for these shapes: depth 20 is 41 +# links (well under 128) but roughly 4.6e7 materialized nodes, expanded on the +# UDF's first call, outside `_transpile_func`'s fallback `except`. +# +# The weight below counts what the JVM materializes, not how many links Python +# lowers: a comparison operand's expanded size is charged once per reference, so +# nesting compounds the weight per level. A flat 64-link chain (all leaf +# operands) stays near 64 * a small constant, while a depth-20 nest crosses the +# bound and is refused at UDF construction. The multiplier per operand +# reference is 2 for the ordering operators and 3 for `==` / `!=` (see +# `_lower_value_compare` / `_lower_eq`). +_MAX_LOWERED_COMPARISON_WEIGHT = 1 << 16 # 65536: flat chains pass, depth-20 nests refuse + +# Per-link overhead in the materialized-weight budget: the ``when`` / raise_error +# wrapper, the NULL-guard ``Or``/``IsNull`` nodes and the comparison itself. A +# rough constant is enough -- the bound is dominated by the per-reference +# multiplier compounding across nesting depth, not by this leaf overhead. +_LINK_OVERHEAD = 4 + # The four ordering comparisons: same lowering, differing only in the operator # and the symbol used to name it in an error message. _ORDERING_OPS: dict[type, Tuple[Callable[[Any, Any], Column], str]] = { @@ -295,10 +331,13 @@ def _is_definitely_boolean(node: ast.AST) -> bool: """Return True when ``node`` is statically guaranteed to produce a Python ``bool`` (or ``None``, which round-trips through ``coalesce``). - Used to gate ``if``/ternary lowering and ``and`` / ``or``: we only allow the - expression into Catalyst's ``when(coalesce(test, false), ...)`` form, or onto - ``&`` / ``|``, when it provably produces a boolean. Everything else (bare - Name, arithmetic, function calls, subscript, ...) must force a fallback to + Gates ``if``/ternary tests and ``not``, and labels boolean-producing + expressions as the "bool" category: we only let the expression into + Catalyst's ``when(coalesce(test, false), ...)`` form when it provably + produces a boolean. ``and`` / ``or`` need more -- Spark's ``NULL AND false`` + is ``false`` where Python's ``None and False`` is ``None`` -- so they gate on + the stricter ``_is_never_null_boolean``. Everything else (bare Name, + arithmetic, function calls, subscript, ...) forces a fallback to interpreted Python instead of silently diverging. """ match node: @@ -399,6 +438,16 @@ def __init__(self) -> None: # refuses part-way through the variant list and silently emits fewer # options. self._lowered_comparisons: int = 0 + # Materialized-weight budget (see ``_MAX_LOWERED_COMPARISON_WEIGHT``). + # Charged only for top-level comparison links (depth 1); nested comparison + # operands contribute through their parent's weight, so charging them + # again would double-count. ``_compare_depth`` tracks the recursion. + self._lowered_comparison_weight: int = 0 + self._compare_depth: int = 0 + # ``_comparison_weight`` recurses into nested comparisons and is called + # once per operand reference, so a shared nested Compare would be + # recomputed exponentially without this id(node) cache. + self._comparison_weight_cache: dict[int, int] = {} # TODO (SPARK-55218): handle implicit-None return bodies like # ``def f(x): x + x`` -- no return statement means return None; @@ -627,10 +676,21 @@ def _lower_compare_pair( (non-decimal ``NumericType`` for "numeric", UTF8_BINARY ``StringType`` for "string") leave numeric widening as the only coercion Catalyst inserts, and under ANSI a failing cast raises. That is what lets the - chain fold below use a plain ``&`` with no ``coalesce``. + chain fold wrap each link in ``NOT link`` with no ``coalesce``. Every comparison link routes through here, which is what makes this the - place to charge ``_MAX_LOWERED_COMPARISONS``. + place to charge ``_MAX_LOWERED_COMPARISONS`` (link count) and + ``_MAX_LOWERED_COMPARISON_WEIGHT`` (materialized size). The link count + catches bodies whose comparison count grows with nesting (a chain + re-lowers each interior operand once per adjacent pair); the weight + catches bodies whose LINK count stays linear but whose materialized + size grows exponentially -- a right-nested ``(x < 0) < ((x < 0) < ...)`` + is ~``2d + 1`` links but ~``2^d`` nodes, because each operand Column is + referenced more than once and the JVM converter does not memoize. The + weight is charged only at top-level links (``_compare_depth == 1``): + a nested comparison's weight is counted through its parent's operand + weight (see ``_comparison_weight``), so charging it again here would + double-count. """ self._lowered_comparisons += 1 if self._lowered_comparisons > _MAX_LOWERED_COMPARISONS: @@ -639,6 +699,27 @@ def _lower_compare_pair( "comparisons, so the UDF falls back to interpreted Python; " "nested comparisons double the count per level" ) + self._compare_depth += 1 + if self._compare_depth == 1: + self._lowered_comparison_weight += self._link_weight(op, left_node, right_node) + if self._lowered_comparison_weight > _MAX_LOWERED_COMPARISON_WEIGHT: + raise UnsupportedOperationException( + f"lowering this body materializes more than " + f"{_MAX_LOWERED_COMPARISON_WEIGHT} comparison nodes, so the " + "UDF falls back to interpreted Python; nested comparisons " + "double the materialized size per level" + ) + result = self._lower_compare_pair_dispatch(params, op, left_node, right_node) + self._compare_depth -= 1 + return result + + def _lower_compare_pair_dispatch( + self, + params: List[str], + op: ast.cmpop, + left_node: ast.AST, + right_node: ast.AST, + ) -> Column: match op: case ast.Is() | ast.IsNot(): # Only lower `x is None` / `None is x` (and their @@ -674,6 +755,56 @@ def _lower_compare_pair( f"comparison operator {type(op).__name__} is not supported by the transpiler" ) + def _link_weight(self, op: ast.cmpop, left_node: ast.AST, right_node: ast.AST) -> int: + """Materialized Catalyst weight of ONE comparison link. + + Counts what the JVM materializes, not how many links Python lowers: each + operand's expanded size is charged once per reference. The ordering + operators reference each operand twice (NULL guard + comparison) and + ``==`` / ``!=`` three times (see ``_lower_value_compare`` / ``_lower_eq``); + ``is`` / ``is not`` reference the subject once. A leaf operand (param, + constant, arithmetic) has weight 1; a nested comparison's weight is + ``_comparison_weight``, so nesting compounds the weight per level. + """ + if isinstance(op, (ast.Eq, ast.NotEq)): + mult = 3 + elif isinstance(op, (ast.Is, ast.IsNot)): + mult = 1 + else: + mult = 2 + return ( + mult + * ( + self._comparison_operand_weight(left_node) + + self._comparison_operand_weight(right_node) + ) + + _LINK_OVERHEAD + ) + + def _comparison_operand_weight(self, node: ast.AST) -> int: + """Materialized weight of a comparison operand: 1 for a leaf, recurses + for a nested ``ast.Compare`` (whose own links compound the weight).""" + if isinstance(node, ast.Compare): + return self._comparison_weight(node) + return 1 + + def _comparison_weight(self, node: ast.Compare) -> int: + """Materialized weight of an ``ast.Compare`` -- the sum of its link + weights (the chain fold adds a ``NOT``/literal branch per non-final link, + folded into ``_LINK_OVERHEAD``). Memoized by node identity: a shared + nested Compare is an operand of multiple links, so without caching the + recursion would be exponential.""" + key = id(node) + cached = self._comparison_weight_cache.get(key) + if cached is not None: + return cached + nodes = [node.left, *node.comparators] + total = 0 + for op, left_node, right_node in zip(node.ops, nodes, nodes[1:]): + total += self._link_weight(op, left_node, right_node) + self._comparison_weight_cache[key] = total + return total + def _category(self, params: List[str], node: ast.AST) -> str: """Infer ``"numeric"`` or ``"string"`` for ``node`` under the current ``self._param_categories`` assumption (set per input-type variant). @@ -891,9 +1022,11 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: ) case ast.Compare(left, ops, comps): # `in` / `not in` are container membership, not a chainable value - # comparison, and need their own NULL analysis: Python's `in` is - # equality-based and never raises on a `None` probe. Not lowered - # yet, so refuse them and fall back. + # comparison, and need their own NULL analysis: Python's `in` + # semantics depend on the container -- equality for tuple/list, + # substring containment for str -- and a `None` probe can raise + # (`None in "abc"` -> TypeError). Not lowered yet, so refuse and + # fall back. if any(isinstance(o, (ast.In, ast.NotIn)) for o in ops): raise UnsupportedOperationException( "`in` / `not in` is not supported by the transpiler; the " @@ -918,21 +1051,38 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: # copy of the argument, `50 < x < 60` over `rand(7) * 100` was ~30% # true where Python is ~10%. # - # Folding into nested CASE WHENs instead -- atomic to predicate - # splitting -- was rejected: it costs pushdown for every - # `and` / `or` UDF, including those over non-nullable columns that - # need no guard at all. + # The chain folds into a sequential CASE WHEN + # (`CASE WHEN NOT link1 THEN false ... WHEN NOT link(n-1) THEN + # false ELSE linkN END`), NOT an `&`. Each later link sits in a + # branch that runs only after every earlier link is true, so a + # NULL guard in a later link cannot fire on rows Python + # short-circuits past -- the same short-circuit `&` gives at + # evaluation time. The difference is what the optimizer can do + # afterwards: `splitConjunctivePredicates` matches only `And`, + # `BooleanSimplification`'s patterns match `And`/`Or`, and + # subexpression elimination recurses only into a CASE WHEN's first + # condition, so no rule can move a later guard ahead of the links + # that gate it. A chain predicate therefore moves only as a whole + # when pushdown applies, on the same key values Python sees. That + # costs the per-link pushdown `&` would give -- but chains had no + # lowering before this PR, so there is no pushdown to lose, and the + # `and` / `or` fold keeps `&` / `|` for the UDFs that did. # - # See the module docstring for the short-circuit and NULL-guard - # semantics this fold relies on. + # Every link is total (never NULL; see `_lower_compare_pair`), so + # `NOT link` is never NULL and the CASE WHEN mirrors Python's + # short-circuit exactly. See the module docstring for the + # short-circuit and NULL-guard semantics this fold relies on. nodes = [left, *comps] - return functools.reduce( - operator.and_, - [ - self._lower_compare_pair(params, op, left_node, right_node) - for op, left_node, right_node in zip(ops, nodes, nodes[1:]) - ], - ) + links = [ + self._lower_compare_pair(params, op, left_node, right_node) + for op, left_node, right_node in zip(ops, nodes, nodes[1:]) + ] + if len(links) == 1: + return links[0] + chain = when(~links[0], lit(False)) + for link in links[1:-1]: + chain = chain.when(~link, lit(False)) + return chain.otherwise(links[-1]) case ast.BinOp(left=left, op=op, right=right): # Operator selection is driven by the operand *categories* under # the current input-type variant (see ``_category``): Python's @@ -1069,6 +1219,9 @@ def _transpile_from_ast( # the budget bounds one body's lowering, and each variant lowers the body # again from scratch. self._lowered_comparisons = 0 + self._lowered_comparison_weight = 0 + self._compare_depth = 0 + self._comparison_weight_cache = {} function_body = function_ast.body if len(function_body) != 1: raise UnsupportedOperationException( From 857d9a16a4170d1dc92045094e4b01e71a124b78 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 24 Oct 2026 05:21:23 +0000 Subject: [PATCH 16/18] [SPARK-58456][PYTHON][SQL][MINOR] Re-assert transpiled on the rebuilt join-test UDF The join half of test_udf_transpile_boolop_short_circuits rebuilds and_udf under the deterministic conf; assert it transpiled there too, so a construction-time fallback fails with a clear message rather than surfacing as an EvalPython count mismatch. Co-authored-by: Holden Karau Co-authored-by: Holden Karau Co-authored-by: Cursor --- python/pyspark/sql/tests/test_udf_transpile_unit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index 173d47ad33d5d..2e05755438b0c 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -656,6 +656,7 @@ def either_positive(a, c): deterministic["spark.sql.autoBroadcastJoinThreshold"] = -1 with self.sql_conf(deterministic): and_udf = UserDefinedFunction(both_positive, BooleanType()) + self.assertTrue(and_udf.transpiled, "rebuilt and_udf should transpile") left = self.spark.createDataFrame([(1, -5)], "k long, a long") right = self.spark.createDataFrame([(1, None)], "k long, c long") joined = left.join(right, "k") From 0a3d392f6b16075385161bbb0810b60b8c29e0df Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 31 Oct 2026 00:42:57 +0000 Subject: [PATCH 17/18] [SPARK-58456][PYTHON][SQL] Bound the whole body's materialized size up front Replace the per-link comparison weight with one memoized pass over the body AST that charges each child once per reference its parent's lowering makes, so ternary, not, len and arithmetic operands are covered too. Drive the caps test through _transpile_from_ast, cover the 128-link budget at N and N+1 and the per-variant reset, run the and/or short-circuit test under both factory modes, and correct comments left over from the earlier & chain fold. Co-authored-by: Holden Karau Co-authored-by: Holden Karau Co-Authored-By: Claude Code --- .../sql/tests/test_udf_transpile_unit.py | 243 +++++++++--------- python/pyspark/sql/transpile.py | 233 ++++++----------- 2 files changed, 202 insertions(+), 274 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index 2e05755438b0c..002c39b781a76 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -509,16 +509,16 @@ def gte_zero(x): def test_udf_transpile_chain_short_circuits(self): # THE test for the chained lowering. Python's chain short-circuits: when # ``a < b`` is false, ``b < c`` is never evaluated, so a NULL ``c`` cannot - # raise. `&` short-circuits at evaluation time too, so the fold matches. - # Nullable *columns* are required: a literal `None` has no column - # category, so `1 < 0 < None` just falls back. + # raise. The chain's CASE WHEN only reaches a later link once every earlier + # one is true, so the fold matches. Nullable *columns* are required: a + # literal `None` has no column category, so `1 < 0 < None` just falls back. # - # Both of `And`'s two laziness implementations are exercised: `doGenCode` + # Both of `CaseWhen`'s laziness implementations are exercised: `doGenCode` # under the default FALLBACK factory mode, and `eval` under NO_CODEGEN. # `spark.sql.codegen.wholeStage=false` is NOT enough for the latter -- it # disables whole-stage fusion, but expressions are still compiled through # the codegen factory, so `doGenCode` runs either way. Survival through the - # OPTIMIZER is separate; see test_udf_transpile_boolop_short_circuits. + # OPTIMIZER is separate; see test_udf_transpile_chain_survives_optimizer. def chained(a, b, c): return a < b < c @@ -545,18 +545,19 @@ def chained(a, b, c): def test_udf_transpile_chain_survives_optimizer(self): # A chain folds to a sequential CASE WHEN, which is atomic to - # `splitConjunctivePredicates` and `BooleanSimplification`. Before this - # fix the chain folded to `&`, which those rules could split or reorder, - # moving a later link's NULL guard ahead of the links that gate it and - # diverging from Python. Each repro below is an optimizer shape that broke - # the `&` fold; all must return Python's answer (or raise like Python) - # with the chain fully lowered (no EvalPython in the plan). - - # PushDownPredicates / Aggregate: a later link over only grouping keys - # used to be pushed below the Aggregate, where it ran on groups Python - # short-circuits past. `sa < b < c` over groupBy(b,c).agg(sum(a)): the - # group (b=1, c=NULL, sa=5) has `sa < b` false, so Python returns False - # without touching `b < c` (which would raise on the NULL c). + # `splitConjunctivePredicates` and `BooleanSimplification`. An `&` fold + # would let those rules split or reorder it, moving a later link's NULL + # guard ahead of the links that gate it -- as they still do for the + # equivalent user-written `and` (see the module docstring). Each repro + # below is such a shape; all must return Python's answer (or raise like + # Python) with the chain fully lowered (no EvalPython in the plan). + + # PushDownPredicates / Aggregate: an `&` fold's later link over only + # grouping keys would be pushed below the Aggregate, where it runs on + # groups Python short-circuits past. `sa < b < c` over + # groupBy(b,c).agg(sum(a)): the group (b=1, c=NULL, sa=6) has `sa < b` + # false, so Python returns False without touching `b < c` (which would + # raise on the NULL c). from pyspark.sql.functions import sum as sum_ def chain_over_agg(sa, b, c): @@ -573,15 +574,15 @@ def chain_over_agg(sa, b, c): .agg(sum_("a").alias("sa")) ) filtered = agg.filter(udf("sa", "b", "c")) - # Python: row (sa=5, b=1, c=None) -> 5 < 1 is False -> False (no raise). + # Python: row (sa=6, b=1, c=None) -> 6 < 1 is False -> False (no raise). # row (sa=2, b=2, c=3) -> 2 < 2 is False -> False. # No row passes, and crucially no raise leaks out of the NULL c. self.assertEqual([], filtered.select("b", "c").collect()) self.assertEqual(0, self._eval_python_count(filtered), "chain must stay lowered") # BooleanSimplification common-factor: `(a < b < c) or (d < b < c)` on - # (5, 1, None, 7). The `&` fold let common-factor elimination hoist the - # shared `b < c` guard to the front, where the NULL c raised where Python + # (5, 1, None, 7). Over an `&` fold, common-factor elimination hoists the + # shared `b < c` guard to the front, where the NULL c raises where Python # returns False. The CASE WHEN fold is atomic, so the shared link cannot # be hoisted out of either chain. def or_shared_chain(a, b, c, d): @@ -598,7 +599,7 @@ def or_shared_chain(a, b, c, d): self.assertEqual(0, self._eval_python_count(projected), "or of chains stays lowered") # BooleanSimplification absorption: `(a < b < c) or (b < c)` on - # (None, 0, 1). The `&` fold's absorption (`(a && b) || b => b`) dropped + # (None, 0, 1). Over an `&` fold, absorption (`(a && b) || b => b`) drops # the first link's guard, returning True where Python raises TypeError on # `None < 0`. The CASE WHEN fold keeps both links, so the first chain # reaches its NULL guard and raises like Python. @@ -635,18 +636,26 @@ def both_positive(a, c): def either_positive(a, c): return a > 0 or c > 0 - with self.sql_conf(_TRANSPILE_ON): - and_udf = UserDefinedFunction(both_positive, BooleanType()) - or_udf = UserDefinedFunction(either_positive, BooleanType()) - self.assertTrue(and_udf.transpiled, "and over compares should transpile") - self.assertTrue(or_udf.transpiled, "or over compares should transpile") - - # Plain select: the second operand is NULL but the first already - # decided the result, so Python's answer must come back unraised. - df = self.spark.createDataFrame([(-5, None)], "a long, c long") - self.assertEqual([False], [r[0] for r in df.select(and_udf("a", "c")).collect()]) - df = self.spark.createDataFrame([(5, None)], "a long, c long") - self.assertEqual([True], [r[0] for r in df.select(or_udf("a", "c")).collect()]) + # Plain select: the second operand is NULL but the first already decided + # the result, so Python's answer must come back unraised. Under both + # factory modes, so `And` / `Or`'s `doGenCode` and `eval` laziness are + # both covered (see test_udf_transpile_chain_short_circuits). + for factory_mode in ("FALLBACK", "NO_CODEGEN"): + conf = dict(_TRANSPILE_ON) + conf["spark.sql.codegen.factoryMode"] = factory_mode + with self.subTest(factory_mode=factory_mode), self.sql_conf(conf): + and_udf = UserDefinedFunction(both_positive, BooleanType()) + or_udf = UserDefinedFunction(either_positive, BooleanType()) + self.assertTrue(and_udf.transpiled, "and over compares should transpile") + self.assertTrue(or_udf.transpiled, "or over compares should transpile") + df = self.spark.createDataFrame([(-5, None)], "a long, c long") + projected = df.select(and_udf("a", "c")) + self.assertEqual([False], [r[0] for r in projected.collect()]) + self.assertEqual(0, self._eval_python_count(projected), "and select lowered") + df = self.spark.createDataFrame([(5, None)], "a long, c long") + projected = df.select(or_udf("a", "c")) + self.assertEqual([True], [r[0] for r in projected.collect()]) + self.assertEqual(0, self._eval_python_count(projected), "or select lowered") # Same thing across a join, where the NULL operand lives entirely on the # right side. Deterministic plan: no AQE, no broadcast, so the pushed-down @@ -849,40 +858,41 @@ def chain_of_compare(a, b, c): self.assertEqual([True], [r[0] for r in df.select(pudf("c", "x", "y")).collect()]) def test_udf_transpile_lowering_caps(self): - # Both bounds on comparison lowering. Exercised directly rather than - # through a UDF so a refusal cannot be confused with a fallback for some - # other reason. No `sql_conf` here on purpose: nothing on this path reads - # SQLConf (only `_get_transpilers` does), so wrapping it would imply the - # caps are conf-gated when they are module constants. + # Every bound on comparison lowering. Exercised through + # `_transpile_from_ast` rather than a UDF so a refusal cannot be confused + # with a fallback for some other reason, while still going through its + # per-variant resets. No `sql_conf` here on purpose: nothing on this path + # reads SQLConf (only `_get_transpilers` does), so wrapping it would imply + # the caps are conf-gated when they are module constants. import ast as _ast from pyspark.errors import UnsupportedOperationException from pyspark.sql.transpile import ( _MAX_COMPARISON_OPS, - _MAX_LOWERED_COMPARISON_WEIGHT, _MAX_LOWERED_COMPARISONS, + _MAX_MATERIALIZED_WEIGHT, CatalystTranspiler, + _materialized_weight, ) transpiler = CatalystTranspiler() - transpiler._param_categories = {0: "numeric"} def compare_node(src): return _ast.parse(src, mode="eval").body - def lower(src): - # `_transpile_from_ast` is what normally zeroes the per-lowering state; - # reset it here so each case is charged on its own. The category cache - # has to go too: it keys on `id(node)`, and CPython reuses the id of a - # collected node for one of a different type, so a cache kept across - # separately parsed sources hands back a stale category. The weight - # cache keys on `id(node)` for the same reason. - transpiler._lowered_comparisons = 0 - transpiler._lowered_comparison_weight = 0 - transpiler._compare_depth = 0 - transpiler._category_cache = {} - transpiler._comparison_weight_cache = {} - return transpiler._convert_chunk(["x"], compare_node(src)) + def lower(src, category="numeric"): + fdef = _ast.parse(f"def f(x):\n return {src}").body[0] + return transpiler._transpile_from_ast( + src, fdef, fdef, ["x"], BooleanType(), {0: category} + ) + + def refused(src): + with self.assertRaises(UnsupportedOperationException) as ctx: + lower(src) + return str(ctx.exception) + + link_msg = f"more than {_MAX_LOWERED_COMPARISONS} comparisons" + weight_msg = f"more than {_MAX_MATERIALIZED_WEIGHT}" def chain_src(n): # n operators over n + 1 operands: x < 0 < 1 < ... < n - 1 @@ -890,94 +900,97 @@ def chain_src(n): # ---- Per-chain operator cap ---- lower(chain_src(_MAX_COMPARISON_OPS)) # at the cap, succeeds - with self.assertRaises(UnsupportedOperationException) as ctx: - lower(chain_src(_MAX_COMPARISON_OPS + 1)) - self.assertIn("operators", str(ctx.exception)) + self.assertIn("operators", refused(chain_src(_MAX_COMPARISON_OPS + 1))) + + # ---- Link budget, at N and N + 1 ---- + def or_src(n): + return " or ".join(["(x == %d)" % i for i in range(n)]) + + lower(or_src(_MAX_LOWERED_COMPARISONS)) + self.assertIn(link_msg, refused(or_src(_MAX_LOWERED_COMPARISONS + 1))) + + # The budget resets per input-type variant: 2 x 100 links both lower. + hundred = " or ".join(["(x == x)"] * 100) + for category in ("numeric", "string"): + self.assertIsNotNone(lower(hundred, category), category) - # ---- Whole-body budget ---- # A comparison may be an operand of another comparison (both category # "bool"), and a chain re-lowers each interior operand once per adjacent - # pair, so nesting doubles the link count per level while every node stays - # far under the per-chain cap. Before the budget existed this reached hours - # of driver-side py4j work from a few hundred characters of source. + # pair, so nesting doubles the link count per level. Before the budget + # existed this reached hours of driver-side py4j work from a few hundred + # characters of source. def nested_src(depth): src = "(x < 0)" for _ in range(depth): src = "((x < 0) < %s < (x < 0))" % src return src - # Shallow nesting stays inside the budget and still lowers. lower(nested_src(2)) - - # Deep nesting is refused by the budget, not by the per-chain cap: every - # `ast.Compare` in here has exactly two operators. + # Depth 5 is 156 links but under the weight bound: only the link budget + # stops it. + self.assertLess(_materialized_weight(compare_node(nested_src(5))), _MAX_MATERIALIZED_WEIGHT) + self.assertIn(link_msg, refused(nested_src(5))) + # ... and the refusal part-way through left no stale count behind. + lower(chain_src(_MAX_COMPARISON_OPS)) + + # ---- Materialized-weight budget ---- + # Every `ast.Compare` here is under the per-chain cap, and the weight is + # checked before the first link is lowered, so the weight is what fires. deep = nested_src(10) - node = compare_node(deep) self.assertTrue( all( len(n.ops) <= _MAX_COMPARISON_OPS - for n in _ast.walk(node) + for n in _ast.walk(compare_node(deep)) if isinstance(n, _ast.Compare) - ), - "every node must be under the per-chain cap, so the budget is what fires", + ) ) - with self.assertRaises(UnsupportedOperationException) as ctx: - lower(deep) - # This shape's LINK count also grows with nesting (a chain re-lowers each - # interior operand once per adjacent pair), so it can trip the link budget - # too; either budget refusing is correct. Assert the materialized-weight - # bound fires, since that is the bound this shape is now judged by. - self.assertIn(str(_MAX_LOWERED_COMPARISON_WEIGHT), str(ctx.exception)) - - # ---- Right-nested body: linear links, exponential materialized size ---- - # `(x < 0) < ((x < 0) < (...))` nests in ONE operand of a SINGLE - # comparison per level, so the link count grows linearly (~d + 1) while - # the materialized size doubles per level (each operand Column is - # referenced twice and the JVM converter does not memoize). The link - # budget above does NOT fire here -- depth 20 is 21 links, well under - # `_MAX_LOWERED_COMPARISONS` -- so the materialized-weight budget is what - # has to refuse it. - def right_nested_src(depth): - src = "(x < 0)" + self.assertIn(weight_msg, refused(deep)) + + # Right-nested: ONE operand of a single comparison per level, so 2d + 1 + # links (41 at depth 20, under the link budget) but ~2^d materialized + # nodes, since each operand Column is referenced twice (3x for `==`) and + # the JVM converter does not memoize. + def right_nested_src(depth, leaf="(x < 0)", op="<"): + src = leaf for _ in range(depth): - src = "(x < 0) < (%s)" % src + src = "%s %s (%s)" % (leaf, op, src) return src - # Shallow right-nesting stays inside the weight budget and still lowers. lower(right_nested_src(2)) - deep_right = right_nested_src(20) - right_node = compare_node(deep_right) - # The link count is linear and well under the link budget, proving the - # weight budget -- not the link budget -- is what refuses this shape. self.assertLess( - sum(len(n.ops) for n in _ast.walk(right_node) if isinstance(n, _ast.Compare)), + sum( + len(n.ops) + for n in _ast.walk(compare_node(deep_right)) + if isinstance(n, _ast.Compare) + ), _MAX_LOWERED_COMPARISONS, - "right-nest link count must stay under the link budget, so the weight " - "budget is what fires", ) - with self.assertRaises(UnsupportedOperationException) as ctx: - lower(deep_right) - self.assertIn(str(_MAX_LOWERED_COMPARISON_WEIGHT), str(ctx.exception)) - - # `==` references each operand three times (vs twice for `<`), so an - # equality right-nest compounds faster; the same depth already refuses. - def eq_right_nested_src(depth): - src = "(x == 1)" - for _ in range(depth): - src = "(x == 1) == (%s)" % src - return src - - with self.assertRaises(UnsupportedOperationException) as ctx: - lower(eq_right_nested_src(20)) - self.assertIn(str(_MAX_LOWERED_COMPARISON_WEIGHT), str(ctx.exception)) + self.assertIn(weight_msg, refused(deep_right)) + self.assertIn(weight_msg, refused(right_nested_src(20, "(x == 1)", "=="))) + + # The same nest routed through a ternary, in each of its three positions: + # the ternary's Column is referenced as often as any other operand. + ternary_shapes = { + "test": "(x < 0) {op} (True if ({inner}) else False)", + "body": "(x < 0) {op} (({inner}) if x > 0 else (x < 0))", + "else": "(x < 0) {op} ((x < 0) if x > 0 else ({inner}))", + } + for position, shape in ternary_shapes.items(): + for op in ("<", "=="): + with self.subTest(position=position, op=op): + src = "(x < 0)" + for _ in range(20): + src = shape.format(op=op, inner=src) + self.assertIn(weight_msg, refused(src)) + # One level over a ternary still lowers. + lower("(x < 0) < ((x < 1) if x > 0 else (x < 2))") + # `is None` reads only its subject, so an identity nest stays linear. + lower("(" * 20 + "x < 0" + ") is None" * 20) # `in` / `not in` are refused outright; membership needs its own NULL - # analysis and is deferred to its own ticket (SPARK-56925 is the - # unrelated bare-`if` truthiness work). - with self.assertRaises(UnsupportedOperationException) as ctx: - lower("x in (1, 2)") - self.assertIn("`in` / `not in`", str(ctx.exception)) + # analysis and is deferred to its own ticket. + self.assertIn("`in` / `not in`", refused("x in (1, 2)")) def test_udf_transpile_multi_row(self): # Every other transpile test uses a 1-row DataFrame; this one runs diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index 51c6a5831d61c..827f4f3121077 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -105,10 +105,9 @@ still raises on its own; only excluding ``BooleanSimplification`` AND disabling one of the subexpression-elimination levers returns Python's answer. -A chain duplicates every interior operand by construction, so where it folds to -``&`` / ``|`` (it does not -- it folds to a CASE WHEN) it would manufacture the -common subexpressions and common factors these rules need. The ``and`` / ``or`` -form does, so treat them as live there. +These rules need a common subexpression or common factor inside an ``&`` / +``|``. A user-written ``and`` / ``or`` that repeats an operand +(``a < b and b < c``) hands them one, so treat them as live there. ``len(s)`` lowers to Catalyst ``length`` for a string operand (SPARK-55214). Other ``len`` arguments (numbers, lists, ...) stay interpreted Python. @@ -238,31 +237,20 @@ # nested case to well under a second. _MAX_LOWERED_COMPARISONS = 128 -# Bound on the materialized Catalyst expression size a body's comparisons may -# emit. Each comparison lowering references its operand Columns more than once -- -# `_lower_value_compare` uses each operand in the NULL guard and the comparison, -# `_lower_eq` uses each one three times -- and `ColumnNodeToExpressionConverter` -# recurses into every child with no memoization, so a shared Column DAG expands -# into a tree. A right-nested body (`(x < 0) < ((x < 0) < ...)`) is charged about -# `2d + 1` links but materializes on the order of `2^d` nodes (`3^d` with `==`), -# so the link budget above fires far too late for these shapes: depth 20 is 41 -# links (well under 128) but roughly 4.6e7 materialized nodes, expanded on the -# UDF's first call, outside `_transpile_func`'s fallback `except`. -# -# The weight below counts what the JVM materializes, not how many links Python -# lowers: a comparison operand's expanded size is charged once per reference, so -# nesting compounds the weight per level. A flat 64-link chain (all leaf -# operands) stays near 64 * a small constant, while a depth-20 nest crosses the -# bound and is refused at UDF construction. The multiplier per operand -# reference is 2 for the ordering operators and 3 for `==` / `!=` (see -# `_lower_value_compare` / `_lower_eq`). -_MAX_LOWERED_COMPARISON_WEIGHT = 1 << 16 # 65536: flat chains pass, depth-20 nests refuse - -# Per-link overhead in the materialized-weight budget: the ``when`` / raise_error -# wrapper, the NULL-guard ``Or``/``IsNull`` nodes and the comparison itself. A -# rough constant is enough -- the bound is dominated by the per-reference -# multiplier compounding across nesting depth, not by this leaf overhead. -_LINK_OVERHEAD = 4 +# Bound on the Catalyst tree a body materializes (see `_materialized_weight`). +# Lowerings reference operand Columns more than once and +# `ColumnNodeToExpressionConverter` does not memoize, so the shared Column DAG +# expands into a tree on the UDF's first call -- outside `_transpile_func`'s +# fallback `except`. `(x < 0) < ((x < 0) < ...)` is 2d + 1 links but ~2^d nodes +# (3^d with `==`), so the link budget alone fires far too late. +_MAX_MATERIALIZED_WEIGHT = 1 << 16 # 65536: flat 64-link chains pass, depth-20 nests refuse + +# Nodes a lowering adds around its operands (`when` / raise_error, NULL guard, ...). +_NODE_OVERHEAD = 4 + +# How often a comparison link reads each operand: `_lower_eq` three times, `is` +# only the subject, `_lower_value_compare` (the default) twice. +_OPERAND_READS: dict[type, int] = {ast.Eq: 3, ast.NotEq: 3, ast.Is: 1, ast.IsNot: 1} # The four ordering comparisons: same lowering, differing only in the operator # and the symbol used to name it in an error message. @@ -307,6 +295,43 @@ def _transpile_from_ast( pass +def _materialized_weight(body: ast.AST) -> int: + """Rough size of the Catalyst tree ``_convert_chunk`` emits for ``body``. + + A child is charged once per reference its parent's lowering makes, since the + JVM expands every reference. Memoized by node, so this stays linear. + """ + memo: dict[int, int] = {} + + def weight(node: Optional[ast.AST]) -> int: + if node is None: + return 1 + cached = memo.get(id(node)) + if cached is not None: + return cached + match node: + case ast.Compare(left=left, ops=ops, comparators=comps): + nodes = [left, *comps] + total = 0 + for op, lhs, rhs in zip(ops, nodes, nodes[1:]): + total += _OPERAND_READS.get(type(op), 2) * (weight(lhs) + weight(rhs)) + total += _NODE_OVERHEAD + case ast.BinOp(left=left, op=ast.Mod(), right=right): + # `sb * pmod(sb * a, abs(b))` with `sb` reading `b` twice. + total = weight(left) + 5 * weight(right) + 2 * _NODE_OVERHEAD + case ast.Call(func=ast.Name(id="len"), args=[arg], keywords=[]): + total = 2 * weight(arg) + _NODE_OVERHEAD + case _: + # A loop, not `sum(...)`: one frame per level, like `_convert_chunk`. + total = 1 + for child in ast.iter_child_nodes(node): + total += weight(child) + memo[id(node)] = total + return total + + return weight(body) + + def _is_definitely_basic_type(node: ast.AST) -> bool: """ Return True when ``node`` is statically guaranteed to produce a Python @@ -438,16 +463,6 @@ def __init__(self) -> None: # refuses part-way through the variant list and silently emits fewer # options. self._lowered_comparisons: int = 0 - # Materialized-weight budget (see ``_MAX_LOWERED_COMPARISON_WEIGHT``). - # Charged only for top-level comparison links (depth 1); nested comparison - # operands contribute through their parent's weight, so charging them - # again would double-count. ``_compare_depth`` tracks the recursion. - self._lowered_comparison_weight: int = 0 - self._compare_depth: int = 0 - # ``_comparison_weight`` recurses into nested comparisons and is called - # once per operand reference, so a shared nested Compare would be - # recomputed exponentially without this id(node) cache. - self._comparison_weight_cache: dict[int, int] = {} # TODO (SPARK-55218): handle implicit-None return bodies like # ``def f(x): x + x`` -- no return statement means return None; @@ -678,19 +693,9 @@ def _lower_compare_pair( inserts, and under ANSI a failing cast raises. That is what lets the chain fold wrap each link in ``NOT link`` with no ``coalesce``. - Every comparison link routes through here, which is what makes this the - place to charge ``_MAX_LOWERED_COMPARISONS`` (link count) and - ``_MAX_LOWERED_COMPARISON_WEIGHT`` (materialized size). The link count - catches bodies whose comparison count grows with nesting (a chain - re-lowers each interior operand once per adjacent pair); the weight - catches bodies whose LINK count stays linear but whose materialized - size grows exponentially -- a right-nested ``(x < 0) < ((x < 0) < ...)`` - is ~``2d + 1`` links but ~``2^d`` nodes, because each operand Column is - referenced more than once and the JVM converter does not memoize. The - weight is charged only at top-level links (``_compare_depth == 1``): - a nested comparison's weight is counted through its parent's operand - weight (see ``_comparison_weight``), so charging it again here would - double-count. + Every comparison link routes through here, so this is where + ``_MAX_LOWERED_COMPARISONS`` is charged. Materialized size is bounded + separately, up front, by ``_MAX_MATERIALIZED_WEIGHT``. """ self._lowered_comparisons += 1 if self._lowered_comparisons > _MAX_LOWERED_COMPARISONS: @@ -699,27 +704,6 @@ def _lower_compare_pair( "comparisons, so the UDF falls back to interpreted Python; " "nested comparisons double the count per level" ) - self._compare_depth += 1 - if self._compare_depth == 1: - self._lowered_comparison_weight += self._link_weight(op, left_node, right_node) - if self._lowered_comparison_weight > _MAX_LOWERED_COMPARISON_WEIGHT: - raise UnsupportedOperationException( - f"lowering this body materializes more than " - f"{_MAX_LOWERED_COMPARISON_WEIGHT} comparison nodes, so the " - "UDF falls back to interpreted Python; nested comparisons " - "double the materialized size per level" - ) - result = self._lower_compare_pair_dispatch(params, op, left_node, right_node) - self._compare_depth -= 1 - return result - - def _lower_compare_pair_dispatch( - self, - params: List[str], - op: ast.cmpop, - left_node: ast.AST, - right_node: ast.AST, - ) -> Column: match op: case ast.Is() | ast.IsNot(): # Only lower `x is None` / `None is x` (and their @@ -755,56 +739,6 @@ def _lower_compare_pair_dispatch( f"comparison operator {type(op).__name__} is not supported by the transpiler" ) - def _link_weight(self, op: ast.cmpop, left_node: ast.AST, right_node: ast.AST) -> int: - """Materialized Catalyst weight of ONE comparison link. - - Counts what the JVM materializes, not how many links Python lowers: each - operand's expanded size is charged once per reference. The ordering - operators reference each operand twice (NULL guard + comparison) and - ``==`` / ``!=`` three times (see ``_lower_value_compare`` / ``_lower_eq``); - ``is`` / ``is not`` reference the subject once. A leaf operand (param, - constant, arithmetic) has weight 1; a nested comparison's weight is - ``_comparison_weight``, so nesting compounds the weight per level. - """ - if isinstance(op, (ast.Eq, ast.NotEq)): - mult = 3 - elif isinstance(op, (ast.Is, ast.IsNot)): - mult = 1 - else: - mult = 2 - return ( - mult - * ( - self._comparison_operand_weight(left_node) - + self._comparison_operand_weight(right_node) - ) - + _LINK_OVERHEAD - ) - - def _comparison_operand_weight(self, node: ast.AST) -> int: - """Materialized weight of a comparison operand: 1 for a leaf, recurses - for a nested ``ast.Compare`` (whose own links compound the weight).""" - if isinstance(node, ast.Compare): - return self._comparison_weight(node) - return 1 - - def _comparison_weight(self, node: ast.Compare) -> int: - """Materialized weight of an ``ast.Compare`` -- the sum of its link - weights (the chain fold adds a ``NOT``/literal branch per non-final link, - folded into ``_LINK_OVERHEAD``). Memoized by node identity: a shared - nested Compare is an operand of multiple links, so without caching the - recursion would be exponential.""" - key = id(node) - cached = self._comparison_weight_cache.get(key) - if cached is not None: - return cached - nodes = [node.left, *node.comparators] - total = 0 - for op, left_node, right_node in zip(node.ops, nodes, nodes[1:]): - total += self._link_weight(op, left_node, right_node) - self._comparison_weight_cache[key] = total - return total - def _category(self, params: List[str], node: ast.AST) -> str: """Infer ``"numeric"`` or ``"string"`` for ``node`` under the current ``self._param_categories`` assumption (set per input-type variant). @@ -1038,40 +972,17 @@ def _convert_chunk(self, params: List[str], body: ast.AST | None) -> Column: f"transpiler lowers at most {_MAX_COMPARISON_OPS}, so the " "UDF falls back to interpreted Python" ) - # Python evaluates `a OP1 b OP2 c` as `(a OP1 b) and (b OP2 c)` - # with `b` evaluated once, so the fold below re-lowers each - # interior operand once per adjacent pair. Safe for the operand - # grammar we accept, all of it deterministic and side-effect free. - # It does mean the body READS a parameter once per adjacent pair, - # which SPARK-58626 is what makes safe: `_udf_param_N` is a - # reference, and `ConvertToCatalyst` gives an argument read more - # than once one column per row -- or declines to lower the call - # where no column fits -- so a non-deterministic argument is drawn - # once and every read agrees. While each read still carried its own - # copy of the argument, `50 < x < 60` over `rand(7) * 100` was ~30% - # true where Python is ~10%. - # - # The chain folds into a sequential CASE WHEN - # (`CASE WHEN NOT link1 THEN false ... WHEN NOT link(n-1) THEN - # false ELSE linkN END`), NOT an `&`. Each later link sits in a - # branch that runs only after every earlier link is true, so a - # NULL guard in a later link cannot fire on rows Python - # short-circuits past -- the same short-circuit `&` gives at - # evaluation time. The difference is what the optimizer can do - # afterwards: `splitConjunctivePredicates` matches only `And`, - # `BooleanSimplification`'s patterns match `And`/`Or`, and - # subexpression elimination recurses only into a CASE WHEN's first - # condition, so no rule can move a later guard ahead of the links - # that gate it. A chain predicate therefore moves only as a whole - # when pushdown applies, on the same key values Python sees. That - # costs the per-link pushdown `&` would give -- but chains had no - # lowering before this PR, so there is no pushdown to lose, and the - # `and` / `or` fold keeps `&` / `|` for the UDFs that did. + # `a OP1 b OP2 c` is `(a OP1 b) and (b OP2 c)` with `b` evaluated + # once; re-lowering `b` per pair is safe because SPARK-58626 gives a + # repeated argument one column per row (before it, `50 < x < 60` + # over `rand(7) * 100` was ~30% true where Python is ~10%). # - # Every link is total (never NULL; see `_lower_compare_pair`), so - # `NOT link` is never NULL and the CASE WHEN mirrors Python's - # short-circuit exactly. See the module docstring for the - # short-circuit and NULL-guard semantics this fold relies on. + # Folds to `CASE WHEN NOT link1 THEN false ... ELSE linkN END`, not + # `&`: a later link runs only once every earlier one is true, and no + # optimizer rule can split it (see the module docstring). The cost: + # no per-link pushdown, and over nullable operands the guards make + # it throwable, so it stays above a join the interpreted UDF went + # below. Every link is total, so `NOT link` is never NULL. nodes = [left, *comps] links = [ self._lower_compare_pair(params, op, left_node, right_node) @@ -1219,15 +1130,19 @@ def _transpile_from_ast( # the budget bounds one body's lowering, and each variant lowers the body # again from scratch. self._lowered_comparisons = 0 - self._lowered_comparison_weight = 0 - self._compare_depth = 0 - self._comparison_weight_cache = {} function_body = function_ast.body if len(function_body) != 1: raise UnsupportedOperationException( "functions with more than one top-level statement are not " "supported by the transpiler" ) + weight = _materialized_weight(function_body[0]) + if weight > _MAX_MATERIALIZED_WEIGHT: + raise UnsupportedOperationException( + f"lowering this body materializes about {weight} expression nodes, " + f"more than {_MAX_MATERIALIZED_WEIGHT}, so the UDF falls back to " + "interpreted Python; nesting multiplies the size per level" + ) # Refuse variants whose body category does not MATCH the declared # return type's category. Two distinct failure modes hide here: # From b50780b5cce7d479bdbbbdbbf4cbb15d38ab5858 Mon Sep 17 00:00:00 2001 From: Holden Karau Date: Sat, 31 Oct 2026 02:23:39 +0000 Subject: [PATCH 18/18] [SPARK-58456][PYTHON][SQL][MINOR] Re-pin the NaN equality chain after SPARK-59091 Merging master brought SPARK-59091's NaN guard in transpiled equality, so `a == b == c` over NaN now returns Python's False. Update the pin and the module docstring; ordering stays the open NaN divergence. Note why the weight pass keeps 3 reads for `==`: the guard's extra read only hits numeric leaves. Co-authored-by: Holden Karau Co-authored-by: Holden Karau Co-Authored-By: Claude Code --- python/pyspark/sql/tests/test_udf_transpile_unit.py | 13 ++++++------- python/pyspark/sql/transpile.py | 10 ++++++---- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py index a88742ab39740..c4346dde426b9 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -2977,11 +2977,10 @@ def test_udf_transpile_known_value_divergences(self): self._vals(eq_strlit, BooleanType(), "a long", [(5,), (3,)], require_lowered=False), [False, False], ) - # Chained comparisons inherit the same NaN divergence one link at a time -- - # no new mechanism, just easier to hit. Spark orders NaN above every value - # (Python's `b < nan` is False) and treats `NaN = NaN` as true (Python's - # `nan == nan` is False). These assertions pin the CURRENT divergent - # behavior so a change is visible; SPARK-58781 decides what to do about it. + # Chained comparisons inherit the NaN ordering divergence one link at a + # time: Spark orders NaN above every value (Python's `b < nan` is False). + # Pinned so a change is visible; SPARK-58781 decides what to do about it. + # Equality agrees with Python since SPARK-59091's NaN guard. nan_chain = lambda a, b, c: a < b < c # noqa: E731 nan_eq_chain = lambda a, b, c: a == b == c # noqa: E731 nan = float("nan") @@ -2994,8 +2993,8 @@ def test_udf_transpile_known_value_divergences(self): self._vals( nan_eq_chain, BooleanType(), "a double, b double, c double", [(nan, nan, nan)] ), - [True], - "Spark treats NaN = NaN as true; Python would return False", + [False], + "NaN == NaN is False, as in Python", ) def test_udf_transpile_overflow_and_modulo_zero_raise(self): diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py index bf567ec1d737c..4b29f9a45fb00 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -95,8 +95,9 @@ living only in a short-circuited position is hoisted and evaluated unconditionally. ``Expression.throwable`` does NOT help -- this rule never consults it, so SPARK-58627 does not cover this one. -* NaN -- OPEN. Spark orders NaN above every value and treats ``NaN = NaN`` as - true where Python does neither (SPARK-58781). +* NaN ordering -- OPEN. Spark orders NaN above every value where Python's NaN + comparisons are all False (SPARK-58781). Equality matches Python since + SPARK-59091's NaN guard. Subexpression elimination and BooleanSimplification are INDEPENDENTLY sufficient, which is worth knowing before reaching for a conf to prove a theory. On the @@ -270,8 +271,9 @@ def _wider_numeric(lc: str, rc: str) -> str: # Nodes a lowering adds around its operands (`when` / raise_error, NULL guard, ...). _NODE_OVERHEAD = 4 -# How often a comparison link reads each operand: `_lower_eq` three times, `is` -# only the subject, `_lower_value_compare` (the default) twice. +# How often a comparison link reads each operand: `_lower_eq` three times (its +# NaN guard adds a fourth read, but only of numeric leaves, which never nest), +# `is` only the subject, `_lower_value_compare` (the default) twice. _OPERAND_READS: dict[type, int] = {ast.Eq: 3, ast.NotEq: 3, ast.Is: 1, ast.IsNot: 1} # The four ordering comparisons: same lowering, differing only in the operator