This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 5d1840dce8c5 [SPARK-59091][SQL] Emit the NaN guard in transpiled
equality only for floating-point columns
5d1840dce8c5 is described below
commit 5d1840dce8c5aecab89f554ef232b766ce035bad
Author: jzhan-2026 <[email protected]>
AuthorDate: Wed Oct 28 16:47:07 2026 -0700
[SPARK-59091][SQL] Emit the NaN guard in transpiled equality only for
floating-point columns
### What changes were proposed in this pull request?
Split the transpiler's `"numeric"` category into `"integer"` (IntegralType)
and `"float"`
(FractionalType, excluding DecimalType). The NaN guard in `_lower_eq` is
now only emitted
when at least one operand has category `"float"`.
- `transpile.py`: `_category` returns `"integer"` / `"float"` instead of
`"numeric"` for
literals and parameters; `_param_category_combos` tries `["integer",
"float", "string"]`
per parameter; `_lower_eq` gates the guard on `lc == "float" or rc ==
"float"`;
`_convert_chunk` BinOp dispatch uses `_is_numeric_cat()` for all three
sub-categories.
- `ResolveTranspiledPythonUDFOptions.scala`: adds `"integer"` →
IntegralType and `"float"`
→ FractionalType (excluding Decimal) cases to `optionMatchesTypes`.
Co-authored with jing.zhansnowflake.com
### Why are the changes needed?
For integer columns (e.g. `BIGINT`), the old plan for `x == y` included
`isnan(cast(x AS DOUBLE))`, which is always false—integers can't be NaN. The
NaN guard only makes sense for floating-point types.
### Does this PR introduce _any_ user-facing change?
No. The generated Catalyst plan for integer-column equality no longer
contains dead `isnan` guards, but the query results are identical.
### How was this patch tested?
New Python tests in `test_udf_transpile_unit`:
- `test_udf_transpile_integer_eq_no_nan_guard` — no `isnan` in plan for
`BIGINT ==`
- `test_udf_transpile_float_eq_nan_semantics` — `isnan` present for `DOUBLE
==`; NaN and null semantics verified end-to-end
- `test_udf_transpile_int_float_annotation_categories` — unit tests for new
category names
New Scala tests in `ResolveTranspiledPythonUDFOptionsSuite`: `"integer"` /
`"float"` matching, cross-type rejection, DecimalType exclusion, FloatType
matching.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored by: Claude Sonnet 4.6
Closes #58901 from jzhan-2026/spark-59091-nan-guard.
Authored-by: jzhan-2026 <[email protected]>
Signed-off-by: Holden Karau <[email protected]>
---
.../pyspark/sql/tests/test_udf_transpile_unit.py | 115 ++++++++++-
python/pyspark/sql/transpile.py | 219 ++++++++++++++-------
.../ResolveTranspiledPythonUDFOptions.scala | 31 +--
.../ResolveTranspiledPythonUDFOptionsSuite.scala | 66 ++++++-
4 files changed, 344 insertions(+), 87 deletions(-)
diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py
b/python/pyspark/sql/tests/test_udf_transpile_unit.py
index 6a650a08a767..5c0309bb53b3 100644
--- a/python/pyspark/sql/tests/test_udf_transpile_unit.py
+++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py
@@ -1640,12 +1640,12 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
# 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))
self.assertTrue(all(r[0] for r in eq_df.collect()))
clamped = rows.select(clamp_udf(rand()).alias("v"))
@@ -2636,17 +2636,120 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
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()]
+ neq_results = [r[0] for r in df.select(pudf_neq("a",
"b")).collect()]
+ self.assertEqual(eq_results, [False, False, True, False, True,
False])
+ self.assertEqual(neq_results, [True, True, False, True, False,
True])
+
+ def test_udf_transpile_int_float_annotation_categories(self):
+ # `int` annotation -> "integer" category; `float` annotation -> "float"
+ # category. For mixed-numeric comparisons the wider category is used.
+ import ast as _ast
+
+ from pyspark.sql.transpile import _annotation_category,
_param_category_combos
+
+ # _annotation_category: int -> "integer", float -> "float".
+ fn_ann = _ast.parse("def f(x: int, y: float): pass").body[0]
+ int_ann = fn_ann.args.args[0].annotation
+ flt_ann = fn_ann.args.args[1].annotation
+ self.assertEqual(_annotation_category(int_ann), "integer")
+ self.assertEqual(_annotation_category(flt_ann), "float")
+
+ # _param_category_combos pins typed params and tries all three for
untyped.
+ fn_typed = _ast.parse("def f(a: int, b: float): return a + b").body[0]
+ combos = _param_category_combos(fn_typed, ["a", "b"])
+ # One combo: both typed -> single variant {0: "integer", 1: "float"}.
+ self.assertEqual(len(combos), 1)
+ self.assertEqual(combos[0][0], "integer")
+ self.assertEqual(combos[0][1], "float")
+
+ fn_untyped = _ast.parse("def f(a): return a").body[0]
+ combos_u = _param_category_combos(fn_untyped, ["a"])
+ # Three combos for a single untyped param: integer / float / string.
+ self.assertEqual(len(combos_u), 3)
+ self.assertEqual({c[0] for c in combos_u}, {"integer", "float",
"string"})
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py
index 21289a15a7e1..9d371a3d64ca 100644
--- a/python/pyspark/sql/transpile.py
+++ b/python/pyspark/sql/transpile.py
@@ -111,6 +111,7 @@ from pyspark.sql.functions import (
coalesce,
col,
concat,
+ isnan,
length,
lit,
pmod,
@@ -131,6 +132,27 @@ if TYPE_CHECKING:
from pyspark.sql import SparkSession
from pyspark.sql._typing import DataTypeOrString
+# Categories that represent numeric (non-string) scalar types.
+_NUMERIC_CATS = frozenset(("numeric", "integer", "float"))
+
+
+def _is_numeric_cat(cat: str) -> bool:
+ """True for any numeric sub-category (``"numeric"``, ``"integer"``,
``"float"``)."""
+ return cat in _NUMERIC_CATS
+
+
+def _wider_numeric(lc: str, rc: str) -> str:
+ """Return the wider of two numeric categories (numeric > float >
integer)."""
+ if not (_is_numeric_cat(lc) and _is_numeric_cat(rc)):
+ raise UnsupportedOperationException(
+ f"_wider_numeric called with non-numeric categories: {lc!r},
{rc!r}"
+ )
+ if "numeric" in (lc, rc):
+ return "numeric"
+ if "float" in (lc, rc):
+ return "float"
+ return "integer"
+
class AbstractTranspiler(object):
"""Base class for transpilers. All experimental."""
@@ -219,16 +241,17 @@ def _truthiness_col(cat: Optional[str], c: Column) ->
Optional[Column]:
Returns ``None`` for unsupported or unknown categories (caller falls back).
Semantics (NULL-as-False throughout, matching Python's ``None`` is falsy):
- "bool" -> coalesce(c, False)
- "string" -> coalesce(length(c) > 0, False) -- empty string is falsy
- "numeric" -> coalesce(c != 0, False) -- zero is falsy
- NaN != 0 is True in Spark so float NaN is truthy, matching
Python.
+ "bool" -> coalesce(c, False)
+ "string" -> coalesce(length(c) > 0, False) -- empty is
falsy
+ "numeric"/"integer"/"float" -> coalesce(c != 0, False) -- zero is
falsy
+ NaN != 0 is True in Spark so float NaN is
+ truthy, matching Python.
"""
if cat == "bool":
return coalesce(c, lit(False))
if cat == "string":
return coalesce(length(c) > lit(0), lit(False))
- if cat == "numeric":
+ if cat is not None and _is_numeric_cat(cat):
return coalesce(c != lit(0), lit(False))
return None
@@ -286,6 +309,12 @@ class CatalystTranspiler(AbstractTranspiler):
body_c = self._safe_category(params, node.body[0]) if node.body
else None
else_c = self._safe_category(params, node.orelse[0]) if
node.orelse else None
if body_c is not None and else_c is not None and body_c != else_c:
+ # Numeric sub-categories are mutually compatible; return the
+ # wider one. Truly mismatched families (e.g. numeric vs string)
+ # return None so the branch-compatibility check in
+ # _convert_if_like raises and the variant is dropped.
+ if _is_numeric_cat(body_c) and _is_numeric_cat(else_c):
+ return _wider_numeric(body_c, else_c)
return None
return body_c if body_c is not None else else_c
if isinstance(node, ast.Constant) and node.value is None:
@@ -341,11 +370,15 @@ class CatalystTranspiler(AbstractTranspiler):
body_cat = self._safe_category(params, body_node)
else_cat = self._safe_category(params, else_node)
if body_cat is not None and else_cat is not None and body_cat !=
else_cat:
- raise UnsupportedOperationException(
- f"if/else branches have incompatible categories ({body_cat} vs
"
- f"{else_cat}); the lowered CASE WHEN has no common type under
ANSI, "
- "so the transpiler falls back to interpreted Python"
- )
+ # Numeric sub-categories (integer/float) are compatible: Spark
widens
+ # to double in the CASE WHEN expression. Truly mismatched families
+ # (e.g. numeric vs string) are still refused.
+ if not (_is_numeric_cat(body_cat) and _is_numeric_cat(else_cat)):
+ raise UnsupportedOperationException(
+ f"if/else branches have incompatible categories
({body_cat} vs "
+ f"{else_cat}); the lowered CASE WHEN has no common type
under ANSI, "
+ "so the transpiler falls back to interpreted Python"
+ )
return when(safe_test, body_col).otherwise(else_col)
def _lower_eq(
@@ -355,7 +388,7 @@ class CatalystTranspiler(AbstractTranspiler):
right_node: ast.AST,
equal: bool,
) -> Column:
- """Lower ``==`` / ``!=`` with Python's None-equality semantics.
+ """Lower ``==`` / ``!=`` with Python's None- and NaN-equality
semantics.
Unlike ordering operators, Python doesn't raise on ``None == x`` /
``None != x``: ``None == None`` is True, ``None == 0`` is False,
@@ -372,31 +405,56 @@ class CatalystTranspiler(AbstractTranspiler):
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)
+ value_cmp: Column = when(nan_cmp, nan_result).otherwise(
+ left_col == right_col if equal else left_col != right_col
+ )
+ else:
+ value_cmp = left_col == right_col if equal else left_col !=
right_col
if equal:
both_null_val: Column = lit(True)
one_null_val: Column = lit(False)
- value_cmp = left_col == right_col
else:
both_null_val = lit(False)
one_null_val = lit(True)
- value_cmp = left_col != right_col
return (
when(left_null & right_null, both_null_val)
.when(left_null | right_null, one_null_val)
@@ -435,7 +493,7 @@ class CatalystTranspiler(AbstractTranspiler):
"""
lc = self._category(params, left_node)
rc = self._category(params, right_node)
- if lc != rc:
+ if lc != rc and not (_is_numeric_cat(lc) and _is_numeric_cat(rc)):
raise UnsupportedOperationException(
f"`{op_repr}` compares operands of different categories "
f"({lc} vs {rc}); Python would raise TypeError, so the "
@@ -452,9 +510,12 @@ class CatalystTranspiler(AbstractTranspiler):
return when(null_guard, raise_error(err)).otherwise(op(left_col,
right_col))
def _category(self, params: List[str], node: ast.AST) -> str:
- """Infer ``"numeric"`` or ``"string"`` for ``node`` under the current
+ """Infer the scalar category of ``node`` under the current
``self._param_categories`` assumption (set per input-type variant).
+ Returns one of ``"integer"``, ``"float"``, ``"numeric"`` (unknown
+ numeric sub-type), ``"string"``, ``"bool"``, or ``"binary"``.
+
Drives operator selection (``+`` -> add vs concat, ``*`` -> multiply vs
repeat) and raises ``UnsupportedOperationException`` when an operator's
operands are type-incompatible, so the caller drops that variant and
the
@@ -471,18 +532,21 @@ class CatalystTranspiler(AbstractTranspiler):
def _category_uncached(self, params: List[str], node: ast.AST) -> str:
match node:
case ast.Constant(value=v):
- # bool subclasses int, so classify it first: int/float ->
numeric,
- # str -> string, bool -> bool, bytes -> binary. None/complex/
- # Ellipsis have no usable Spark column type, so raise to drop
this
- # variant and fall back rather than emit an option that fails
- # CheckAnalysis or silently diverges (e.g. `x + None` -> NULL
where
- # Python raises TypeError).
+ # bool subclasses int, so classify it first.
+ # int -> "integer", float -> "float", str -> "string",
+ # bool -> "bool", bytes -> "binary".
+ # None/complex/Ellipsis have no usable Spark column type, so
+ # raise to drop this variant and fall back rather than emit an
+ # option that fails CheckAnalysis or silently diverges (e.g.
+ # `x + None` -> NULL where Python raises TypeError).
if isinstance(v, bool):
return "bool"
if isinstance(v, bytes):
return "binary"
- if isinstance(v, (int, float)):
- return "numeric"
+ if isinstance(v, float):
+ return "float"
+ if isinstance(v, int):
+ return "integer"
if isinstance(v, str):
return "string"
raise UnsupportedOperationException(
@@ -496,15 +560,25 @@ class CatalystTranspiler(AbstractTranspiler):
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
+ # will validate against the actual column type at query
time).
+ if lc == "string" and rc in ("integer", "numeric"):
+ return "string"
+ if rc == "string" and lc in ("integer", "numeric"):
+ return "string"
+ if isinstance(op, (ast.Sub, ast.Mod)):
+ if _is_numeric_cat(lc) and _is_numeric_cat(rc):
+ return _wider_numeric(lc, rc)
raise UnsupportedOperationException(
f"operands of `{type(op).__name__}` are not
type-compatible "
"for this input-type variant"
@@ -537,6 +611,11 @@ class CatalystTranspiler(AbstractTranspiler):
body_cat = branch_category(if_body)
else_cat = branch_category(if_orelse)
if body_cat is not None and else_cat is not None and body_cat
!= else_cat:
+ # Numeric sub-categories (integer/float) are mutually
+ # compatible: `1 if c else 1.0` is valid Python and Spark
+ # widens to double in the CASE WHEN.
+ if _is_numeric_cat(body_cat) and _is_numeric_cat(else_cat):
+ return _wider_numeric(body_cat, else_cat)
raise UnsupportedOperationException(
f"ternary branches have mismatched categories
({body_cat} "
f"vs {else_cat}) and cannot drive operator selection"
@@ -591,7 +670,7 @@ class CatalystTranspiler(AbstractTranspiler):
# of falling back, since the option is type-checked as a
# child of TranspiledPythonUDF before ConvertToCatalyst can
# drop it. Fail closed for every non-numeric category.
- if self._category(params, operand) != "numeric":
+ if not _is_numeric_cat(self._category(params, operand)):
raise UnsupportedOperationException(
"unary `+`/`-` is only supported for numeric operands "
"(Python raises TypeError on strings, and Spark would "
@@ -754,20 +833,20 @@ class CatalystTranspiler(AbstractTranspiler):
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"):
return repeat(left_col, right_col.cast("int"))
- if lc == "numeric" and rc == "string":
+ if lc in ("integer", "numeric") and rc == "string":
return repeat(right_col, left_col.cast("int"))
case ast.Mod():
- if lc == rc == "numeric":
+ if _is_numeric_cat(lc) and _is_numeric_cat(rc):
# Python's `%` takes the sign of the divisor;
Spark's
# takes the dividend's. `sign(b) * pmod(sign(b) *
a,
# abs(b))` reproduces Python for every non-zero
divisor
@@ -895,7 +974,7 @@ class CatalystTranspiler(AbstractTranspiler):
cast_ok = (
body_cat is None
or (
- body_cat == "numeric"
+ _is_numeric_cat(body_cat)
and isinstance(returnType, NumericType)
and not isinstance(returnType, DecimalType)
)
@@ -936,21 +1015,24 @@ def _get_transpilers(session: "SparkSession") ->
List[AbstractTranspiler]:
def _annotation_category(annotation: Optional[ast.AST]) -> Optional[str]:
- """Map a parameter's type annotation to a category
- (``"numeric"``/``"string"``/``"bool"``/``"binary"``), or ``None`` when it's
- absent or unrecognised (the caller then tries both numeric and string)."""
+ """Map a parameter's type annotation to a category, or ``None`` when absent
+ or unrecognised (the caller then tries all numeric and string variants).
+
+ Returns one of ``"integer"``, ``"float"``, ``"string"``, ``"bool"``, or
+ ``"binary"`` (matching the constant handling in ``_category``). ``int``
+ maps to ``"integer"`` and ``float`` maps to ``"float"`` so the JVM can
+ select the variant whose NaN semantics match the bound column type."""
name: Optional[str] = None
if isinstance(annotation, ast.Name):
name = annotation.id
elif isinstance(annotation, ast.Constant) and isinstance(annotation.value,
str):
name = annotation.value # stringized annotation, e.g. def f(a: "int")
- # str -> "string", int/float -> "numeric", bool -> "bool", bytes ->
"binary"
- # (matching the constant handling in ``_category``). complex and anything
- # unrecognised return None so the caller tries both numeric and string.
if name == "str":
return "string"
- if name in ("int", "float"):
- return "numeric"
+ if name == "int":
+ return "integer"
+ if name == "float":
+ return "float"
if name == "bool":
return "bool"
if name == "bytes":
@@ -959,14 +1041,16 @@ def _annotation_category(annotation: Optional[ast.AST])
-> Optional[str]:
def _param_category_combos(function_ast: ast.FunctionDef, public_params:
List[str]) -> List[dict]:
- """Per-variant maps ``{public_param_index -> category}`` where category is
- one of ``"numeric"``/``"string"``/``"bool"``/``"binary"``.
-
- A typed param (``def f(a: str, b: int)``) is pinned to its category; an
- untyped param is tried as both numeric and string. To cap plan growth, when
- more than three params are untyped we collapse the untyped ones to the
- all-numeric and all-string variants (encourage typing inputs to keep the
- matrix small) while keeping every typed param pinned.
+ """Per-variant maps ``{public_param_index -> category}``.
+
+ A typed param (``def f(a: str, b: int)``) is pinned to its category
+ (``"integer"``/``"float"``/``"string"``/``"bool"``/``"binary"``); an
+ untyped param is tried as ``"integer"``, ``"float"``, and ``"string"``
+ separately so the JVM can select the variant that matches the bound column
+ type and apply the right NaN semantics. To cap plan growth, when more than
+ three params are untyped we collapse the untyped ones to one variant per
+ base category (encourage typing inputs to keep the matrix small) while
+ keeping every typed param pinned.
"""
n = len(public_params)
all_args = _positional_args(function_ast)
@@ -976,17 +1060,16 @@ def _param_category_combos(function_ast:
ast.FunctionDef, public_params: List[st
for arg in public_args:
cat = _annotation_category(arg.annotation)
if cat is None:
- candidates.append(["numeric", "string"])
+ candidates.append(["integer", "float", "string"])
untyped += 1
else:
candidates.append([cat])
if untyped > 3:
- # Cap the 2**untyped blow-up, but keep each typed param pinned to its
- # category (a single-element ``candidates`` entry); only the untyped
- # params collapse to the all-numeric / all-string pair.
+ # Cap the 3**untyped blow-up; collapse untyped params to one variant
+ # per base fill while keeping typed params pinned.
return [
{i: c[0] if len(c) == 1 else fill for i, c in
enumerate(candidates)}
- for fill in ("numeric", "string")
+ for fill in ("integer", "float", "string")
]
return [{i: choice[i] for i in range(n)} for choice in
itertools.product(*candidates)] or [{}]
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala
index d9b07093fb89..4559831f767f 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala
@@ -21,7 +21,9 @@ import
org.apache.spark.sql.catalyst.expressions.{TranspiledPythonUDF, Transpile
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.TRANSPILED_PYTHON_UDF
-import org.apache.spark.sql.types.{BinaryType, BooleanType, DataType,
DecimalType, NumericType, StringType}
+import org.apache.spark.sql.types.{
+ BinaryType, BooleanType, DataType, DecimalType, FractionalType,
+ IntegralType, NumericType, StringType}
/**
* Prunes the per-input-type options carried by a [[TranspiledPythonUDF]] down
to those whose
@@ -76,18 +78,22 @@ object ResolveTranspiledPythonUDFOptions extends
Rule[LogicalPlan] {
}
}
- // True when each declared category matches the corresponding argument type:
- // "numeric" -> NumericType, "string" -> StringType, "bool" -> BooleanType,
- // "binary" -> BinaryType. "string" matches only StringType (not
BinaryType): a
- // bytes/BinaryType column is tagged "binary" instead, so the string
lowerings
- // (e.g. `repeat`) never see it. Empty categories means "no restriction", so
the
- // option is kept.
+ // True when each declared category matches the corresponding argument type.
+ // Empty categories means "no restriction" and the option is kept.
+ //
+ // Category -> Spark type mapping:
+ // "integer" -> IntegralType (byte/short/int/long; NaN impossible)
+ // "float" -> FractionalType minus DecimalType (float/double; NaN
possible)
+ // "numeric" -> any NumericType minus DecimalType (legacy catch-all)
+ // "string" -> StringType with UTF8_BINARY collation only
+ // "bool" -> BooleanType
+ // "binary" -> BinaryType
//
// Two deliberate exclusions keep the transpiled semantics faithful to
Python:
- // - DecimalType is NOT "numeric": Python receives decimal.Decimal objects,
- // which raise TypeError when mixed with float literals and carry different
- // precision semantics than Spark's decimal arithmetic, so decimal columns
- // fall back to interpreted Python.
+ // - DecimalType is NOT matched by any numeric category: Python receives
+ // decimal.Decimal objects, which raise TypeError when mixed with float
+ // literals and carry different precision semantics than Spark's decimal
+ // arithmetic, so decimal columns fall back to interpreted Python.
// - "string" requires the default UTF8_BINARY collation: under a non-binary
// collation (e.g. UTF8_LCASE) Spark's `=`/`<`/`concat` follow collation
// rules while Python compares codepoints, so `'abc' == 'ABC'` would return
@@ -99,6 +105,9 @@ object ResolveTranspiledPythonUDFOptions extends
Rule[LogicalPlan] {
false
} else {
categories.zip(argTypes).forall {
+ case ("integer", dt) => dt.isInstanceOf[IntegralType]
+ case ("float", dt) =>
+ dt.isInstanceOf[FractionalType] && !dt.isInstanceOf[DecimalType]
case ("numeric", dt) => dt.isInstanceOf[NumericType] &&
!dt.isInstanceOf[DecimalType]
case ("string", st: StringType) => st.isUTF8BinaryCollation
case ("bool", dt) => dt.isInstanceOf[BooleanType]
diff --git
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala
index f9295eb17b00..4a12ecbbcb4c 100644
---
a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala
+++
b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptionsSuite.scala
@@ -19,10 +19,11 @@ package org.apache.spark.sql.catalyst.analysis
import org.apache.spark.api.python.PythonEvalType
import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{Add, Alias, Concat,
Expression, Literal, PythonUDF, TranspiledPythonUDF}
+import org.apache.spark.sql.catalyst.expressions.{
+ Add, Alias, AttributeReference, Concat, Expression, Literal, PythonUDF,
TranspiledPythonUDF}
import org.apache.spark.sql.catalyst.plans.PlanTest
import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project}
-import org.apache.spark.sql.types.LongType
+import org.apache.spark.sql.types.{DecimalType, DoubleType, FloatType,
LongType}
/**
* Unit tests for [[ResolveTranspiledPythonUDFOptions]], which prunes a
@@ -92,4 +93,65 @@ class ResolveTranspiledPythonUDFOptionsSuite extends
PlanTest {
val pruned = prune(node, LocalRelation(a))
assert(pruned.transpiledOptions == List(onlyOpt))
}
+
+ test("keeps the integer option for integral columns and drops float and
string") {
+ val a = $"a".long
+ val intOpt = Add(a, Literal(1L))
+ val floatOpt = Add(a, Literal(1.0))
+ val stringOpt = Concat(Seq(a, a))
+ val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)),
+ List(intOpt, floatOpt, stringOpt),
+ List(List("integer"), List("float"), List("string")))
+ val pruned = prune(node, LocalRelation(a))
+ assert(pruned.transpiledOptions == List(intOpt))
+ assert(pruned.optionInputCategories.isEmpty)
+ }
+
+ test("keeps the float option for fractional columns and drops integer and
string") {
+ val a = AttributeReference("a", DoubleType)()
+ val intOpt = Add(a, Literal(1L))
+ val floatOpt = Add(a, Literal(1.0))
+ val stringOpt = Concat(Seq(a, a))
+ val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)),
+ List(intOpt, floatOpt, stringOpt),
+ List(List("integer"), List("float"), List("string")))
+ val pruned = prune(node, LocalRelation(a))
+ assert(pruned.transpiledOptions == List(floatOpt))
+ assert(pruned.optionInputCategories.isEmpty)
+ }
+
+ test("integer does not match float columns; float does not match integer
columns") {
+ val aLong = $"a".long
+ val aDouble = AttributeReference("a", DoubleType)()
+ val intOpt = Add(aLong, Literal(1L))
+ val floatOpt = Add(aLong, Literal(1.0))
+ // float column, integer-only option -> dropped
+ val intOnlyForDouble = TranspiledPythonUDF("udf", pyUDF(Seq(aDouble)),
+ List(intOpt), List(List("integer")))
+ assert(prune(intOnlyForDouble,
LocalRelation(aDouble)).transpiledOptions.isEmpty)
+ // long column, float-only option -> dropped
+ val floatOnlyForLong = TranspiledPythonUDF("udf", pyUDF(Seq(aLong)),
+ List(floatOpt), List(List("float")))
+ assert(prune(floatOnlyForLong,
LocalRelation(aLong)).transpiledOptions.isEmpty)
+ }
+
+ test("neither integer nor float matches decimal columns") {
+ val a = AttributeReference("a", DecimalType(10, 2))()
+ val intOpt = Add(a, Literal(1L))
+ val floatOpt = Add(a, Literal(1.0))
+ val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)),
+ List(intOpt, floatOpt),
+ List(List("integer"), List("float")))
+ val pruned = prune(node, LocalRelation(a))
+ assert(pruned.transpiledOptions.isEmpty)
+ }
+
+ test("float matches FloatType columns") {
+ val a = AttributeReference("a", FloatType)()
+ val floatOpt = Add(a, Literal(1.0f))
+ val node = TranspiledPythonUDF("udf", pyUDF(Seq(a)),
+ List(floatOpt), List(List("float")))
+ val pruned = prune(node, LocalRelation(a))
+ assert(pruned.transpiledOptions == List(floatOpt))
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
