diff --git a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py index c0db082821ca..1cce17d5e79d 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py +++ b/python/pyspark/sql/tests/test_udf_transpile_hypothesis.py @@ -206,6 +206,25 @@ def _env_opts_in(value: Optional[str]) -> bool: (1, -_LONG_BOUND), ) + # 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, 0, 1), + (0, None, 1), + (0, 1, None), + (1, 0, None), # first link false -> None never reached + # Boundary orderings. + (0, 1, 2), + (2, 1, 0), + (_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 +235,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 -------------------------------------- # @@ -231,6 +256,16 @@ def wrapper(method): # (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 (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): @@ -343,6 +378,35 @@ def neq_pair(x, y): return 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; ``_run``'s "both raised" equivalence covers that. + return 0 < x < 10 + + +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 + # transpiled form must not raise there either. + return x < y < z + + +def chained_eq(x: int, y: int, z: int) -> bool: + # 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 links + # disagree on direction and short-circuit on different rows again. + return x <= y > z + + # A lambda captured at module scope so ``inspect.getsource`` can read # its definition. Exercises the ``ast.Lambda`` branch in # ``_get_function_from_ast``. @@ -440,9 +504,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, @@ -466,6 +530,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 +738,52 @@ 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, 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) + @_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_eq_matches_python(self, x, y, z): + 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})" + ) + 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 0df1f4b969fb..a94dc1530a91 100644 --- a/python/pyspark/sql/tests/test_udf_transpile_unit.py +++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py @@ -54,6 +54,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.", @@ -328,8 +342,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 `&`/`|`, + # 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 @@ -353,11 +368,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), @@ -549,39 +564,491 @@ 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. + 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 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 `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_chain_survives_optimizer. + def chained(a, b, c): + return a < b < c + + schema = "a long, b long, c long" + 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): + pudf = UserDefinedFunction(chained, BooleanType()) + self.assertTrue(pudf.transpiled, "3-column chain should transpile") + df = self.spark.createDataFrame([(1, 0, None)], 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([(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_chain_survives_optimizer(self): + # A chain folds to a sequential CASE WHEN, which is atomic to + # `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): + 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=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). 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): + 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). 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. + 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 + # result. + # + # 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: 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 + + def either_positive(a, c): + return a > 0 or c > 0 + + # 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 + # 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()) + 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") + 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. + # 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) -- every NULL pattern 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), + ] + 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. + 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 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=label): + self.assertEqual( + self._vals(func, B, "a long, b long, c long", rows), expected, label + ) + 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 well under the cap and must transpile", + ) + + def test_udf_transpile_comparison_fallbacks(self): + # 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 - from pyspark.sql.types import StructField, StructType + # `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 + 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 + + # (label, func, schema, rows, expected). `expected` is what interpreted + # Python produces, so the query must still run and return it. + cases = [ + ("`in` on a str", str_in, "a string", [("b",), ("z",)], [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]), + ] + 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) 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()], + expected, + f"{label}: interpreted result diverged", + ) + # `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"), + (in_long_rt, LongType(), "an `in` with a LongType return type"), + ]: + with self.subTest(label=label): + 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 chained(x): - return 0 < x < 10 + def test_udf_transpile_boolean_body_requires_boolean_return_type(self): + # 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 - 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.subTest(value=value): - df = self.spark.createDataFrame([Row(a=value)], schema=schema) - [row] = df.select(pudf("a")).collect() - self.assertEqual(row[0], expected) + def ternary_compare(c, x, y): + return (x if c > 0 else y) == 1 + + def eq_true(a, b): + # 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 comparison operands. + return (a < b) <= (b < c) + + cases = [ + (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)]), + ] + 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_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): + # 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_COMPARISONS, + _MAX_MATERIALIZED_WEIGHT, + CatalystTranspiler, + _materialized_weight, + ) + + transpiler = CatalystTranspiler() + + def compare_node(src): + return _ast.parse(src, mode="eval").body + + 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 + return " < ".join(["x"] + [str(i) for i in range(n)]) + + # ---- Per-chain operator cap ---- + lower(chain_src(_MAX_COMPARISON_OPS)) # at the cap, succeeds + 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) + + # 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. 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 + + lower(nested_src(2)) + # 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) + self.assertTrue( + all( + len(n.ops) <= _MAX_COMPARISON_OPS + for n in _ast.walk(compare_node(deep)) + if isinstance(n, _ast.Compare) + ) + ) + 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 = "%s %s (%s)" % (leaf, op, src) + return src + + lower(right_nested_src(2)) + deep_right = right_nested_src(20) + self.assertLess( + sum( + len(n.ops) + for n in _ast.walk(compare_node(deep_right)) + if isinstance(n, _ast.Compare) + ), + _MAX_LOWERED_COMPARISONS, + ) + 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. + 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 @@ -1084,9 +1551,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; 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, 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 @@ -1094,12 +1562,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]) @@ -2012,6 +2490,10 @@ 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 + 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: @@ -2038,6 +2520,26 @@ 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 the `and`/`or` gate accepts it and it + # 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]), ( nested, B, @@ -2091,6 +2593,15 @@ 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 is fully lowered, so the EvalPython node + # disappears for it too. + bounded = lambda x: 0 < x < 10 # noqa: E731 + 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_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. @@ -2488,6 +2999,25 @@ 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 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") + 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)] + ), + [False], + "NaN == NaN is False, as in Python", + ) 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 02c2ba9a25bd..90e0699480b3 100644 --- a/python/pyspark/sql/transpile.py +++ b/python/pyspark/sql/transpile.py @@ -35,6 +35,81 @@ 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`` -- 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 + 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. ``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 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 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 + unconditionally. ``Expression.throwable`` does NOT help -- this rule never + consults it, so SPARK-58627 does not cover this one. +* 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 +``(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. + +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. ``len(None)`` raises (via ``raise_error`` on NULL), matching CPython. @@ -94,8 +169,10 @@ import ast import contextlib +import functools import inspect import itertools +import operator import sys import textwrap import threading @@ -154,6 +231,61 @@ def _wider_numeric(lc: str, rc: str) -> str: return "integer" +# The most operators we will lower in one chained comparison (64 operators over +# 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"), + ast.GtE: (operator.ge, ">="), +} + + class AbstractTranspiler(object): """Base class for transpilers. All experimental.""" @@ -187,6 +319,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 @@ -214,17 +383,21 @@ 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. + 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: 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 + # 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) @@ -238,6 +411,48 @@ 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. ``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* 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"``. + """ + 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. @@ -309,6 +524,15 @@ def __init__(self) -> None: self._category_cache: dict[int, str] = {} self._nonnull_params: frozenset = frozenset() self._null_guard_emitted: bool = False + # 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; @@ -380,7 +604,7 @@ def _safe_category(self, params: List[str], node: Optional[ast.AST]) -> Optional return _wider_numeric(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). @@ -575,6 +799,72 @@ def _lower_value_compare( self._null_guard_emitted = True 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``. + + 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 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 wrap each link in ``NOT link`` with no ``coalesce``. + + 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: + 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 + # `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) + 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(): + 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 by the caller. + raise UnsupportedOperationException( + f"comparison operator {type(op).__name__} is not supported by the transpiler" + ) + def _category(self, params: List[str], node: ast.AST) -> str: """Infer the scalar category of ``node`` under the current ``self._param_categories`` assumption (set per input-type variant). @@ -670,7 +960,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) @@ -692,9 +982,17 @@ def branch_category(b: ast.AST) -> Optional[str]: "ternary with all-None branches has no usable column category" ) return result_cat + case ast.Compare(): + # 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): - # 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" @@ -752,49 +1050,35 @@ 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 `&` / `|`. 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. + # + # 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( - "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): - 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. @@ -828,60 +1112,45 @@ 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: + # `in` / `not in` are container membership, not a chainable value + # 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( - "chained comparisons (e.g. `a < b < c`) are not supported by the transpiler" + "`in` / `not in` is not supported by the transpiler; the " + "UDF falls back to interpreted Python" ) - 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 _: - raise UnsupportedOperationException( - f"comparison operator {type(ops[0]).__name__} " - "is not supported by the transpiler" - ) + if len(ops) > _MAX_COMPARISON_OPS: + raise UnsupportedOperationException( + f"chained comparison has {len(ops)} operators; the " + f"transpiler lowers at most {_MAX_COMPARISON_OPS}, so the " + "UDF falls back to interpreted Python" + ) + # `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%). + # + # 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) + 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 @@ -1017,6 +1286,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 # Reset per-transpilation null-narrowing state. self._nonnull_params = frozenset() self._null_guard_emitted = False @@ -1036,6 +1309,13 @@ def _transpile_from_ast( "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: #