spark-test commented on code in PR #59898:
URL: https://github.com/apache/spark/pull/59898#discussion_r4408012971
##########
python/pyspark/sql/transpile.py:
##########
@@ -372,31 +405,56 @@ def _lower_eq(
where Python's ``==`` is simply False. Refuse those so the UDF falls
back to interpreted Python. A ``None`` literal operand stays allowed
(the four-branch NULL handling above reproduces Python exactly).
+ Numeric sub-categories (``"integer"`` vs ``"float"``) are treated as
+ compatible; Spark widens to double in the comparison.
+
+ A NaN guard is added unless both operands are provably ``"integer"``:
+ ``NaN == NaN`` is ``False`` in Python (IEEE 754), but Spark's
+ ``EqualTo`` returns ``True``. ``"numeric"`` is an unknown sub-type
+ that may resolve to a float column, so it keeps the guard.
- One value-level difference remains (needs runtime values, so it is
- documented, not guarded): Spark treats ``NaN = NaN`` as true, while
- Python's ``nan == nan`` is False.
+ One value-level difference remains for ordering comparisons (tracked
+ separately): Spark orders NaN as greater than every value, whereas
+ Python's NaN comparisons are all False.
"""
lc = self._safe_category(params, left_node)
rc = self._safe_category(params, right_node)
if lc is not None and rc is not None and lc != rc:
- raise UnsupportedOperationException(
- f"`==`/`!=` operands have incompatible categories ({lc} vs
{rc}); "
- "Python compares across types as unequal while Spark would
coerce "
- "or fail analysis, so the transpiler falls back to interpreted
Python"
- )
+ # Numeric sub-categories are mutually compatible for equality.
+ if not (_is_numeric_cat(lc) and _is_numeric_cat(rc)):
+ raise UnsupportedOperationException(
+ f"`==`/`!=` operands have incompatible categories ({lc} vs
{rc}); "
+ "Python compares across types as unequal while Spark would
coerce "
+ "or fail analysis, so the transpiler falls back to
interpreted Python"
+ )
left_col = self._convert_chunk(params, left_node)
right_col = self._convert_chunk(params, right_node)
left_null = left_col.isNull()
right_null = right_col.isNull()
+ # NaN guard: Python's `NaN == NaN` is False (IEEE 754 reflexivity
fails),
+ # but Spark's EqualTo returns True. Emit only for numeric operands, and
+ # skip only when both are provably integral ("integer"); "numeric" is
an
+ # unknown sub-type that may resolve to a float column, so it keeps the
guard.
+ if (
+ lc is not None
+ and rc is not None
+ and _is_numeric_cat(lc)
+ and _is_numeric_cat(rc)
+ and not (lc == "integer" and rc == "integer")
+ ):
+ nan_result = lit(not equal)
+ nan_cmp = isnan(left_col) | isnan(right_col)
Review Comment:
**Non-blocking (P2):** This guard trusts `_safe_category`, but `"numeric"`
is also `_category_uncached`'s catch-all (lines 635-639), and that function has
no `ast.Compare` arm. A comparison is labeled `"bool"` only when
`_is_definitely_boolean` accepts it, which requires basic operands, so `len(s)
> 0`, `(a if c else b) > 0` or a nested comparison is labeled `"numeric"` even
though `_lower_value_compare` lowers it to a BOOLEAN column.
For example, with transpilation and ANSI on, `def f(s: str, t: str): return
(len(s) > 0) == (len(t) > 0)` registered as `udf(f, "int")` over two string
columns passes the return-type check as numeric, keeps its `[string, string]`
option, and now emits `isnan(<boolean>)`. `IsNaN` accepts only DOUBLE/FLOAT
(`nullExpressions.scala` 394-397), and ANSI implicit casting finds no
BOOLEAN-to-DOUBLE cast (`Cast.canANSIStoreAssign` returns false). Because
options are children of `TranspiledPythonUDF`, CheckAnalysis rejects the whole
query with a DATATYPE_MISMATCH on an `isnan` the user never wrote, instead of
falling back to Python. On branch-4.x before this change the same option was
`EqualTo(boolean, boolean)` cast to int, which analyzes. The trigger is narrow
(compound comparisons on both sides of `==`/`!=` with a numeric declared return
type), but it is exactly the whole-query failure the fallback design exists to
prevent.
**Recommended change:** Give every comparison node the boolean category in
the transpiler's category inference, ahead of the numeric catch-all, so
compound comparisons are labeled as the booleans they lower to. Add Python UDF
transpilation tests for equality between compound comparisons.
**Why this works:** Every comparison the transpiler lowers produces a
BOOLEAN column: is/is not None, ==, !=, <, <=, > and >=. Once such nodes are
labeled boolean, the NaN guard, which fires only when both operands are labeled
numeric, no longer sees them. A boolean-versus-boolean equality lowers to a
plain equality, and the existing category checks treat these nodes as booleans
for arithmetic, truthiness, branch compatibility and the return-type match.
**Scope:** Category inference for comparison nodes in the Python UDF
transpiler, plus regression tests for equality between compound comparisons.
**Compatibility:** NaN equality semantics for floating-point operands, and
the existing lowering of basic comparisons, stay unchanged.
**Risks:** UDFs whose body is a compound comparison declared with a numeric
return type currently lower and return the boolean cast to a number. Once their
body is labeled boolean, the return-type check refuses them and they run as
interpreted Python, which returns NULL for a bool returned under a numeric type.
**Constraints:** Keep the NaN guard for genuinely numeric operands,
including negated floating-point operands that are still labeled numeric, so
NaN == NaN remains False in lowered equality.
**Exclusions:**
- Making unary plus and minus inherit their operand's category, and labeling
len() as integral, is left out. Both would also tighten the integral-only
repeat check and drop guards that can never fire, but neither is needed to fix
this analysis failure.
**Success:** A UDF such as `def f(s: str, t: str): return (len(s) > 0) ==
(len(t) > 0)` with a numeric return type no longer fails the query with an
isnan type mismatch; its body is labeled boolean and the call runs as
interpreted Python. The same body declared with a boolean return type lowers to
a plain boolean equality with no isnan. Lowered floating-point equality still
returns False for NaN == NaN and True for NaN != NaN.
**Verification:**
- **Regression:** With transpilation and ANSI mode enabled, an equality or
inequality between two comparisons over len() or ternary operands analyzes
successfully and returns the interpreted-Python results instead of failing the
query.
- **Behavior:** A boolean-returning equality between compound comparisons is
lowered, with no Python evaluation in the executed plan and no isnan, while
lowered floating-point equality over NaN rows still returns Python's results.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -1640,12 +1640,12 @@ def test_udf_transpile_evaluates_inputs_once(self):
# unresolved and ResolveRandomSeed then gives each spliced copy
its OWN seed. Both have
# to end up sharing one evaluation, or `x == x` compares two
independent draws.
#
- # Sharing leaves `col = col`, which SimplifyBinaryComparison folds
to true and column
- # pruning then drops the draw entirely -- so zero draws in the
plan is the proof that
- # both sides became the same column. Two independent draws would
leave both.
+ # The NaN guard for the untyped (numeric) parameter keeps one
`isnan` reference to the
+ # shared arg, so one draw remains in the plan after
SimplifyBinaryComparison folds
+ # `col == col` to true. Two independent draws would leave both.
for arg in (rand(), expr("rand()")):
eq_df = rows.select(eq_udf(arg).alias("v"))
- self.assertEqual(0, self._draw_count(eq_df))
+ self.assertEqual(1, self._draw_count(eq_df))
Review Comment:
**Nit (P3):** Before this change, `0` draws could only come from the lowered
path with both reads on one shared column. One draw is also what the optimized
plan shows when the JVM discards the option and runs the UDF as interpreted
Python, which also returns True on every row. Since the test only checks
`eq_udf.transpiled`, a regression that drops the option for these calls
(including the unresolved-seed `expr("rand()")` case only this test covers)
would pass. The sibling sharing tests pair their one-draw counts with
`_shares_an_argument`; adding that and `self.assertEqual(0,
self._eval_python_count(eq_df))` here restores the signal. Also, the guard that
keeps the draw comes from the float variant the JVM keeps for the DOUBLE
`rand()` argument, not from a "numeric" parameter as the new comment says.
**Verification:**
- **Regression:** The shared-draw check for `x == x` over rand() and
expr('rand()') fails when the transpiled option is discarded and the call runs
as interpreted Python.
##########
python/pyspark/sql/tests/test_udf_transpile_unit.py:
##########
@@ -2636,17 +2636,120 @@ def bytes_ident(x: bytes):
def test_param_category_combos_caps_preserve_typed_pins(self):
# With more than three untyped params the cap collapses the untyped
ones
- # to numeric/string but keeps each typed param pinned (here a: str).
+ # to integer/float/string but keeps each typed param pinned (here a:
str).
import ast as _ast
from pyspark.sql.transpile import _param_category_combos
fn = _ast.parse("def f(a: str, b, c, d, e): return a").body[0]
combos = _param_category_combos(fn, ["a", "b", "c", "d", "e"])
- self.assertEqual(len(combos), 2)
+ # Three fill values: integer / float / string -> 3 combos.
+ self.assertEqual(len(combos), 3)
for combo in combos:
self.assertEqual(combo[0], "string")
+ def test_udf_transpile_integer_eq_no_nan_guard(self):
+ # `a == b` on BIGINT columns must NOT emit a NaN guard. Integers can
never
+ # be NaN, so emitting `isnan(cast(bigint as double))` is unnecessary
dead
+ # code. Verify via the logical plan: neither "isnan" nor an implicit
double
+ # cast should appear for the integer equality variant.
+ from pyspark.sql.types import StructField, StructType
+
+ def x_eq_y(x: int, y: int) -> bool:
+ return x == y
+
+ schema = StructType(
+ [
+ StructField("a", LongType(), nullable=True),
+ StructField("b", LongType(), nullable=True),
+ ]
+ )
+ with self.sql_conf(_TRANSPILE_ON):
+ pudf = UserDefinedFunction(x_eq_y, BooleanType())
+ self.assertTrue(pudf.transpiled, "annotated int == int should
transpile")
+ df = self.spark.createDataFrame([(1, 1), (1, 2), (None, None)],
schema=schema)
+ projected = df.select(pudf("a", "b").alias("r"))
+ # Verify no Python eval node (plan was rewritten to Catalyst).
+ self.assertEqual(0, self._eval_python_count(projected))
+ # Verify isnan is NOT in the logical plan.
+ logical_plan =
projected._jdf.queryExecution().analyzed().toString()
+ self.assertNotIn("isnan", logical_plan.lower(), "integer == must
not emit NaN guard")
+ # Verify correct results: (1==1, 1==2, None==None in Python ==
True).
+ results = [r[0] for r in projected.collect()]
+ self.assertEqual(results, [True, False, True])
+
+ def test_udf_transpile_float_eq_nan_semantics(self):
+ # `a == b` on DOUBLE columns must correctly implement Python's NaN
semantics:
+ # IEEE 754 says NaN != NaN (Python returns False), while Spark's
EqualTo
+ # returns True for NaN = NaN. The NaN guard fixes this for float
columns.
+ from pyspark.sql.types import DoubleType, StructField, StructType
+
+ def x_eq_y(x: float, y: float) -> bool:
+ return x == y
+
+ def x_neq_y(x: float, y: float) -> bool:
+ return x != y
+
+ schema = StructType(
+ [
+ StructField("a", DoubleType(), nullable=True),
+ StructField("b", DoubleType(), nullable=True),
+ ]
+ )
+ nan = float("nan")
+ with self.sql_conf(_TRANSPILE_ON):
+ pudf_eq = UserDefinedFunction(x_eq_y, BooleanType())
+ pudf_neq = UserDefinedFunction(x_neq_y, BooleanType())
+ self.assertTrue(pudf_eq.transpiled, "annotated float == float
should transpile")
+ self.assertTrue(pudf_neq.transpiled, "annotated float != float
should transpile")
+ # Verify isnan IS present in the logical plan (NaN guard was
emitted).
+ df = self.spark.createDataFrame([(1.0, 1.0)], schema=schema)
+ projected_eq = df.select(pudf_eq("a", "b").alias("r"))
+ logical_plan =
projected_eq._jdf.queryExecution().analyzed().toString()
+ self.assertIn("isnan", logical_plan.lower(), "float == must emit
NaN guard")
+ # Verify semantics: NaN == NaN -> False (Python), None == None ->
True.
+ rows = [
+ (nan, nan), # NaN == NaN -> False in Python
+ (nan, 1.0), # NaN == 1.0 -> False in Python
+ (1.0, 1.0), # 1.0 == 1.0 -> True
+ (1.0, 2.0), # 1.0 == 2.0 -> False
+ (None, None), # None == None -> True (Python semantics)
+ (None, 1.0), # None == 1.0 -> False
+ ]
+ df = self.spark.createDataFrame(rows, schema=schema)
+ eq_results = [r[0] for r in df.select(pudf_eq("a", "b")).collect()]
Review Comment:
**Nit (P3):** These result assertions run on a DataFrame whose plan is never
checked, and interpreted Python produces exactly these lists, so they also pass
if the call silently falls back. The `isnan` check above covers only `pudf_eq`
on a different one-row frame, and `pudf_neq` has no plan check; `_vals(...,
require_lowered=True)` in this file exists for exactly this case. Coverage is
also limited to bare, same-family annotated parameters. Nothing exercises an
untyped parameter reaching the float variant with NaN data, mixed equality such
as `x == 0` on DOUBLE, or derived operands such as `(a + 1) == (b + 1)`, whose
guard decision depends on `_wider_numeric`; if integer-plus-float widening
regressed to `"integer"`, those would return Spark's `NaN == NaN` result with
every test green. The header comment of
`test_udf_transpile_int_float_annotation_categories` ("the wider category is
used") is also not backed by its assertions, which only check helper outputs.
**Verification:**
- **Regression:** The NaN-equality tests fail when the guard is skipped for
float-derived operands, for example if integer-plus-float widening returned the
integer category, and when the call falls back to interpreted Python.
##########
python/pyspark/sql/transpile.py:
##########
@@ -496,15 +560,25 @@ def _category_uncached(self, params: List[str], node:
ast.AST) -> str:
case ast.BinOp(left=left, op=op, right=right):
lc = self._category(params, left)
rc = self._category(params, right)
- if isinstance(op, ast.Add) and lc == rc:
- return lc # str + str -> str, num + num -> num
+ if isinstance(op, ast.Add):
+ if lc == rc:
+ return lc # str+str, integer+integer, float+float
+ if _is_numeric_cat(lc) and _is_numeric_cat(rc):
+ return _wider_numeric(lc, rc) # integer+float -> float
if isinstance(op, ast.Mult):
- if {lc, rc} == {"numeric", "numeric"}:
- return "numeric"
- if {lc, rc} == {"numeric", "string"}:
- return "string" # str * int / int * str -> repeat
- if isinstance(op, (ast.Sub, ast.Mod)) and lc == rc ==
"numeric":
- return "numeric"
+ if _is_numeric_cat(lc) and _is_numeric_cat(rc):
+ return _wider_numeric(lc, rc)
+ # repeat(str, n): second operand must be integral (not
float).
+ # Python raises TypeError for `str * float`, so we allow
only
+ # "integer" or the legacy "numeric" catch-all (which the
JVM
Review Comment:
**Nit (P3):** No JVM step validates this. The JVM only receives
per-parameter categories from `_param_category_combos`, which never assigns
`"numeric"` after this change, and even a parameter-level `"numeric"` matches
FloatType/DoubleType in `optionMatchesTypes`. At this check `"numeric"` is
always an expression-level label: unary `+`/`-` (including negative literals),
`len()`, or `_wider_numeric` absorbing `"float"`. So `def f(s: str, n: float):
return s * -n` still lowers to `repeat(s, cast(-n as int))` and returns a
string where Python raises TypeError. That divergence predates this PR, but the
comment now tells maintainers it is closed. Please describe what is actually
enforced: only `"integer"` comes from a validated parameter type, and an
expression-level `"numeric"` multiplier can still be a float.
**Verification:**
- **Inspection:** The comment on the string-repetition category check
matches the categories the check accepts and claims no JVM validation for
expression-level categories.
##########
python/pyspark/sql/transpile.py:
##########
@@ -754,20 +833,20 @@ def _convert_chunk(self, params: List[str], body: ast.AST
| None) -> Column:
case ast.Add():
if lc == rc == "string":
return concat(left_col, right_col)
- if lc == rc == "numeric":
+ if _is_numeric_cat(lc) and _is_numeric_cat(rc):
return left_col.__add__(right_col)
case ast.Sub():
- if lc == rc == "numeric":
+ if _is_numeric_cat(lc) and _is_numeric_cat(rc):
return left_col.__sub__(right_col)
case ast.Mult():
- if lc == "numeric" and rc == "numeric":
+ if _is_numeric_cat(lc) and _is_numeric_cat(rc):
return left_col.__mul__(right_col)
- if lc == "string" and rc == "numeric":
+ if lc == "string" and rc in ("integer", "numeric"):
Review Comment:
**Nit (P3):** This check (and the matching one at lines 575-578) newly
refuses a float multiplier. On branch-4.x before this change, a
string-returning UDF such as `a * 2.0` on a STRING column, or `a * b` over
(string, double), lowered to `repeat(s, cast(n as int))` and silently returned
a truncated repetition where Python raises TypeError. That is a real
correctness fix, but nothing tests it:
`test_udf_transpile_string_operands_fall_back` has no `str * float` case, so
reverting to `_is_numeric_cat(rc)` would leave every test green. When adding
the case, declare a string return type: the existing `_raises` helper declares
`LongType`, under which the repeat variant already fails the return-type check,
so a case built on it would pass with or without this fix.
**Verification:**
- **Regression:** A string-times-float UDF declared with a string return
type, with either a float literal or a DOUBLE-bound multiplier, runs as
interpreted Python and raises TypeError, and the test fails if the multiplier
check is reverted so the call lowers to repeat.
##########
python/pyspark/sql/transpile.py:
##########
@@ -372,31 +405,56 @@ def _lower_eq(
where Python's ``==`` is simply False. Refuse those so the UDF falls
back to interpreted Python. A ``None`` literal operand stays allowed
(the four-branch NULL handling above reproduces Python exactly).
+ Numeric sub-categories (``"integer"`` vs ``"float"``) are treated as
+ compatible; Spark widens to double in the comparison.
+
+ A NaN guard is added unless both operands are provably ``"integer"``:
+ ``NaN == NaN`` is ``False`` in Python (IEEE 754), but Spark's
+ ``EqualTo`` returns ``True``. ``"numeric"`` is an unknown sub-type
+ that may resolve to a float column, so it keeps the guard.
- One value-level difference remains (needs runtime values, so it is
- documented, not guarded): Spark treats ``NaN = NaN`` as true, while
- Python's ``nan == nan`` is False.
+ One value-level difference remains for ordering comparisons (tracked
+ separately): Spark orders NaN as greater than every value, whereas
+ Python's NaN comparisons are all False.
"""
lc = self._safe_category(params, left_node)
rc = self._safe_category(params, right_node)
if lc is not None and rc is not None and lc != rc:
- raise UnsupportedOperationException(
- f"`==`/`!=` operands have incompatible categories ({lc} vs
{rc}); "
- "Python compares across types as unequal while Spark would
coerce "
- "or fail analysis, so the transpiler falls back to interpreted
Python"
- )
+ # Numeric sub-categories are mutually compatible for equality.
+ if not (_is_numeric_cat(lc) and _is_numeric_cat(rc)):
+ raise UnsupportedOperationException(
+ f"`==`/`!=` operands have incompatible categories ({lc} vs
{rc}); "
+ "Python compares across types as unequal while Spark would
coerce "
+ "or fail analysis, so the transpiler falls back to
interpreted Python"
+ )
left_col = self._convert_chunk(params, left_node)
right_col = self._convert_chunk(params, right_node)
left_null = left_col.isNull()
right_null = right_col.isNull()
+ # NaN guard: Python's `NaN == NaN` is False (IEEE 754 reflexivity
fails),
+ # but Spark's EqualTo returns True. Emit only for numeric operands, and
+ # skip only when both are provably integral ("integer"); "numeric" is
an
+ # unknown sub-type that may resolve to a float column, so it keeps the
guard.
+ if (
+ lc is not None
+ and rc is not None
+ and _is_numeric_cat(lc)
+ and _is_numeric_cat(rc)
+ and not (lc == "integer" and rc == "integer")
+ ):
+ nan_result = lit(not equal)
+ nan_cmp = isnan(left_col) | isnan(right_col)
Review Comment:
**Non-blocking (P2):** Spark's `EqualTo` diverges from Python only when both
operands are NaN; with one NaN operand both already return false for `==`. This
guard fires whenever either side may be NaN, so it is dead for `x == 0.5` or `x
== 0` on DOUBLE, or `x == 0.5` on BIGINT. Its CASE form also survives
optimization: `isnan(0.5)` folds to false, but `CASE WHEN isnan(x) THEN false
ELSE x = 0.5 END` remains, and `SimplifyConditionals` only rewrites the
true/false CASE pairs.
Compared with branch-4.x before this change, which emitted no guard:
- A filter `udf(is_half, "boolean")(t.x)` with `def is_half(x: float) ->
bool: return x == 0.5` over a JDBC catalog table (DSv2; JDBC columns are always
nullable) used to reduce to `CASE WHEN isnull(x) THEN false ELSE x = 0.5 END`,
which `V2ExpressionBuilder` translates and JDBC can push. With the guard the
predicate contains `isnan`, which has no V2 translation, so it is not pushed
and the source returns every row for Spark to filter.
- A join condition `udf(same, "boolean")(l.a, r.b)` with `def same(a: float,
b: float) -> bool: return a == b` over non-nullable DOUBLE keys (for example
computed with `coalesce`) used to fold to `a = b`, which `ExtractEquiJoinKeys`
takes as an equi-key. Now no top-level `EqualTo` remains, so the join becomes a
broadcast nested-loop or Cartesian join.
Guarding only the both-NaN case inside the existing NULL handling keeps the
semantics: `(l = r) AND NOT (isnan(l) AND isnan(r))` for `==` and `(l != r) OR
(isnan(l) AND isnan(r))` for `!=`, emitted only when neither operand is
`"integer"`. Against a non-NaN literal this folds back to the plain comparison,
and between columns it keeps the equality conjunct the planner needs.
**Verification:**
- **Behavior:** Lowered equality and inequality over floating-point rows
containing NaN, NULL and ordinary values return Python's results.
- **Regression:** The optimized plan of a lowered floating-point
column-versus-literal equality contains no isnan, and a lowered equality over
non-nullable floating-point keys used as a join condition yields an equi-join
key.
--
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]
