HyukjinKwon commented on code in PR #59339:
URL: https://github.com/apache/spark/pull/59339#discussion_r4242511500
##########
python/pyspark/sql/udf.py:
##########
@@ -329,6 +338,24 @@ def _conf_is_true(key: str, default: Optional[str] = None)
-> bool:
self._transpiled_param_names = []
self._transpiled_input_categories = []
self._positional_only_param_names = frozenset()
+ if pending_guard_warning and self.transpiled:
+ # Advice, not an error: the lowering is fine, so warn in its own
handler --
+ # under warnings-as-errors the handler above would have thrown the
+ # lowering away, which is the wrong fix. "may not" because a
non-nullable
+ # bound column (NullPropagation) or ``asNondeterministic()`` can
still drop
+ # the check after this point.
+ try:
+ warnings.warn(
+ f"Transpiled UDF {func} still checks for NULL in "
+ f"{', '.join(pending_guard_warning)} and raises if it
finds one, so "
+ "Spark may not be able to push a filter on this UDF
through a join "
+ "or combine it with an adjacent filter. Guard the
parameter "
+ "(`if x is not None:`) or bind a non-nullable column to
drop the "
+ "check; binding one already drops it.",
Review Comment:
**Nit (P3):** As written, the closing advice contradicts itself: it suggests
binding a non-nullable column to drop the check, then says binding one already
drops it. It also says "the parameter" even when the label lists several names.
Splitting it into one action and one caveat would read better, for example:
"Guard the parameters named above (e.g. `if x is not None:`); if the bound
column is non-nullable, the optimizer already folds the check away and this
warning can be ignored."
**Verification:**
- **Inspection:** The emitted warning text reads as a single coherent
instruction plus caveat for both single- and multi-parameter labels.
- **Behavior:** The existing warning assertions (label text, `is not None`
hint, silence for guarded UDFs) still hold.
##########
python/pyspark/sql/transpile.py:
##########
@@ -239,8 +329,116 @@ class CatalystTranspiler(AbstractTranspiler):
variety = "catalyst"
def __init__(self) -> None:
+ # Category inference depends on the per-variant assumptions set in
+ # ``_transpile_from_ast``; both are reset there per variant.
self._param_categories: dict[int, str] = {}
self._category_cache: dict[int, str] = {}
+ # Instance attrs (not class defaults): ``_raise_on_null`` does ``|=``
on
+ # ``_pending_null_guards``, which on a class-default set would mutate
the
+ # class dict and leak state across UDFs. Both reset per variant in
+ # ``_transpile_from_ast``.
+ self._non_null: frozenset = frozenset()
+ self._pending_null_guards: frozenset = frozenset()
+ #: Raising NULL checks the last lowered variant needed; empty when none
+ #: survived. ``_transpile_func`` unions this across kept variants and
warns
+ #: once per UDF. Assigned, not accumulated, so nothing leaks between
UDFs.
+ self.null_guards: frozenset = frozenset()
+
+ @contextlib.contextmanager
+ def _narrowed(self, proven_non_null: frozenset) -> Iterator[None]:
+ """Lower the enclosed nodes with ``proven_non_null`` added to what is
known.
+
+ Save/restore suffices: lowering walks straight down, so a fact's scope
is
+ exactly the subtree we recurse into.
+ """
+ previous = self._non_null
+ self._non_null = previous | proven_non_null
+ try:
+ yield
+ finally:
+ self._non_null = previous
+
+ def _is_never_null(self, params: List[str], node: ast.AST) -> bool:
+ """Whether ``node`` provably cannot evaluate to NULL here.
+
+ Answers "do we still need a NULL check?" Every arm is an explicit
proof and
+ the catch-all is ``False``: a missed proof costs a redundant check, a
wrong
+ one drops a check Python needs. Expressions only.
+ """
+ match node:
+ case ast.Constant(value=value):
+ return value is not None
+ case ast.Name(id=name):
+ return name in params and name in self._non_null
+ case ast.Compare(ops=[op]):
+ # is/is not lower to isNull/isNotNull; ==/!= cover NULL with
boolean
+ # literals; ordering raises on NULL rather than returning it.
+ return isinstance(
+ op, (ast.Is, ast.IsNot, ast.Eq, ast.NotEq, ast.Lt,
ast.LtE, ast.Gt, ast.GtE)
+ )
+ case ast.UnaryOp(op=ast.Not()):
+ # Either inverts a non-NULL operand or coalesces against a
literal.
+ return True
+ case ast.UnaryOp(operand=operand):
+ return self._is_never_null(params, operand)
+ case ast.BoolOp(values=values):
+ return all(self._is_never_null(params, v) for v in values)
+ case ast.BinOp(left=left, right=right):
+ # Catalyst arithmetic and concat propagate NULL from either
side.
+ return self._is_never_null(params, left) and
self._is_never_null(params, right)
+ case ast.Call(func=ast.Name(id="len"), args=[_], keywords=[]):
+ # The lowering raises on NULL or returns ``length``; neither
is NULL,
+ # so a parent (``len(x) > 0``) must not add a second check. A
non-string
+ # argument never produces a column -- that variant is dropped
first.
+ return True
+ case ast.IfExp(test=test, body=body, orelse=orelse):
+ # Each branch is lowered under the facts its test proves, so
ask
+ # about the branches the same way.
+ when_true, when_false = _null_facts(test)
+ with self._narrowed(when_true):
+ body_ok = self._is_never_null(params, body)
+ with self._narrowed(when_false):
+ return body_ok and self._is_never_null(params, orelse)
+ case _:
+ return False
+
+ @staticmethod
+ def _any_null(columns: List[Column]) -> Column:
+ """``columns[0] IS NULL OR ...``; raises on an empty list."""
+ return functools.reduce(operator.or_, (column.isNull() for column in
columns))
+
+ def _raise_on_null(
+ self,
+ params: List[str],
+ operands: List[Tuple[ast.AST, Column]],
+ label: str,
+ message: str,
+ otherwise: Column,
+ ) -> Column:
+ """``otherwise``, guarded so that a NULL operand raises ``message``
instead.
+
+ The only place that emits a ``raise_error``; future lowerings needing
one
+ (SPARK-55210's arithmetic/unary/concat guards) must come through here
or they
+ skip the narrowing and the warning. Operands proven non-NULL
contribute no
+ check, and with none left the guard is dropped -- ``RaiseError`` is
+ ``throwable`` (SPARK-58627), so a plan holding one cannot be pushed
through a
+ join. ``label`` names the construct plus the parameters checked
(reached
+ through the whole operand, so ``a > (b + 1)`` names ``b``).
+ """
+ checked = [(node, c) for node, c in operands if not
self._is_never_null(params, node)]
+ if not checked:
+ return otherwise
+ named = sorted(
+ {
+ inner.id
+ for node, _ in checked
+ for inner in ast.walk(node)
+ if isinstance(inner, ast.Name) and inner.id in params
Review Comment:
**Nit (P3):** This filter keeps every parameter that appears anywhere in a
checked operand, including parameters an enclosing guard has already proven
non-NULL. Take `lambda a, b: a is not None and a + b > 0`. `a + b` is still
checked because of `b`, so the label becomes "comparison `>` on a, b" and the
warning asks the user to guard `a`, which they already did.
`test_udf_transpile_warns_only_when_a_null_check_survives` states this property
("It must name the parameter still checked, not both"), but it only covers a
bare-name operand. To match that intent, also filter out `inner.id in
self._non_null`, and skip sub-expressions that `_is_never_null` already proves,
such as a nested `len(...)` that has its own label. A composite-operand case in
the warning test would then pin the behavior.
**Verification:**
- **Regression:** A UDF whose checked operand mixes a proven and an unproven
parameter warns about only the unproven one.
- **Behavior:** An unguarded multi-parameter comparison names every
parameter it checks, and the existing single-name labels are unchanged.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -2634,6 +2663,440 @@ def bytes_ident(x: bytes):
self.assertFalse(UserDefinedFunction(bool_add,
LongType()).transpiled)
self.assertTrue(UserDefinedFunction(bytes_ident,
BinaryType()).transpiled)
+ def _transpiled_plan(self, func, return_type, schema):
+ """The optimized plan of ``func`` applied to every column of
``schema``.
+
+ The lowered NULL checks are what these tests are about, so they have
to be
+ read off the plan: the values a guarded UDF returns are the same
whether or
+ not the plan carries a check that can never fire.
+
+ Asserting the option was actually APPLIED is what keeps every
+ ``assertNotIn(..., plan)`` built on this from passing for the wrong
reason
+ -- an interpreted plan contains none of the strings we look for, so a
UDF
+ that quietly fell back would satisfy them all (see ``_vals``, which
guards
+ the same way for the same reason).
+ """
+ with self.sql_conf(_TRANSPILE_ON):
+ u = self._transpiled_udf(func, return_type)
+ df = self.spark.createDataFrame([], schema)
+ projected = df.select(u(*df.columns))
+ self.assertEqual(0, self._eval_python_count(projected), str(func))
+ return projected._jdf.queryExecution().optimizedPlan().toString()
+
+ def test_udf_transpile_drops_null_checks_a_guard_already_made(self):
+ # SPARK-58628. Python raises TypeError on `None > 0` where Spark
returns
+ # NULL, so an ordering comparison carries a check that raises. When
the UDF
+ # has already tested the parameter the check can never fire, and
leaving it
+ # in is not free: RaiseError is throwable (SPARK-58627), which stops
the
+ # optimizer from moving a predicate containing it. Every shape below
proves
+ # non-NULL-ness, so none of them should lower to a raise -- and all of
them
+ # must still agree with Python, including on the None row.
+ def if_guard(x):
+ if x is not None:
+ return x > 0
+ else:
+ return None
+
+ def else_of_is_none(x):
+ if x is None:
+ return None
+ else:
+ return x > 0
+
+ def not_is_none(x):
+ if not (x is None): # noqa: E714 the `not` form is the point
+ return x > 0
+ else:
+ return None
+
+ def reversed_test(x):
+ if None is not x:
+ return x > 0
+ else:
+ return None
+
+ def ne_none_test(x):
+ # The `!= None` spelling of the same guard. PEP 8 prefers `is not
None`,
+ # but both are common and both must narrow, or the equality form
keeps a
+ # raising check and gets told to add the guard it already has.
+ if x != None: # noqa: E711
+ return x > 0
+ else:
+ return None
+
+ def eq_none_test(x):
+ if x == None: # noqa: E711
+ return None
+ else:
+ return x > 0
+
+ ternary = lambda x: (x > 0) if x is not None else None # noqa: E731
+ short_circuit_and = lambda x: x is not None and x > 0 # noqa: E731
+ short_circuit_or = lambda x: x is None or x > 0 # noqa: E731
+ # `<=` / `>=` share `_lower_value_compare` but are otherwise only
exercised
+ # where a guard IS emitted, so cover the narrowed side for them too.
+ lte_guarded = lambda x: (x <= 0) if x is not None else None # noqa:
E731
+ gte_guarded = lambda x: (x >= 0) if x is not None else None # noqa:
E731
+ lt_guarded = lambda x: (x < 0) if x is not None else None # noqa: E731
+
+ rows = [(-1,), (0,), (5,), (None,)]
+ cases = [
+ (if_guard, [False, False, True, None]),
+ (else_of_is_none, [False, False, True, None]),
+ (not_is_none, [False, False, True, None]),
+ (reversed_test, [False, False, True, None]),
+ (ne_none_test, [False, False, True, None]),
+ (eq_none_test, [False, False, True, None]),
+ (ternary, [False, False, True, None]),
+ (lte_guarded, [True, True, False, None]),
+ (gte_guarded, [False, True, True, None]),
+ (lt_guarded, [True, False, False, None]),
+ # `and` / `or` return the guard's own result for None, as Python
does.
+ (short_circuit_and, [False, False, True, False]),
+ (short_circuit_or, [False, False, True, True]),
+ ]
+ for func, expected in cases:
+ with self.subTest(func=getattr(func, "__name__", "lambda")):
+ self.assertEqual(
+ expected,
+ [func(v) for (v,) in rows],
+ "the expectation should be what Python actually does",
+ )
+ self.assertEqual(expected, self._vals(func, BooleanType(), "a
long", rows))
+ plan = self._transpiled_plan(func, BooleanType(), "a long")
+ self.assertNotIn("raise_error", plan)
+
+ def test_udf_transpile_keeps_the_null_check_it_needs(self):
+ # The other side of the previous test: nothing here proves the
parameter
+ # non-NULL, so the check has to stay and the UDF has to raise like
Python.
+ # `x is None and ...` is the interesting one -- an `and` being false
says
+ # only that SOME operand was false, so it narrows nothing.
+ unguarded = lambda x: x > 0 # noqa: E731
+ wrong_way_and = lambda x: x is None and x > 0 # noqa: E731
+ two_columns = lambda a, b: a > b # noqa: E731
+
+ for func, schema in ((unguarded, "a long"), (two_columns, "a long, b
long")):
+ with self.subTest(func="lambda", schema=schema):
+ plan = self._transpiled_plan(func, BooleanType(), schema)
+ self.assertIn("raise_error", plan)
+ boolean = BooleanType()
+ self._raises(unguarded, "a long", [(None,)], "cannot compare null",
boolean)
+ self._raises(two_columns, "a long, b long", [(1, None)], "cannot
compare null", boolean)
+ # Each ordering operator names itself in the guard's message and in the
+ # warning's label; only `>` was covered before.
+ lt_zero = lambda x: x < 0 # noqa: E731
+ lte_zero = lambda x: x <= 0 # noqa: E731
+ gte_zero = lambda x: x >= 0 # noqa: E731
+ for func, op_text in ((lt_zero, "`<`"), (lte_zero, "`<=`"), (gte_zero,
"`>=`")):
+ with self.subTest(op=op_text):
+ with self.sql_conf(_TRANSPILE_ON):
+ _, warned = self._udf_and_warnings(func, boolean)
+ self.assertIn(f"comparison {op_text} on x", warned)
+ self._raises(func, "a long", [(None,)], f"operator {op_text}",
boolean)
+ # `x is None and x > 0` is False for every non-None x without
evaluating the
+ # comparison, and raises for None -- exactly like Python.
+ with self.sql_conf(_TRANSPILE_ON):
+ u = self._transpiled_udf(wrong_way_and, BooleanType())
+ self.assertIn("raise_error", str(u.transpiled[0]))
+ df = self.spark.createDataFrame([(1,)], "a long")
+ self.assertEqual([False], [r[0] for r in
df.select(u("a")).collect()])
+ self._raises(wrong_way_and, "a long", [(None,)], "cannot compare
null", boolean)
+
+ def test_udf_transpile_narrows_only_the_proven_parameter(self):
+ # With two parameters the narrowing has to be per-name: guarding `a`
must
+ # not license dropping the check on `b`. Both guarded, no check at all.
+ one_guarded = lambda a, b: a is not None and a > b # noqa: E731
+ both_guarded = lambda a, b: a is not None and b is not None and a > b
# noqa: E731
+
+ # The same two-parameter guards written as an `if` TEST rather than as
the
+ # expression itself. This is the only shape that asks `_null_facts`
about a
+ # composite `and` / `or` node instead of about a single comparison,
and it
+ # is the idiomatic way to guard two columns.
+ def composite_and(a, b):
+ if a is not None and b is not None:
+ return a > b
+ else:
+ return None
+
+ def composite_or(a, b):
+ if a is None or b is None:
+ return None
+ else:
+ return a > b
+
+ def composite_not_or(a, b):
+ if not (a is None or b is None):
+ return a > b
+ else:
+ return None
+
+ schema = "a long, b long"
+ partial = self._transpiled_plan(one_guarded, BooleanType(), schema)
+ self.assertIn("raise_error", partial)
+ # Only `b` is still checked. Note `isnull(a` does not match the
guard's own
+ # `isnotnull(a...)`, which is the test rather than a check on the
operand.
+ self.assertIn("isnull(b", partial)
+ self.assertNotIn("isnull(a", partial)
+ rows = [(5, 1), (1, 5), (1, None), (None, 1)]
+ for func in (both_guarded, composite_and, composite_or,
composite_not_or):
+ with self.subTest(func=getattr(func, "__name__", "lambda")):
+ self.assertNotIn("raise_error", self._transpiled_plan(func,
BooleanType(), schema))
+ self.assertEqual(
+ [func(a, b) for a, b in rows],
+ self._vals(func, BooleanType(), schema, rows),
+ )
+ # Partial narrowing was only plan-asserted, but it creates two runtime
+ # invariants: the surviving check must still fire for a NULL `b`, and
the
+ # guard on `a` must still short-circuit a NULL `a` to False without
reaching
+ # the comparison. This is the one live consumer of the claim that
Catalyst's
+ # `And` short-circuits on false the way Python does.
+ self._raises(one_guarded, schema, [(1, None)], "cannot compare null",
BooleanType())
+ self.assertEqual([False], self._vals(one_guarded, BooleanType(),
schema, [(None, None)]))
+ self.assertEqual([False], self._vals(one_guarded, BooleanType(),
schema, [(None, 1)]))
+
+ # The UNSOUND direction, end to end: a true `or` says only that ONE
operand
+ # held, so it narrows nothing and the check must stay. Covered here as
well as
+ # in `test_null_facts_narrow_by_outcome`, because a plausible
"symmetry" edit
+ # to `_null_facts` would delete a needed raise, and the fact-table
unit test
+ # sits one remove from the plan that would go wrong.
+ def or_test(a, b):
+ if a is not None or b is not None:
+ return a > 0
+ else:
+ return None
+
+ def not_and_test(a, b):
+ if not (a is None and b is None):
+ return a > 0
+ else:
+ return None
+
+ for func in (or_test, not_and_test):
+ with self.subTest(func=func.__name__):
+ self.assertIn("raise_error", self._transpiled_plan(func,
BooleanType(), schema))
+ with self.assertRaises(Exception):
+ func(None, 1)
+ self._raises(func, schema, [(None, 1)], "cannot compare null",
BooleanType())
+
+ def test_udf_transpile_drops_the_coalesce_when_nothing_can_be_null(self):
+ # The `if` test and the `not` operand each get coalesced against a
literal
+ # so that a NULL follows Python's "None is falsy" rule. Neither
coalesce
+ # does anything once the operand provably cannot be NULL, and dropping
it
+ # leaves a plainer expression for the optimizer.
+ def non_null_test(x):
+ # `x is not None` lowers to isNotNull, which is never NULL itself.
+ if x is not None:
+ return 1
+ else:
+ return 2
+
+ def nullable_test(x):
+ # A ternary arm of literal None can be NULL, so the coalesce stays.
+ if (x > 0) if x is not None else None:
+ return 1
+ else:
+ return 2
+
+ # `not (x is None)` on purpose, rather than `x is not None`: it is the
`not`
+ # arm of _convert_chunk that this case is about.
+ negate_compare = lambda x: not (x is None) # noqa: E714,E731
+ # The `not` arm's negative case, and unlike the `if` test above this
+ # coalesce is load-bearing: Python's `not None` is True while
Catalyst's
+ # `~NULL` is NULL, so dropping it here would return NULL for a None
row.
+ negate_nullable = lambda x: not ((x > 0) if x is not None else None)
# noqa: E731
+
+ with self.sql_conf(_TRANSPILE_ON):
+ proven = str(self._transpiled_udf(non_null_test,
LongType()).transpiled[0])
+ self.assertNotIn("coalesce", proven)
+ self.assertIn("isNotNull", proven)
+ unproven = str(self._transpiled_udf(nullable_test,
LongType()).transpiled[0])
+ self.assertIn("coalesce", unproven)
+ negated = str(self._transpiled_udf(negate_compare,
BooleanType()).transpiled[0])
+ self.assertNotIn("coalesce", negated)
+ # Positive companion, so this cannot pass on a plan that no longer
+ # resembles the `not` lowering at all.
+ self.assertIn("!(isNull", negated)
+ still = str(self._transpiled_udf(negate_nullable,
BooleanType()).transpiled[0])
+ self.assertIn("coalesce", still)
+ # Values still follow Python, coalesce or not.
+ rows = [(5,), (-5,), (None,)]
+ for func, return_type in (
+ (non_null_test, LongType()),
+ (nullable_test, LongType()),
+ (negate_compare, BooleanType()),
+ (negate_nullable, BooleanType()),
+ ):
+ with self.subTest(func=getattr(func, "__name__", "lambda")):
+ self.assertEqual(
+ [func(v) for (v,) in rows],
+ self._vals(func, return_type, "a long", rows),
+ )
+
+ def test_udf_transpile_does_not_null_check_a_literal(self):
+ # A literal operand is never NULL, so only the column is checked. This
used
+ # to emit `isNull(0)` alongside it -- constant-folded later, but it
meant a
+ # comparison of two literals still produced a guard.
+ unguarded = lambda x: x > 0 # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ lowered = str(self._transpiled_udf(unguarded,
BooleanType()).transpiled[0])
+ self.assertIn("isNull(_udf_param_0)", lowered)
+ self.assertNotIn("isNull(0)", lowered)
+
+ def test_udf_transpile_warns_only_when_a_null_check_survives(self):
+ # The (gentle) warning SPARK-58628 asks for: the UDF transpiled fine,
but
+ # the check it kept will stop filters moving, so say so once. A
guarded UDF
+ # must not warn -- the warning is a nudge toward guarding, and firing
it on
+ # already-guarded code would train people to ignore it.
+ def guarded(x):
+ if x is not None:
+ return x > 0
+ else:
+ return None
+
+ unguarded = lambda x: x > 0 # noqa: E731
+ half_guarded = lambda a, b: a is not None and a > b # noqa: E731
+ needle = "still checks for NULL"
+ with self.sql_conf(_TRANSPILE_ON):
+ _, warned = self._udf_and_warnings(unguarded, BooleanType())
+ self.assertIn(needle, warned)
+ # Actionable: it should name what is checked, and the guard to add.
+ self.assertIn("comparison `>` on x", warned)
+ self.assertIn("is not None", warned)
+ # It must name the parameter still checked, not both -- `a` is
guarded
+ # here, so advising the reader to guard "the parameter" is only
useful
+ # if they can tell which one is left.
+ _, partial = self._udf_and_warnings(half_guarded, BooleanType())
+ self.assertIn("comparison `>` on b", partial)
+ # The full label, not " on a": that would match any prose word
starting
+ # with "a" elsewhere in the message.
+ self.assertNotIn("comparison `>` on a", partial)
+ # `_transpiled_udf`, not `_udf_and_warnings`: udf.py reports the
fallback
+ # and this warning through mutually exclusive branches, so a
`guarded`
+ # that stopped transpiling would satisfy assertNotIn for the wrong
+ # reason. Assert it transpiled, then that it stayed quiet.
+ quiet_udf, quiet = self._udf_and_warnings(guarded, BooleanType())
+ self.assertTrue(quiet_udf.transpiled, f"guarded must transpile:
{quiet}")
+ self.assertNotIn(needle, quiet)
+ # ``len`` raises on NULL the same way ``>`` does, so it warns the
same
+ # way -- and only for the check it actually kept. ``len(x) > 0``
must
+ # not also warn about the comparison: the ``length`` result is not
NULL.
+ unguarded_len = lambda x: len(x) # noqa: E731
+ len_cmp = lambda x: len(x) > 0 # noqa: E731
+ _, len_warned = self._udf_and_warnings(unguarded_len, LongType())
+ self.assertIn("`len` on x", len_warned)
+ # LongType, not BooleanType: this compare is not classified as
boolean
+ # (``len`` is a call), so a boolean return type refuses the
lowering.
+ cmp_udf, cmp_warned = self._udf_and_warnings(len_cmp, LongType())
+ self.assertTrue(cmp_udf.transpiled, f"len compare must transpile:
{cmp_warned}")
+ self.assertIn("`len` on x", cmp_warned)
+ self.assertNotIn("comparison `>`", cmp_warned)
+ self.assertFalse(UserDefinedFunction(len_cmp,
BooleanType()).transpiled)
+
+ def guarded_len(x):
+ if x is not None:
+ return len(x)
+ else:
+ return None
+
+ quiet_len, quiet_len_w = self._udf_and_warnings(guarded_len,
LongType())
+ self.assertTrue(quiet_len.transpiled, f"guarded len must
transpile: {quiet_len_w}")
+ self.assertNotIn(needle, quiet_len_w)
+
+ def test_udf_transpile_null_check_blocks_predicate_pushdown(self):
+ # Why the checks are worth dropping. RaiseError is throwable, and
+ # PushPredicateThroughJoin refuses to push a throwable condition to
either
+ # side (Optimizer.scala). So an unguarded UDF pins its filter above the
+ # join while a guarded one pushes below it -- the "filter cannot
bubble up"
+ # in SPARK-58628.
+ def guarded(x):
+ if x is not None:
+ return x > 0
+ else:
+ return None
+
+ unguarded = lambda x: x > 0 # noqa: E731
+
+ conf = dict(_TRANSPILE_ON)
+ # Broadcast would reshape the plan and hide where the Filter landed.
+ conf["spark.sql.autoBroadcastJoinThreshold"] = -1
+ with self.sql_conf(conf):
+ left = self.spark.createDataFrame([(1,), (2,), (None,)], "a long")
+ right = self.spark.createDataFrame([(1,), (2,)], "k long")
+ for func, pushed in ((guarded, True), (unguarded, False)):
+ with self.subTest(pushed=pushed):
+ u = self._transpiled_udf(func, BooleanType())
+ joined = left.join(right, left["a"] ==
right["k"]).filter(u("a"))
+ # This method builds its own plan rather than going through
Review Comment:
**Nit (P3):** `_optimized_plan` (line 1617) just returns
`optimizedPlan().toString()` and has no applied-or-not check. The helper that
asserts `_eval_python_count == 0` before reading the plan is the new
`_transpiled_plan`. As written, the comment suggests switching to
`_optimized_plan` would be safe, but that would drop exactly the protection the
comment is explaining. Pointing it at `_transpiled_plan` instead would fix it.
See **Shared repair plan 1** in the review body.
##########
python/pyspark/sql/udf.py:
##########
@@ -329,6 +338,24 @@ def _conf_is_true(key: str, default: Optional[str] = None)
-> bool:
self._transpiled_param_names = []
self._transpiled_input_categories = []
self._positional_only_param_names = frozenset()
+ if pending_guard_warning and self.transpiled:
Review Comment:
**Nit (P3):** This deferral matters. If the `warnings.warn` moved back into
the `try`, then under warnings-as-errors the generic `except Exception` handler
would clear `transpiled`, and its own `warnings.warn(...)` would then raise out
of the constructor. No test exercises this: the only `simplefilter("error")` in
the suite is the CHAR/VARCHAR case, and `_udf_and_warnings` uses `"always"`. A
small test would pin it: construct an unguarded `lambda x: x > 0` (BooleanType)
under `warnings.simplefilter("error")` and assert that it constructs and still
has `transpiled` options.
**Verification:**
- **Regression:** Under an error warning filter, an unguarded UDF that keeps
a NULL check constructs successfully with its transpiled options intact.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -2634,6 +2663,440 @@ def bytes_ident(x: bytes):
self.assertFalse(UserDefinedFunction(bool_add,
LongType()).transpiled)
self.assertTrue(UserDefinedFunction(bytes_ident,
BinaryType()).transpiled)
+ def _transpiled_plan(self, func, return_type, schema):
+ """The optimized plan of ``func`` applied to every column of
``schema``.
+
+ The lowered NULL checks are what these tests are about, so they have
to be
+ read off the plan: the values a guarded UDF returns are the same
whether or
+ not the plan carries a check that can never fire.
+
+ Asserting the option was actually APPLIED is what keeps every
+ ``assertNotIn(..., plan)`` built on this from passing for the wrong
reason
+ -- an interpreted plan contains none of the strings we look for, so a
UDF
+ that quietly fell back would satisfy them all (see ``_vals``, which
guards
+ the same way for the same reason).
+ """
+ with self.sql_conf(_TRANSPILE_ON):
+ u = self._transpiled_udf(func, return_type)
+ df = self.spark.createDataFrame([], schema)
+ projected = df.select(u(*df.columns))
+ self.assertEqual(0, self._eval_python_count(projected), str(func))
+ return projected._jdf.queryExecution().optimizedPlan().toString()
+
+ def test_udf_transpile_drops_null_checks_a_guard_already_made(self):
+ # SPARK-58628. Python raises TypeError on `None > 0` where Spark
returns
+ # NULL, so an ordering comparison carries a check that raises. When
the UDF
+ # has already tested the parameter the check can never fire, and
leaving it
+ # in is not free: RaiseError is throwable (SPARK-58627), which stops
the
+ # optimizer from moving a predicate containing it. Every shape below
proves
+ # non-NULL-ness, so none of them should lower to a raise -- and all of
them
+ # must still agree with Python, including on the None row.
+ def if_guard(x):
+ if x is not None:
+ return x > 0
+ else:
+ return None
+
+ def else_of_is_none(x):
+ if x is None:
+ return None
+ else:
+ return x > 0
+
+ def not_is_none(x):
+ if not (x is None): # noqa: E714 the `not` form is the point
+ return x > 0
+ else:
+ return None
+
+ def reversed_test(x):
+ if None is not x:
+ return x > 0
+ else:
+ return None
+
+ def ne_none_test(x):
+ # The `!= None` spelling of the same guard. PEP 8 prefers `is not
None`,
+ # but both are common and both must narrow, or the equality form
keeps a
+ # raising check and gets told to add the guard it already has.
+ if x != None: # noqa: E711
+ return x > 0
+ else:
+ return None
+
+ def eq_none_test(x):
+ if x == None: # noqa: E711
+ return None
+ else:
+ return x > 0
+
+ ternary = lambda x: (x > 0) if x is not None else None # noqa: E731
+ short_circuit_and = lambda x: x is not None and x > 0 # noqa: E731
+ short_circuit_or = lambda x: x is None or x > 0 # noqa: E731
+ # `<=` / `>=` share `_lower_value_compare` but are otherwise only
exercised
+ # where a guard IS emitted, so cover the narrowed side for them too.
+ lte_guarded = lambda x: (x <= 0) if x is not None else None # noqa:
E731
+ gte_guarded = lambda x: (x >= 0) if x is not None else None # noqa:
E731
+ lt_guarded = lambda x: (x < 0) if x is not None else None # noqa: E731
+
+ rows = [(-1,), (0,), (5,), (None,)]
+ cases = [
+ (if_guard, [False, False, True, None]),
+ (else_of_is_none, [False, False, True, None]),
+ (not_is_none, [False, False, True, None]),
+ (reversed_test, [False, False, True, None]),
+ (ne_none_test, [False, False, True, None]),
+ (eq_none_test, [False, False, True, None]),
+ (ternary, [False, False, True, None]),
+ (lte_guarded, [True, True, False, None]),
+ (gte_guarded, [False, True, True, None]),
+ (lt_guarded, [True, False, False, None]),
+ # `and` / `or` return the guard's own result for None, as Python
does.
+ (short_circuit_and, [False, False, True, False]),
+ (short_circuit_or, [False, False, True, True]),
+ ]
+ for func, expected in cases:
+ with self.subTest(func=getattr(func, "__name__", "lambda")):
+ self.assertEqual(
+ expected,
+ [func(v) for (v,) in rows],
+ "the expectation should be what Python actually does",
+ )
+ self.assertEqual(expected, self._vals(func, BooleanType(), "a
long", rows))
+ plan = self._transpiled_plan(func, BooleanType(), "a long")
+ self.assertNotIn("raise_error", plan)
+
+ def test_udf_transpile_keeps_the_null_check_it_needs(self):
+ # The other side of the previous test: nothing here proves the
parameter
+ # non-NULL, so the check has to stay and the UDF has to raise like
Python.
+ # `x is None and ...` is the interesting one -- an `and` being false
says
+ # only that SOME operand was false, so it narrows nothing.
+ unguarded = lambda x: x > 0 # noqa: E731
+ wrong_way_and = lambda x: x is None and x > 0 # noqa: E731
+ two_columns = lambda a, b: a > b # noqa: E731
+
+ for func, schema in ((unguarded, "a long"), (two_columns, "a long, b
long")):
+ with self.subTest(func="lambda", schema=schema):
+ plan = self._transpiled_plan(func, BooleanType(), schema)
+ self.assertIn("raise_error", plan)
+ boolean = BooleanType()
+ self._raises(unguarded, "a long", [(None,)], "cannot compare null",
boolean)
+ self._raises(two_columns, "a long, b long", [(1, None)], "cannot
compare null", boolean)
+ # Each ordering operator names itself in the guard's message and in the
+ # warning's label; only `>` was covered before.
+ lt_zero = lambda x: x < 0 # noqa: E731
+ lte_zero = lambda x: x <= 0 # noqa: E731
+ gte_zero = lambda x: x >= 0 # noqa: E731
+ for func, op_text in ((lt_zero, "`<`"), (lte_zero, "`<=`"), (gte_zero,
"`>=`")):
+ with self.subTest(op=op_text):
+ with self.sql_conf(_TRANSPILE_ON):
+ _, warned = self._udf_and_warnings(func, boolean)
+ self.assertIn(f"comparison {op_text} on x", warned)
+ self._raises(func, "a long", [(None,)], f"operator {op_text}",
boolean)
+ # `x is None and x > 0` is False for every non-None x without
evaluating the
+ # comparison, and raises for None -- exactly like Python.
+ with self.sql_conf(_TRANSPILE_ON):
+ u = self._transpiled_udf(wrong_way_and, BooleanType())
+ self.assertIn("raise_error", str(u.transpiled[0]))
+ df = self.spark.createDataFrame([(1,)], "a long")
+ self.assertEqual([False], [r[0] for r in
df.select(u("a")).collect()])
+ self._raises(wrong_way_and, "a long", [(None,)], "cannot compare
null", boolean)
+
+ def test_udf_transpile_narrows_only_the_proven_parameter(self):
+ # With two parameters the narrowing has to be per-name: guarding `a`
must
+ # not license dropping the check on `b`. Both guarded, no check at all.
+ one_guarded = lambda a, b: a is not None and a > b # noqa: E731
+ both_guarded = lambda a, b: a is not None and b is not None and a > b
# noqa: E731
+
+ # The same two-parameter guards written as an `if` TEST rather than as
the
+ # expression itself. This is the only shape that asks `_null_facts`
about a
+ # composite `and` / `or` node instead of about a single comparison,
and it
+ # is the idiomatic way to guard two columns.
+ def composite_and(a, b):
+ if a is not None and b is not None:
+ return a > b
+ else:
+ return None
+
+ def composite_or(a, b):
+ if a is None or b is None:
+ return None
+ else:
+ return a > b
+
+ def composite_not_or(a, b):
+ if not (a is None or b is None):
+ return a > b
+ else:
+ return None
+
+ schema = "a long, b long"
+ partial = self._transpiled_plan(one_guarded, BooleanType(), schema)
+ self.assertIn("raise_error", partial)
+ # Only `b` is still checked. Note `isnull(a` does not match the
guard's own
+ # `isnotnull(a...)`, which is the test rather than a check on the
operand.
+ self.assertIn("isnull(b", partial)
+ self.assertNotIn("isnull(a", partial)
+ rows = [(5, 1), (1, 5), (1, None), (None, 1)]
+ for func in (both_guarded, composite_and, composite_or,
composite_not_or):
+ with self.subTest(func=getattr(func, "__name__", "lambda")):
+ self.assertNotIn("raise_error", self._transpiled_plan(func,
BooleanType(), schema))
+ self.assertEqual(
+ [func(a, b) for a, b in rows],
+ self._vals(func, BooleanType(), schema, rows),
+ )
+ # Partial narrowing was only plan-asserted, but it creates two runtime
+ # invariants: the surviving check must still fire for a NULL `b`, and
the
+ # guard on `a` must still short-circuit a NULL `a` to False without
reaching
+ # the comparison. This is the one live consumer of the claim that
Catalyst's
+ # `And` short-circuits on false the way Python does.
+ self._raises(one_guarded, schema, [(1, None)], "cannot compare null",
BooleanType())
+ self.assertEqual([False], self._vals(one_guarded, BooleanType(),
schema, [(None, None)]))
+ self.assertEqual([False], self._vals(one_guarded, BooleanType(),
schema, [(None, 1)]))
+
+ # The UNSOUND direction, end to end: a true `or` says only that ONE
operand
+ # held, so it narrows nothing and the check must stay. Covered here as
well as
+ # in `test_null_facts_narrow_by_outcome`, because a plausible
"symmetry" edit
+ # to `_null_facts` would delete a needed raise, and the fact-table
unit test
+ # sits one remove from the plan that would go wrong.
+ def or_test(a, b):
+ if a is not None or b is not None:
+ return a > 0
+ else:
+ return None
+
+ def not_and_test(a, b):
+ if not (a is None and b is None):
+ return a > 0
+ else:
+ return None
+
+ for func in (or_test, not_and_test):
+ with self.subTest(func=func.__name__):
+ self.assertIn("raise_error", self._transpiled_plan(func,
BooleanType(), schema))
+ with self.assertRaises(Exception):
+ func(None, 1)
+ self._raises(func, schema, [(None, 1)], "cannot compare null",
BooleanType())
+
+ def test_udf_transpile_drops_the_coalesce_when_nothing_can_be_null(self):
+ # The `if` test and the `not` operand each get coalesced against a
literal
+ # so that a NULL follows Python's "None is falsy" rule. Neither
coalesce
+ # does anything once the operand provably cannot be NULL, and dropping
it
+ # leaves a plainer expression for the optimizer.
+ def non_null_test(x):
+ # `x is not None` lowers to isNotNull, which is never NULL itself.
+ if x is not None:
+ return 1
+ else:
+ return 2
+
+ def nullable_test(x):
+ # A ternary arm of literal None can be NULL, so the coalesce stays.
+ if (x > 0) if x is not None else None:
+ return 1
+ else:
+ return 2
+
+ # `not (x is None)` on purpose, rather than `x is not None`: it is the
`not`
+ # arm of _convert_chunk that this case is about.
+ negate_compare = lambda x: not (x is None) # noqa: E714,E731
+ # The `not` arm's negative case, and unlike the `if` test above this
+ # coalesce is load-bearing: Python's `not None` is True while
Catalyst's
+ # `~NULL` is NULL, so dropping it here would return NULL for a None
row.
+ negate_nullable = lambda x: not ((x > 0) if x is not None else None)
# noqa: E731
+
+ with self.sql_conf(_TRANSPILE_ON):
+ proven = str(self._transpiled_udf(non_null_test,
LongType()).transpiled[0])
+ self.assertNotIn("coalesce", proven)
+ self.assertIn("isNotNull", proven)
+ unproven = str(self._transpiled_udf(nullable_test,
LongType()).transpiled[0])
+ self.assertIn("coalesce", unproven)
+ negated = str(self._transpiled_udf(negate_compare,
BooleanType()).transpiled[0])
+ self.assertNotIn("coalesce", negated)
+ # Positive companion, so this cannot pass on a plan that no longer
+ # resembles the `not` lowering at all.
+ self.assertIn("!(isNull", negated)
+ still = str(self._transpiled_udf(negate_nullable,
BooleanType()).transpiled[0])
+ self.assertIn("coalesce", still)
+ # Values still follow Python, coalesce or not.
+ rows = [(5,), (-5,), (None,)]
+ for func, return_type in (
+ (non_null_test, LongType()),
+ (nullable_test, LongType()),
+ (negate_compare, BooleanType()),
+ (negate_nullable, BooleanType()),
+ ):
+ with self.subTest(func=getattr(func, "__name__", "lambda")):
+ self.assertEqual(
+ [func(v) for (v,) in rows],
+ self._vals(func, return_type, "a long", rows),
+ )
+
+ def test_udf_transpile_does_not_null_check_a_literal(self):
+ # A literal operand is never NULL, so only the column is checked. This
used
+ # to emit `isNull(0)` alongside it -- constant-folded later, but it
meant a
+ # comparison of two literals still produced a guard.
+ unguarded = lambda x: x > 0 # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ lowered = str(self._transpiled_udf(unguarded,
BooleanType()).transpiled[0])
+ self.assertIn("isNull(_udf_param_0)", lowered)
+ self.assertNotIn("isNull(0)", lowered)
+
+ def test_udf_transpile_warns_only_when_a_null_check_survives(self):
+ # The (gentle) warning SPARK-58628 asks for: the UDF transpiled fine,
but
+ # the check it kept will stop filters moving, so say so once. A
guarded UDF
+ # must not warn -- the warning is a nudge toward guarding, and firing
it on
+ # already-guarded code would train people to ignore it.
+ def guarded(x):
+ if x is not None:
+ return x > 0
+ else:
+ return None
+
+ unguarded = lambda x: x > 0 # noqa: E731
+ half_guarded = lambda a, b: a is not None and a > b # noqa: E731
+ needle = "still checks for NULL"
+ with self.sql_conf(_TRANSPILE_ON):
+ _, warned = self._udf_and_warnings(unguarded, BooleanType())
+ self.assertIn(needle, warned)
+ # Actionable: it should name what is checked, and the guard to add.
+ self.assertIn("comparison `>` on x", warned)
+ self.assertIn("is not None", warned)
+ # It must name the parameter still checked, not both -- `a` is
guarded
+ # here, so advising the reader to guard "the parameter" is only
useful
+ # if they can tell which one is left.
+ _, partial = self._udf_and_warnings(half_guarded, BooleanType())
+ self.assertIn("comparison `>` on b", partial)
+ # The full label, not " on a": that would match any prose word
starting
+ # with "a" elsewhere in the message.
+ self.assertNotIn("comparison `>` on a", partial)
+ # `_transpiled_udf`, not `_udf_and_warnings`: udf.py reports the
fallback
Review Comment:
**Nit (P3):** This reads "`_transpiled_udf`, not `_udf_and_warnings`", but
the next line deliberately calls `_udf_and_warnings` and asserts `transpiled`
explicitly. That is the right choice, because the check needs the captured
warning text and `_transpiled_udf` drops it. Only the helper names are
reversed. Something like "`_udf_and_warnings` plus an explicit `transpiled`
assertion, not `_transpiled_udf` (which discards the warnings)" would match the
code.
See **Shared repair plan 1** in the review body.
--
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]
