spark-test commented on code in PR #59978:
URL: https://github.com/apache/spark/pull/59978#discussion_r4466812337
##########
python/pyspark/sql/transpile.py:
##########
@@ -203,6 +203,9 @@ def _is_definitely_basic_type(node: ast.AST) -> bool:
return _is_definitely_basic_type(operand)
case ast.Name():
return True
+ case ast.Call(func=ast.Name(id="len"), args=[_], keywords=[]):
Review Comment:
**Blocking (P1):** Making `len()` definitely basic lets `and`/`or` chains
over `len()` comparisons lower, which the base branch refused. Each later
operand keeps its `len()` NULL guard (`CASE WHEN isnull(s) THEN
raise_error(...) ELSE length(s) END`), and the lowering relies on Catalyst's
And/Or to short-circuit past it. Catalyst gives no such guarantee. With the
default `spark.sql.subexpressionElimination.skipForShortcutExpr=false`,
subexpression elimination treats every And/Or operand as always evaluated. It
finds `isnull(<len guard>)` in both comparison guards and evaluates it for
every row before the projection runs. On Spark 4.3.0, this branch's option for
`def f(s): return s is not None and len(s) > 3 and len(s) < 10` failed the
query over a string column with a NULL row, raising `[USER_RAISED_EXCEPTION]
Python UDF transpiler: cannot call len() on NULL`, where Python returns False.
A single comparison, or the same chain with subexpression elimination disabled,
returns Python'
s values. The same failure occurs for `x > 0 and len(s) > 3 and len(s) < 10`,
for `s is None or len(s) < 3 or len(s) > 10`, for the chain used as an `if`
test, and for two calls such as `s is not None and len(s) > 3` and `s is not
None and len(s) < 10` in one select.
**Recommended change:** In the `and`/`or` lowering of `_convert_chunk`,
refuse the chain, so the variant falls back to interpreted Python, when an
operand after the first still emits a `len()` NULL guard after narrowing. Count
the `len()` guards the lowering emits and compare the count before and after
converting each later operand. Add a UDF transpiler test that runs `s is not
None and len(s) > 3 and len(s) < 10` over a string column holding NULL and
asserts interpreted Python's results.
**Why this works:** Python evaluates only the first operand of an `and`/`or`
unconditionally. Keeping raising `len()` guards out of the later operands
leaves nothing that Catalyst can evaluate ahead of the short-circuit, whatever
subexpression elimination shares within a chain or across calls.
**Scope:** Keep `len()` NULL guards out of `and`/`or` operands after the
first.
**Compatibility:** Chains whose later operands carry no `len()` guard lower
as before, and `len()` comparisons outside `and`/`or` keep their guard.
**Risks:** Chains whose later operands keep a `len()` guard, such as `x > 0
and len(s) > 3`, run as interpreted Python, as on the base branch. Narrowing
later operands by earlier None checks is what keeps `s is not None and ...`
chains lowered.
**Constraints:** Count only `len()` guards that remain after narrowing, so a
chain whose earlier None checks prove the parameter non-null still lowers. Run
the regression test where subexpression elimination applies: a
whole-stage-codegen projection with default settings.
**Exclusions:**
- Guards that the base branch already places in `and`/`or` operands, on
ordering comparisons over parameters, are left as they are.
**Success:** `s is not None and len(s) > 3 and len(s) < 10`, `x > 0 and
len(s) > 3 and len(s) < 10`, and `s is None or len(s) < 3 or len(s) > 10`
return interpreted Python's results on NULL rows under default subexpression
elimination, lowered where narrowing removes every later `len()` guard and
interpreted otherwise. No lowered `and`/`or` operand after the first contains a
`len()` NULL guard.
**Verification:**
- **Regression:** Over a string column holding NULL, with transpilation and
ANSI on and default subexpression elimination in a whole-stage-codegen
projection, the three chains above, and the two calls `s is not None and len(s)
> 3` and `s is not None and len(s) < 10` selected together, return interpreted
Python's values.
- **Behavior:** Constructing `x > 0 and len(s) > 3` produces no transpiled
option.
##########
python/pyspark/sql/transpile.py:
##########
@@ -256,6 +259,46 @@ def _truthiness_col(cat: Optional[str], c: Column) ->
Optional[Column]:
return None
+def _null_narrowed(params: List[str], test_node: ast.AST) -> Tuple[frozenset,
frozenset]:
+ """Return ``(body_nonnull, else_nonnull)`` — frozensets of parameter names
+ that are proven non-null in the true branch and the false branch
+ respectively of an ``if``/ternary with this test.
+
+ Recognised patterns (``p`` must be in ``params``):
+ ``p is not None`` → body: {p}
+ ``p is None`` → else: {p}
+ ``None is not p`` → body: {p}
+ ``None is p`` → else: {p}
+ ``A and B and …`` → body: union of each clause's body set
+ All other nodes yield empty sets (conservative; no narrowing assumed).
+ """
+ match test_node:
+ case (
+ ast.Compare(
+ left=ast.Name(id=name),
+ ops=[op],
+ comparators=[ast.Constant(value=None)],
+ )
+ | ast.Compare(
+ left=ast.Constant(value=None),
+ ops=[op],
+ comparators=[ast.Name(id=name)],
+ )
+ ) if name in params:
+ if isinstance(op, ast.IsNot):
+ return frozenset({name}), frozenset()
+ if isinstance(op, ast.Is):
+ return frozenset(), frozenset({name})
+ case ast.BoolOp(op=ast.And(), values=values):
Review Comment:
**Non-blocking (P2):** Only `and` is handled here, and `_null_narrowed`
returns nothing for an `or`. So `if a is None or b is None: return None` /
`return a < b`, which this change newly lowers, still guards `a < b` with `CASE
WHEN isnull(a) OR isnull(b) THEN raise_error`, and emits the NULL-guard warning
once per lowered variant (five for an untyped two-parameter UDF). The
`and`/`or` lowering converts every operand with the outer non-null set, so `x
is not None and x > 0` keeps the guard on `x > 0`. `(x + 1 > 0) if x is not
None else None` keeps it too, because arithmetic over a narrowed parameter is
not treated as non-null. Each of these already excludes NULL, so it still has
the problem this change fixes for the plain ternary.
**Recommended change:** Extend null narrowing in three ways. First, have
`_null_narrowed` also return false-branch narrowing for an `or`-chain of `is
None` checks (their union). Second, when lowering an `and`, convert each later
operand with the parameters that earlier `is not None` operands prove non-null;
for an `or`, use those that earlier `is None` operands prove. Restore the outer
set afterwards. Third, have `_is_nonnull_node` treat `+`, `-`, `*`, `%`, and
unary minus over non-null operands as non-null. Add UDF transpiler tests for
each idiom asserting no raise_error, no advisory, Python-matching results, and
a lowered plan.
**Why this works:** Each extension proves non-null only where Python itself
never evaluates the operand on None. Under SQL three-valued logic,
`isnotnull(x) AND x > 0` and `isnull(x) OR x > 0` return Python's results
whatever the evaluation order. Arithmetic over non-null operands either returns
a non-null value or raises under ANSI.
**Scope:** Extend null narrowing to `or`-chains of `is None` checks, to
short-circuit operands, and to arithmetic over narrowed parameters.
**Compatibility:** Narrowing for the if/ternary shapes this change already
recognizes is unchanged.
**Risks:** Plans change for these idioms. Their operands no longer contain
raise_error, so filter splitting may place them more freely; that is correct,
because they can no longer raise.
**Constraints:** Narrow an operand only by operands that precede it in the
same `and`/`or`, and restore the outer set afterwards, as the if/ternary arms
do. Keep the docstring and comments of `_null_narrowed` ASCII.
**Exclusions:**
- Short-circuit protection that comes from a condition other than a None
check, such as `x > 0 and len(s) > 3`, is not narrowed; the fix that keeps
`len()` guards out of later `and`/`or` operands covers that case.
**Success:** `return x is not None and x > 0`, `if a is None or b is None:
return None` / `return a < b`, their ternary twins, `(x + 1 > 0) if x is not
None else None`, and `s is not None and len(s) > 3 and len(s) < 10` all lower
with no raise_error and emit no NULL-guard advisory. Comparisons whose operands
are not proven non-null keep the guard.
**Verification:**
- **Behavior:** For each idiom, the optimized plan has no raise_error, the
JVM keeps the lowered option, no advisory is recorded, and the results match
interpreted Python on NULL and non-NULL rows.
- **Regression:** Cases that prove nothing keep the guard and still raise on
NULL: tests that `or` together `is not None` checks, the true branch of an `or`
of `is None` checks, and operands that come before the guarding check.
##########
python/pyspark/sql/transpile.py:
##########
@@ -938,7 +1017,20 @@ def _transpile_from_ast(
# Category inference depends on the per-variant assumptions above.
Cache
# each AST node only for this lowering so recursive conversion stays
linear.
self._category_cache = {}
+ # Reset per-transpilation null-narrowing state.
+ self._nonnull_params = frozenset()
+ self._null_guard_emitted = False
function_body = function_ast.body
+ # Normalize `if <test>: return y\nreturn body` (2-stmt early-return)
+ # into a single if/else so null narrowing can apply to the main body.
+ match function_body:
+ case [
+ ast.If(test=test, body=[ast.Return() as early_ret], orelse=[])
as if_node,
Review Comment:
**Blocking (P1):** This pattern constrains only the statement shape, so any
`if <test>: return a` / `return b` body is now lowered, not just the null-guard
idiom it exists for. A test that is not a None check gets no narrowing benefit,
yet the body stops running as interpreted Python, as it did on the base branch,
and picks up the lowered semantics instead. For example, `def f(x): if x > 0:
return 1` / `return 0` (LongType) over a DOUBLE column keeps the float option,
where Spark orders NaN above every value: on Spark 4.3.0, this branch's float
option returned `[1, 1, 0]` for `[NaN, 2.0, -3.0]`, where Python returns `[0,
1, 0]`. Likewise, `if x > 0: return x + y` / `return 0` with `y` NULL returns
NULL where Python raises TypeError. Results therefore change silently for
inputs that were handled correctly before this change.
**Recommended change:** Apply the early-return rewrite only when
`_null_narrowed` proves some parameter non-null in either branch of the test.
Other two-statement bodies keep failing the one-statement check and run as
interpreted Python, as on the base branch. Add UDF transpiler regression tests
for a non-None-test early return over a double column holding NaN, and for
unguarded arithmetic with a NULL operand.
**Why this works:** Gating the rewrite on the narrowing result confines it
to the null-guard idiom, the only input that needs it. Bodies whose test proves
nothing non-null keep the base branch's interpreted results.
**Scope:** Gate the two-statement early-return rewrite on null narrowing.
**Compatibility:** Null-guard early returns, including those narrowed
through `and`-chains, keep lowering exactly as they do now.
**Risks:** Early returns whose test is not a None check run as interpreted
Python again, as on the base branch. That gives up the lowering this change
added for them; it can be relaxed later once the documented differences are
handled.
**Constraints:** Use `_null_narrowed`'s result itself as the gate, so that
any extension of narrowing (for example to `or`-chains of `is None` checks)
automatically extends which early returns are rewritten.
**Exclusions:**
- The NaN-ordering and NULL-arithmetic differences that the null-guard idiom
itself inherits are left to the author's decision on that idiom.
**Success:** `if x > 0: return 1` / `return 0` produces no transpiled
options and returns Python's result on a double NaN row. `if x is None: return
None` / `return x > 0` still lowers with no raise_error in its main body.
**Verification:**
- **Regression:** A two-statement early return whose test is not a None
check falls back to interpreted Python and returns Python's results for a NaN
input and for a NULL operand of unguarded arithmetic.
- **Behavior:** Null-guard early returns keep lowering without a guard, and
the JVM keeps the lowered option.
##########
python/pyspark/sql/transpile.py:
##########
@@ -991,6 +1083,15 @@ def _transpile_from_ast(
"interpreted Python"
)
converted = self._convert_chunk(params, function_body[0])
+ if self._null_guard_emitted:
+ warnings.warn(
Review Comment:
**Non-blocking (P2):** This `warnings.warn` runs inside `_transpile_func`'s
per-variant `try ... except Exception`, after the variant's Column is built,
and `UserWarning` is an `Exception`. Under a filter that escalates it, the
built option is recorded as an error and dropped. On Spark 4.3.0 with this
branch's transpile.py, constructing `UserDefinedFunction(unguarded_gt,
BooleanType())` for `def unguarded_gt(x): return x > 0` under
`warnings.simplefilter('error')` raised `Exception transpiling UDF ...: Unable
to transpile UDF ...: ['Python UDF transpiler: the transpiled UDF includes a
NULL guard ...']`; the base branch's file constructs it with both options.
Under a filter that escalates only this advisory, construction succeeds with
zero options and only a fallback warning, so the UDF runs as interpreted
Python. A blanket error filter would surface a correctly placed advisory too.
What this placement adds is throwing away a successful lowering and reporting
it as a transpilation
failure.
See **Shared repair plan 1** in the review body.
##########
python/pyspark/sql/transpile.py:
##########
@@ -501,12 +564,15 @@ def _lower_value_compare(
)
left_col = self._convert_chunk(params, left_node)
right_col = self._convert_chunk(params, right_node)
+ if self._is_nonnull_node(params, left_node) and
self._is_nonnull_node(params, right_node):
Review Comment:
**Nit (P3):** This fast path now skips the guard when both operands are
proven non-null. The `_lower_value_compare` docstring above (lines 539-544)
still says such callers 'take the otherwise branch' of a runtime guard, and the
`len()` comment at 978-980 says the same next to the matching fast path at
987-988. A reader would miss that the guard is now omitted statically, and that
it stays unless every operand is proven non-null: `x > y` with only `x`
narrowed keeps it. Both descriptions should state when the guard is omitted,
for example by referring to `_is_nonnull_node`.
**Verification:**
- **Inspection:** Each description states that the guard is omitted when
every operand is proven non-null, and neither still claims that proven callers
take a runtime otherwise branch.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -2750,6 +2750,148 @@ def
test_udf_transpile_int_float_annotation_categories(self):
self.assertEqual(len(combos_u), 3)
self.assertEqual({c[0] for c in combos_u}, {"integer", "float",
"string"})
+ def test_udf_transpile_null_narrowing_skips_null_guard(self):
+ # When a parameter is guarded by `x is not None` in an
+ # if/ternary test, the transpiler should skip the raise_error null
guard
+ # for ordering comparisons on that parameter within the guarded branch.
+ def if_guard_gt(x):
+ return (x > 0) if x is not None else None
+
+ def if_guard_and_lt(a, b):
+ return (a < b) if a is not None and b is not None else None
+
+ def ternary_guard_ge(x):
+ return (x >= 0) if x is not None else None
+
+ def is_none_ternary(x):
+ return None if x is None else x > 0
+
+ cases_if = [
+ # if/else with `is not None` guard
+ ("if_guard_gt", if_guard_gt),
+ # if/else with `is not None` and-chain (two params)
+ ("if_guard_and_lt", if_guard_and_lt),
+ # ternary with `is not None` guard
+ ("ternary_guard_ge", ternary_guard_ge),
+ # `is None` guard: non-null in the else branch
+ ("is_none_ternary", is_none_ternary),
+ ]
+ with self.sql_conf(_TRANSPILE_ON):
+ for name, func in cases_if:
+ with self.subTest(func=name):
+ u, warnings_text = self._udf_and_warnings(func,
BooleanType())
+ self.assertTrue(u.transpiled, f"{name} did not transpile:
{warnings_text}")
+ arity = func.__code__.co_argcount
+ if arity == 1:
+ df = self.spark.createDataFrame([(1,), (-1,),
(None,)], "a long")
+ plan = self._optimized_plan(df.select(u("a")))
+ self.assertNotIn("raise_error", plan, f"{name}: plan
had raise_error")
Review Comment:
**Non-blocking (P2):** If the JVM discarded these options, the optimized
plan would hold the PythonUDF, which has no `raise_error`, and interpreted
Python would return the same values, so every assertion here would still pass.
`_vals` documents exactly this trap, and the file's signal that the lowered
option was kept, `_eval_python_count(...) == 0`, is not used here. Every case
is also positive, so narrowing that leaked into the wrong branch would go
unnoticed: `(x > 0) if x is None else None` would return NULL where Python
raises TypeError. The captured warning text is never asserted on either, so a
warning that fires for every narrowed UDF would pass too.
See **Shared repair plan 2** in the review body.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -2750,6 +2750,148 @@ def
test_udf_transpile_int_float_annotation_categories(self):
self.assertEqual(len(combos_u), 3)
self.assertEqual({c[0] for c in combos_u}, {"integer", "float",
"string"})
+ def test_udf_transpile_null_narrowing_skips_null_guard(self):
Review Comment:
**Nit (P3):** Several boundaries the new logic introduces have no test.
Nothing exercises the reversed `None is not p` / `None is p` forms or narrowing
in an elif test. The early-return rewrite is safe only because it requires
`orelse=[]`; without that, `if x is None: return 1` / `elif x > 0: return 2` /
`return 3` would drop the elif and return 3 where Python returns 2, and no test
pins that this body is still refused. The `len()` reclassification changes
three outcomes, none of them tested: `return len(x) > 3` declared LongType now
falls back, which matches the interpreted NULL; BooleanType `len(x) > 3` and
`len(x) == 3` now lower; and `if len(x) > 3:` now lowers as a boolean test.
See **Shared repair plan 2** in the review body.
##########
python/pyspark/sql/transpile.py:
##########
@@ -203,6 +203,9 @@ def _is_definitely_basic_type(node: ast.AST) -> bool:
return _is_definitely_basic_type(operand)
case ast.Name():
return True
+ case ast.Call(func=ast.Name(id="len"), args=[_], keywords=[]):
Review Comment:
**Blocking (P1):** With this arm, `len(s) > 10` is definitely boolean and
`_category` labels it 'bool'. Both Mult arms accept only an 'integer' or
'numeric' repeat count, so `def mark(s): return s + '*' * (len(s) > 10)` now
raises at the Mult arm in every input-type variant and falls back to
interpreted Python. On the base branch, the comparison fell into the 'numeric'
catch-all, and the string variant lowered to `concat(s, repeat('*',
CAST(<guarded comparison> AS INT)))`, an ANSI-legal cast that matches Python's
`str * bool`. Running `udf()` on Spark 4.3.0 with each revision's transpile.py
confirms this: with the base branch's file the plan has no BatchEvalPython and
returns Python's values, and with this branch's file it plans BatchEvalPython.
Results are unchanged, but every row now makes a pickled round trip through a
Python worker. Arithmetic over the comparison, such as `(len(s) > 10) + 1`, is
not affected, because its old option failed ANSI analysis.
**Recommended change:** In both string-repeat arms, accept a count whose
node is a comparison the transpiler classifies as boolean, and lower it as
`repeat(str, CAST(count AS INT))`, as the base branch did for `len()`
comparisons. Keep refusing other boolean counts, such as a ternary that can
yield None. Add a UDF transpiler test showing `s + '*' * (len(s) > 10)` lowers
and matches Python.
**Why this works:** A lowered comparison is never NULL: ordering comparisons
are guarded or narrowed, equality uses the four-branch None handling, and `is`
checks become isNull or isNotNull. So `CAST(comparison AS INT)` is a 0 or 1
count equal to Python's `str * bool`, and ANSI allows that cast.
**Scope:** Accept comparison-valued string-repeat counts again.
**Compatibility:** Numeric `*`, integer repeat counts, and the boolean
classification of `len()` comparisons are unchanged.
**Risks:** Repeat counts given by other comparisons, such as `'s' * (n !=
1)`, which the base branch refused, now lower natively too. Their results match
Python, and removing the comparison case rolls this back.
**Constraints:** Limit the new case to comparison nodes and keep refusing
other boolean categories.
**Success:** `s + '*' * (len(s) > 10)` declared StringType lowers to a
native projection again and returns Python's values.
**Verification:**
- **Regression:** A UDF that uses a `len()` comparison as a string-repeat
count produces a lowered option, the JVM keeps it, and the results match
interpreted Python.
- **Behavior:** A string-repeat count given by a ternary with a None branch
still falls back.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -2750,6 +2750,148 @@ def
test_udf_transpile_int_float_annotation_categories(self):
self.assertEqual(len(combos_u), 3)
self.assertEqual({c[0] for c in combos_u}, {"integer", "float",
"string"})
+ def test_udf_transpile_null_narrowing_skips_null_guard(self):
+ # When a parameter is guarded by `x is not None` in an
+ # if/ternary test, the transpiler should skip the raise_error null
guard
+ # for ordering comparisons on that parameter within the guarded branch.
+ def if_guard_gt(x):
+ return (x > 0) if x is not None else None
+
+ def if_guard_and_lt(a, b):
+ return (a < b) if a is not None and b is not None else None
+
+ def ternary_guard_ge(x):
+ return (x >= 0) if x is not None else None
+
+ def is_none_ternary(x):
+ return None if x is None else x > 0
+
+ cases_if = [
+ # if/else with `is not None` guard
Review Comment:
**Nit (P3):** `if_guard_gt` and `if_guard_and_lt` are conditional
expressions (`(x > 0) if x is not None else None`), the same shape as the case
labelled 'ternary', so despite its name and comments this list never reaches
the `ast.If` narrowing arm. Separately, the assertion message in
`test_udf_transpile_null_guard_warning` says 'under strict or loose', but
transpile.py, udf.py, and the transpiler confs define no such modes.
See **Shared repair plan 2** in the review body.
##########
python/pyspark/sql/transpile.py:
##########
@@ -256,6 +259,46 @@ def _truthiness_col(cat: Optional[str], c: Column) ->
Optional[Column]:
return None
+def _null_narrowed(params: List[str], test_node: ast.AST) -> Tuple[frozenset,
frozenset]:
+ """Return ``(body_nonnull, else_nonnull)`` — frozensets of parameter names
Review Comment:
**Nit (P3):** This docstring and the inline comment at line 293 add an em
dash (U+2014), right arrows (U+2192, lines 268-272 and 293), and ellipses
(U+2026). They are the only non-ASCII characters in transpile.py. Repository
guidance asks for ASCII in code and comments, and the rest of the file uses
`--`, `->`, and `...`, for example in `_truthiness_col`'s table.
**Verification:**
- **Inspection:** A scan for non-ASCII characters finds none in the
transpiler module.
##########
python/pyspark/sql/transpile.py:
##########
@@ -991,6 +1083,15 @@ def _transpile_from_ast(
"interpreted Python"
)
converted = self._convert_chunk(params, function_body[0])
+ if self._null_guard_emitted:
+ warnings.warn(
+ "Python UDF transpiler: the transpiled UDF includes a NULL
guard "
Review Comment:
**Nit (P3):** The message never says which UDF, parameter, or operator needs
the guard. Because `stacklevel=2` resolves to `_transpile_func`'s call at line
1566, Python attributes every copy to transpile.py rather than to user code.
One copy is issued for each variant that kept a guard: two for an untyped `x >
0`, five for an untyped two-parameter comparison. A UDF that only calls
`len(s)` is told to guard comparisons it does not contain. The warning is also
issued before the options are final, so `udf(f).asNondeterministic()` warns
about a lowering that never runs. udf.py's own construction warnings all name
the UDF (`{func}`).
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]
