spark-test commented on code in PR #59216:
URL: https://github.com/apache/spark/pull/59216#discussion_r4417732379
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -491,39 +506,478 @@ 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. `&` 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.
+ #
+ # Both of `And`'s two laziness implementations are exercised:
`doGenCode`
Review Comment:
**Nit (P3):** A chain now lowers to `CASE WHEN NOT link1 THEN false ... ELSE
linkN END` (transpile.py:1080-1085), so this FALLBACK / NO_CODEGEN loop
exercises `CaseWhen`'s laziness, and no `And` is evaluated here. Line 512 ("`&`
short-circuits at evaluation time too, so the fold matches") has the same
problem. Lines 520-521 still point to
`test_udf_transpile_boolop_short_circuits` for optimizer survival, which
`test_udf_transpile_chain_survives_optimizer` now covers for chains. The and/or
short-circuit test runs only under the default factory mode, so no test in the
transpiler suites exercises `And.eval`'s laziness, which the module docstring
says the `and` / `or` fold relies on. A regression that made interpreted `And`
eager would raise `cannot compare NULL` where Python returns False or True, and
no transpiler test would fail.
Related leftovers of the earlier `&` chain fold:
- `test_udf_transpile_chain_survives_optimizer` (548-556) says "Before this
fix the chain folded to `&`" and that a later link "used to be pushed below the
Aggregate". Relative to the merge base, chains never lowered at all, so that
history never shipped.
- Its worked example (558, 576) uses the group `(b=1, c=NULL, sa=5)`. The
rows `(1, 1, None)` and `(5, 1, None)` fall into the same group, so `sa = 6`.
The assertion still holds.
- The module docstring sentence "where it folds to `&` / `|` (it does not --
it folds to a CASE WHEN) it would manufacture ..." (transpile.py:108-111)
describes a fold that the same sentence denies.
Suggested fix: here, describe `CaseWhen` laziness and point at the
optimizer-survival test. In that test, recast the history as shapes that would
break an `&` fold (and do break the `and` spelling), and correct `sa` to 6. In
the module docstring, state the point directly for the user-written `and` /
`or` form. Finally, run `test_udf_transpile_boolop_short_circuits` under both
FALLBACK and NO_CODEGEN.
**Verification:**
- **Regression:** and/or cases whose second operand is NULL, after the first
operand has already decided the result, return Python's answer without raising
under both generated and interpreted expression evaluation.
- **Inspection:** The chain test comments, the aggregate repro's worked
example and the module docstring describe the CASE WHEN fold and the group sum
the test data produces.
##########
python/pyspark/sql/transpile.py:
##########
@@ -451,6 +658,153 @@ 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``.
+
+ 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, 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.
+ """
+ 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"
+ )
+ 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
+ # `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 _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
Review Comment:
**Nit (P3):** `_comparison_operand_weight` charges every non-`Compare`
operand as 1, and `_lower_compare_pair` adds weight only at `_compare_depth ==
1`. A comparison nested under a ternary (in its test, body or else branch) or a
`len()` argument is therefore lowered at depth 2 or more and is never charged,
neither as its own link nor through its parent's operand weight. The enclosing
link still references the ternary's Column twice (three times for `==`). The
new `case ast.Compare(): return "bool"` arm also lets such a ternary pass the
category check and the `BooleanType` return-type check.
Example, declared `BooleanType`: `R(0) = (x < 0)` and `R(k) = (x < 0) <
((R(k-1)) if x > 0 else (x < 0))`. At depth 20 the body lowers with 81 links
(under 128) and a charged weight of 22 (65536 allowed). I ran this file's
lowering against stub Columns that record the shared DAG, and the expanded tree
comes to about 9.4e7 nodes. The `==` form is charged a constant 37 and triples
per level (about 1.3e9 nodes at depth 15). On every UDF application,
`UserDefinedPythonFunction.builder` converts each option through
`ColumnNodeToExpressionConverter`, which recurses into shared children without
memoizing. So the UDF constructs as transpiled, then hangs or runs the driver
out of memory on first use, outside `_transpile_func`'s fallback. That is the
outcome the `_MAX_LOWERED_COMPARISON_WEIGHT` comment says this budget prevents.
The merge base refused these shapes. Two caveats on scope: I did not run the
JVM expansion end to end, since the node counts come from the stubbed lowering
plus the
converter source. And the shape needs the experimental flag with ANSI plus a
deliberately deep nest, so ordinary bodies are not affected.
Suggested fix: make the operand weight recurse through every operand kind
that can contain a comparison. That means an `IfExp`'s test, body and else
branch, plus `not`, `len()` and arithmetic, each weighted by the number of
references its lowering emits, memoized by node identity like
`_comparison_weight`. Then add caps-test cases that expect the weight refusal
for ternary-routed nests.
This extends an [existing
discussion](https://github.com/apache/spark/pull/59216#pullrequestreview-5505733895).
**Verification:**
- **Regression:** Lowering a comparison nest of about 20 levels routed
through a ternary operand, in its test, body and else positions and with both <
and ==, raises the materialized-weight refusal instead of returning a column.
- **Compatibility:** A flat chain at the 64-operator cap, the shallow nested
and right-nested shapes, and single-level comparisons over ternaries still
lower, and their transpiled results still match interpreted Python.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -491,39 +506,478 @@ 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. `&` 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.
+ #
+ # 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 = "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`. 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
+ # 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
+
+ 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. 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(<boolean> 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):
+ # 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 (
+ _MAX_COMPARISON_OPS,
+ _MAX_LOWERED_COMPARISON_WEIGHT,
+ _MAX_LOWERED_COMPARISONS,
+ CatalystTranspiler,
+ )
+
+ 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 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
+ 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)
+ # 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))
Review Comment:
**Nit (P3):** No case in this test makes the 128-link budget the refusing
bound. `nested_src(10)` is refused by the weight check at its first top-level
link (charged about 1.1e7), which is why this line asserts the weight message.
The right-nested and `==` right-nested cases are weight refusals too, and
`_MAX_LOWERED_COMPARISONS` appears only as the `assertLess` bound on an AST
operator count. UDF-level tests stay far below 128 links, and `lower()` resets
the per-lowering fields by hand, so nothing drives the budgets through
`_transpile_from_ast` across input-type variants. Deleting the link-count check
in `_lower_compare_pair` (transpile.py:695-701), or the per-variant resets
(1221-1224), would leave every current assertion passing.
The link budget is still load-bearing. In the stubbed lowering, only the
link budget refuses `nested_src(5)` (156 links, weight 21,832). The same holds
for a chain with a ternary interior operand, `(x < 0) < ((<nest>) if x > 0 else
(x < 0)) < (x < 0)`, whose links double per level at a constant weight. Without
the reset, a variant refused inside a link leaves `_compare_depth` raised, so
the next variant skips weight charging. The PR description says both caps are
tested "at N and N+1". Two comments here also misstate the shapes:
- `right_nested_src(20)` has 2d+1 = 41 links, not "~d + 1" / "21 links"
(934, 937).
- `nested_src`'s `(x < 0)` operands have one operator, not "exactly two"
(913).
Suggested fix: add a link-bound case such as `nested_src(5)` that asserts
the link-budget message, plus one just under 128 links that still lowers. Add
UDF-construction cases that depend on the per-variant reset: an untyped body
whose two variants each need about 100 links must yield both options, and a
variant refused part-way through a link must not leave the next one uncharged.
Then correct the two comments.
**Verification:**
- **Regression:** A body that needs more than 128 comparison links while
staying under the materialized-weight bound is refused with the link-budget
message, and a body just under 128 links still lowers.
- **Regression:** An untyped body whose numeric and string variants each
stay under the budgets yields both options. A variant refused part-way through
a link does not leave the next variant uncharged against the weight bound.
- **Inspection:** The caps-test comments give the right-nested shape's link
count as 2d+1 and describe the per-chain-cap condition that the neighbouring
assertion checks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
