[SPARK-55206][PYTHON][FOLLOWUP] Preserve Python semantics in UDF transpilation

[cloud-fan]

What changes were proposed in this pull request?

Followup to #56327.

This change makes the experimental Python UDF transpiler preserve Python semantics more
conservatively. Numeric and non-literal unary arithmetic, modulo, and string repetition fall back
to interpreted Python when fixed-width Catalyst operations cannot preserve Python behavior.
String concatenation raises on null inputs as Python would, and numeric equality and ordering
handle NaN according to Python semantics.

The transpilation unit tests are updated to assert the safe fallback paths and the corrected NaN
behavior.

Why are the changes needed?

Python integers have arbitrary precision, while Catalyst numeric arithmetic uses fixed-width
types. Lowering these expressions can therefore overflow or otherwise produce behavior different
from the original Python UDF. Spark also treats NaN differently from Python for equality and
ordering, and Catalyst normally propagates null through concatenation where Python raises a
TypeError. The transpiler must fail closed whenever it cannot preserve the source UDF's behavior.

Does this PR introduce any user-facing change?

Yes. When the unreleased experimental Python UDF transpilation feature is explicitly enabled,
unsafe arithmetic expressions now remain interpreted Python UDFs, and transpiled comparisons and
string concatenation more closely match Python behavior.

How was this patch tested?

Updated pyspark.sql.tests.test_udf_transpile_unit with positive and negative cases covering
numeric fallback, unary operations, modulo, string operations, nulls, and NaN comparisons.

• build/sbt -Phive package
• python/run-tests --testnames pyspark.sql.tests.test_udf_transpile_unit

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-5)

[cloud-fan]

cc @holdenk

[holdenk]

Oooh very cool :) let me take a look

[holdenk]

Hey @cloud-fan see some CI failures that seem related, can you take a look?

[holdenk]

Ok so this is a design choice we can think about, I'm in favor of transpiling the numeric operators even if overflow is possible provided we document that behavior. Otherwise we could catch the overflow and promote (for example try_... and then on nulls we promote).

[holdenk]

Or more simply add a catalyst expr which does the automatic promotion instead of encoding that in weird branching. But I think (for v0) document and throw is ok.

[cloud-fan]

Agreed. For v0, I restored numeric and unary transpilation and kept ANSI overflow as a documented divergence. Exact promotion via try_* or a dedicated Catalyst expression can follow under SPARK-55210.

The PR still fixes NULL handling, NaN comparisons, and string concatenation; unsafe string repetition falls back to Python. I also fixed the related CI and formatting failures in cb4b842. Thanks!

[holdenk]

@cloud-fan still seeing a lot of related unit test failures around UDF presence in plan.

[cloud-fan]

@holdenk CI green now :)

[dongjoon-hyun]

Thanks for the follow-up. The direction looks right to me -- failing closed on NULL arithmetic, handling NaN in comparisons, and guarding string concat are all real fidelity improvements, and CI is green. Three things I'd like to raise before this goes in.

1. The NULL guard duplicates operand subtrees, so nested arithmetic grows exponentially

In _convert_chunk's BinOp arm, left_col / right_col appear both in the guard condition and in the otherwise branch:

null_guard=left_col.isNull() |right_col.isNull()
...
returnwhen(null_guard, raise_error(null_error)).otherwise(left_col.__add__(right_col))

Each nesting level therefore emits two copies of each child subtree. Counting node occurrences, a_n = 7 + 2*a_(n-1) + 2, i.e. a_n = 20 * 2^(n-1) - 9:

expression	before this PR	after this PR
x + 1	3	11
x + 1 chained 10x	21	~10,200
x + 1 chained 20x	41	~10,500,000

Catalyst expressions are immutable case classes, so the initial object graph is shared, but transform / canonicalized / codegen all visit each occurrence, so analysis and optimization time grows with the occurrence count, not the shared node count. Subexpression elimination only mitigates runtime CPU, not planning; and the generated code can blow past the 64KB method limit and fall back to interpreted eval. The Mod case is worse: sb references right_col twice, so the right operand is duplicated ~6x per level.

To be clear about what is new here: the same duplication pattern already existed in _lower_value_compare and _lower_eq on master. It did not bite, because a comparison usually sits once at the top of a UDF body while the arithmetic underneath it was linear. Extending the pattern to + - * % and unary -- where nesting actually happens -- is what turns it from linear into exponential.

Suggestion: track nullability statically on the Python side and only emit an isNull() term for operands that can actually be NULL:

• ast.Constant (non-None) -> not nullable
• an already-guarded arithmetic / concat result -> not nullable (under ANSI, Add/concat/pmod over non-NULL operands cannot produce NULL)
• ast.Name (a bound parameter) -> nullable

With that, x + 1 + 1 + ... guards x once at the innermost level and emits no guard above it, which brings the chain back to linear. This has to live in the transpiler: Catalyst sees CaseWhen.nullable = true because of the RaiseError branch, so it cannot infer it for us. The unary arm has the same shape (operand_col referenced twice) and would benefit from the same treatment.

2. isnan on integral columns adds a per-row cast whose result is always false

IsNaN declares inputTypes = Seq(TypeCollection(DoubleType, FloatType)) with ImplicitCastInputTypes, while the "numeric" category on the JVM side matches NumericType && !DecimalType (ResolveTranspiledPythonUDFOptions.optionMatchesTypes) -- so Byte/Short/Int/Long bind to it too.

That means a plain lambda x: x > 0 on a long column now evaluates IsNaN(cast(a as double)) for every row, and the answer is always false. Analysis does resolve (Cast.canANSIStoreAssign(LongType, DoubleType) is true under ANSI), so this is a cost issue rather than a correctness one, but it lands on what is probably the most common transpiled shape.

Would it be worth splitting the numeric category into integral / fractional variants so the NaN guard is only emitted where NaN is representable? If that is too much for v0, a comment noting the overhead plus a follow-up JIRA would at least keep it visible.

3. Leftover repeat references now that the lowering is gone

These lines were accurate on master and became stale when this PR dropped the repeat lowering:

• _category still returns "string" for {numeric, string}Mult with the comment # str * int / int * str -> repeat. Behavior is still fail-closed (_convert_chunk raises and the variant is dropped), but it is now a dead branch with a misleading comment.
• the _category catch-all comment still says "don't drive concat/repeat selection".
• ResolveTranspiledPythonUDFOptions.scala: "so the string lowerings (e.g. repeat) never see it" now points at a lowering that no longer exists.

Worth cleaning up in this PR since it is the one removing the feature.

[holdenk]

For #1 I've got https://issues.apache.org/jira/browse/SPARK-58628 filed I'm happy to take it on post merge of this if you're ok with that @dongjoon-hyun ? I could also try and land it first and we could go rebase on top

[cloud-fan]

Thanks @dongjoon-hyun. I addressed these in 2f700fd2f04:

1. Added a TODO(SPARK-58628) at the NULL guards to track eliminating the expression growth. As Holden mentioned above, is it okay to handle this post-merge, or would you prefer SPARK-58628 to land first?
2. Split numeric comparisons into integral and fractional variants. Integral inputs no longer emit isnan; Catalyst and PySpark regression tests cover this.
3. Removed the dead string-multiplication categorization and the stale repeat comments.

Verification passed: build/sbt -Phive package, the Catalyst transpilation suite (6/6), the focused PySpark regression test, and the full pyspark.sql.tests.test_udf_transpile_unit suite.

[holdenk]

LGTM if @dongjoon-hyun is ok with the defered null guard resolution.