HyukjinKwon commented on code in PR #59216:
URL: https://github.com/apache/spark/pull/59216#discussion_r4260389198
##########
python/pyspark/sql/transpile.py:
##########
@@ -673,60 +890,49 @@ 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`
is
+ # equality-based and never raises on a `None` probe. Not
lowered
Review Comment:
**Nit (P3):** Python's `in` is not always equality-based, and it can raise
on a `None` probe. For a str container it is substring containment, and `None
in "abc"` raises `TypeError: 'in <string>' requires string as left operand, not
NoneType`. The comment in `test_udf_transpile_comparison_fallbacks` already
says str `in` is substring containment. Tuples and lists do use equality.
This comment states the semantics a future `in` lowering would follow, so an
implementation built from it would return a value where Python raises. Consider
saying that the semantics depend on the container, which is why membership
needs its own NULL analysis.
**Verification:**
- **Inspection:** The refusal comment's description of Python `in` matches
CPython for str and tuple containers.
##########
python/pyspark/sql/transpile.py:
##########
@@ -132,6 +199,45 @@
from pyspark.sql._typing import DataTypeOrString
+# 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<b) < L(d-1) < (a<b))"`:
+# UDF construction took 8.0s at d=8, 32.0s at d=10 and 63.4s at d=11 -- exact
+# doubling, from a 249-character body, and extrapolating to hours by d=20. Each
+# link is a py4j round trip on the driver, so that is a wedged driver at
+# `udf(...)` time, which `_transpile_func`'s blanket `except` cannot recover
+# from. Counting links once across the whole lowering turns it into a refusal
+# and a fallback to interpreted Python.
+#
+# 128 leaves flat chains at the cap untouched (64 links) while holding the
worst
+# nested case to well under a second.
+_MAX_LOWERED_COMPARISONS = 128
Review Comment:
**Nit (P3):** This budget is charged once per link in `_lower_compare_pair`,
but each lowering references its operand Columns more than once.
`_lower_value_compare` uses each operand in both the NULL guard and the
comparison, and `_lower_eq` uses each one three times.
`UserDefinedPythonFunction` converts every option with
`ColumnConversions.expression`, and `ColumnNodeToExpressionConverter` recurses
into every child with no memoization, so a shared Column DAG is expanded into a
tree.
The new `case ast.Compare(): return "bool"` arm lets comparisons nest at any
depth. A right-nested body such as `(x < 0) < ((x < 0) < (...))` is charged
about 2d+1 links but materializes on the order of 2^d nodes (3^d with `==`). At
depth 20 (41 links, about 250 characters) that is roughly 4.6e7 nodes. The
expansion happens on the UDF's first call, outside `_transpile_func`'s fallback
`except`, so for these shapes the budget does not deliver the refusal and
fallback this comment promises.
I did not reproduce the JVM hang end to end; the mechanism comes from the
converter source and a Python-side node count. Inputs like this are
pathological and the path is experimental, but they are the input class the
budget was added for.
**Recommended change:** Charge the comparison budget with the materialized
(tree) size each comparison lowering emits, counting each operand's expanded
size once per reference, and refuse once a body's materialized size passes a
bound; correct the budget comments; and add transpile unit tests for nested
comparison shapes whose link count grows linearly.
**Why this works:** Measuring what the JVM materializes, rather than how
many links Python lowers, makes every accepted option bounded in size, so
pathological nesting is refused at UDF construction and falls back to
interpreted Python, as the PR intends.
**Scope:** Make the comparison budget measure materialized expression size
and cover linearly nested comparison shapes in tests.
**Compatibility:** Preserve behavior outside the stated repair boundary.
**Risks:** A size bound set too low could refuse realistic flat chains; it
must keep the 64-operator flat chain and the existing shallow nested cases
lowering.
**Constraints:** Keep resetting the budget per input-type variant in
_transpile_from_ast. Account for any structure the chain-fold repair adds.
**Success:** A right-nested `(x < 0) < (... (x < 0))` body and a `((x == 1)
== True) == ...` body of about 20 levels are refused at UDF construction and
fall back to interpreted Python. A flat 64-operator chain and the existing
shallow nested bodies still lower.
**Verification:**
- **Regression:** Nested comparison bodies whose link count grows linearly
but whose materialized size grows exponentially are refused by the budget.
- **Compatibility:** The existing cap-boundary cases (flat chain at the cap,
shallow and deep doubling nests) keep their current outcomes.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -491,39 +506,345 @@ 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_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())
+ left = self.spark.createDataFrame([(1, -5)], "k long, a long")
+ right = self.spark.createDataFrame([(1, None)], "k long, c long")
+ joined = left.join(right, "k")
+ self.assertEqual([False], [r[0] for r in
joined.select(and_udf("a", "c")).collect()])
+
+ # Filtering must return Python's answer rather than raising. This
is
+ # the assertion SPARK-58627 fixed; see above.
+ filtered = joined.filter(and_udf("a", "c"))
Review Comment:
**Nit (P3):** The join block asserts only `[False]` for the select and `[]`
for this filter. Interpreted Python returns the same values. If the JVM
discards the transpiled options for the joined query, every assertion still
passes, and the pushed-down-guard regression this test exists to catch goes
unnoticed. The block also rebuilds `and_udf` without re-asserting `transpiled`.
The file's `_vals` helper already explains why `transpiled` alone is not
enough, and it checks `self._eval_python_count(...) == 0`. The same check here,
for both the joined select and the filter, would make this test fail when the
fallback runs.
**Verification:**
- **Regression:** The joined select and filter each assert zero EvalPython
nodes alongside their result checks.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -491,39 +506,345 @@ 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_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())
+ left = self.spark.createDataFrame([(1, -5)], "k long, a long")
+ right = self.spark.createDataFrame([(1, None)], "k long, c long")
+ joined = left.join(right, "k")
+ self.assertEqual([False], [r[0] for r in
joined.select(and_udf("a", "c")).collect()])
+
+ # Filtering 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())
+
+ 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 (SPARK-56925), so these cases
all
Review Comment:
**Nit (P3):** This comment, and the one in
`test_udf_transpile_lowering_caps` ("refused outright, deferred to
SPARK-56925"), name SPARK-56925 as the follow-up for `in` / `not in`. In this
tree SPARK-56925 is the already-landed bare `if x` truthiness work:
`transpile.py` cites it for `_truthiness_col`, and
`test_udf_transpile_bare_if_truthiness` says those bodies now transpile under
it. Please point both comments at the ticket that tracks membership support
(the PR description says it is deferred to its own ticket), or drop the
reference.
**Verification:**
- **Inspection:** Both comments reference the membership follow-up rather
than the truthiness ticket.
##########
python/pyspark/sql/transpile.py:
##########
@@ -673,60 +890,49 @@ 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`
is
+ # equality-based and never raises on a `None` probe. Not
lowered
+ # yet, so refuse them and fall back.
+ if any(isinstance(o, (ast.In, ast.NotIn)) for o in ops):
raise UnsupportedOperationException(
- "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"
+ )
+ # Python evaluates `a OP1 b OP2 c` as `(a OP1 b) and (b OP2 c)`
+ # with `b` evaluated once, so the fold below re-lowers each
+ # interior operand once per adjacent pair. Safe for the operand
+ # grammar we accept, all of it deterministic and side-effect
free.
+ # It does mean the body READS a parameter once per adjacent
pair,
+ # which SPARK-58626 is what makes safe: `_udf_param_N` is a
+ # reference, and `ConvertToCatalyst` gives an argument read
more
+ # than once one column per row -- or declines to lower the call
+ # where no column fits -- so a non-deterministic argument is
drawn
+ # once and every read agrees. While each read still carried
its own
+ # copy of the argument, `50 < x < 60` over `rand(7) * 100` was
~30%
+ # true where Python is ~10%.
+ #
+ # Folding into nested CASE WHENs instead -- atomic to predicate
+ # splitting -- was rejected: it costs pushdown for every
+ # `and` / `or` UDF, including those over non-nullable columns
that
+ # need no guard at all.
+ #
+ # See the module docstring for the short-circuit and NULL-guard
+ # semantics this fold relies on.
+ nodes = [left, *comps]
+ return functools.reduce(
Review Comment:
**Non-blocking (P2):** Before this PR every chained comparison fell back to
interpreted Python, so it always returned Python's answer. This fold turns a
chain into a splittable `And` of NULL-guarded links. That is only correct while
every later link stays behind the links that gate it, and the optimizer passes
that run after `ConvertToCatalyst` do not keep that order:
- The Aggregate arm of `PushDownPredicates` splits conjuncts and pushes the
deterministic ones that reference only grouping keys, explicitly including
throwable ones (Optimizer.scala:2672-2696). Take `udf(lambda sa, b, c: sa < b <
c)` applied as `groupBy("b", "c").agg(sum("a").alias("sa")).filter(...)`:
`link(b < c)` runs below the Aggregate. A group with `b = 1, c = NULL, sa = 5`
then raises `cannot compare NULL` where Python returns False.
- `BooleanSimplification` common-factor elimination and subexpression
elimination hoist a shared link, so `(a < b < c) or (d < b < c)` on `(5, 1,
None, 7)` raises. Absorption drops a guard instead: `(a < b < c) or (b < c)` on
`(None, 0, 1)` returns True where Python raises.
- File-source partition filtering
(`DataSourceUtils.getPartitionFiltersAndDataFilters`) has no `throwable` check,
so a later link over a partition column runs during partition pruning.
The module docstring records these as OPEN for the `and` spelling. That
spelling already transpiled before this PR, but chains did not, so this PR is
what exposes chains to them. The comment above rejects nested CASE WHENs
because they "cost pushdown for every `and` / `or` UDF". That cost does not
apply to chains, which had no lowering to lose. The issue only affects users
who enable the experimental flag with ANSI on. The evaluation-time
short-circuit itself is correct.
**Recommended change:** Fold chain links with a sequential CASE WHEN (`CASE
WHEN NOT link1 THEN false ... WHEN NOT link(n-1) THEN false ELSE linkN END`)
instead of `&`, update the module docstring and the chain-arm comment to
describe that fold, and add transpile unit tests for the optimizer shapes that
break the `&` fold.
**Why this works:** Each later link sits in a CASE WHEN branch that runs
only after every earlier link is true. Every link is total (never NULL), so
`NOT link` is never NULL and the sequence mirrors Python's short-circuit
exactly. 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. When pushdown moves the whole predicate, it moves as one unit
evaluated on the same key values Python sees.
**Scope:** Replace the chain's `&` fold with a sequential CASE WHEN and pin
the optimizer shapes that broke it.
**Compatibility:** Chained comparisons still transpile to native expressions
and keep Python's evaluation-time short-circuit semantics.
**Risks:** A chain predicate no longer splits for partial pushdown and moves
only as a whole; this is no regression versus the merge target, where chains
did not transpile.
**Constraints:** Each link from _lower_compare_pair must stay total (never
NULL), since `NOT link` must be non-NULL for the CASE WHEN to mirror Python.
The comparison budget must keep charging every lowered link and any structure
the new fold adds.
**Success:** A chain UDF used as a filter over groupBy/agg whose later link
references only grouping keys returns Python's answer for groups with a NULL
key instead of raising. `(a < b < c) or (d < b < c)` on (5, 1, None, 7) returns
False with transpilation on, matching Python. `(a < b < c) or (b < c)` on
(None, 0, 1) raises like Python instead of returning True. Chains still lower
with no EvalPython node and keep their existing evaluation-time short-circuit,
NULL-pattern, cap and NaN behavior.
**Verification:**
- **Regression:** A chain UDF filtered over an Aggregate on its grouping
keys, with a NULL grouping value whose earlier link is false, returns Python's
answer and runs fully lowered.
- **Regression:** Disjunctions of chains that share a link return Python's
answer for the common-factor and absorption repros recorded in the module
docstring.
- **Compatibility:** Existing chain coverage (short-circuit under both
codegen factory modes, every `==` NULL pattern, operator mixes, cap boundaries,
EvalPython elision, NaN pins and the hypothesis chain templates) keeps passing.
--
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]
