holdenk opened a new pull request, #59724:
URL: https://github.com/apache/spark/pull/59724
### What changes were proposed in this pull request?
Widens the eval types the Python UDF transpiler may rewrite, from the pickled
`SQL_BATCHED_UDF` alone to also cover `SQL_ARROW_BATCHED_UDF` (101) and
`SQL_SCALAR_PANDAS_UDF` (200). The two are admitted for different reasons,
and that
asymmetry is the substance of the change.
`SQL_ARROW_BATCHED_UDF` is not really vectorized -- Arrow is only the
transport, and the
worker still calls the function once per row with Python scalars. Its
Python-level
semantics are exactly what the transpiler already reproduces, so every
existing lowering
applies unchanged.
`SQL_SCALAR_PANDAS_UDF` is genuinely vectorized (one `pandas.Series` per
argument, a
Series required back), so most scalar lowerings do **not** carry over. Only
an
element-wise subset is rewritten, chosen by a strict allowlist checked
before any
lowering runs: parameter references, int/float/str constants, `+` `-` `*`
`%`, unary
`+`/`-`, and `<param>.isnull()` (plus the `.isna()` / `.notnull()` /
`.notna()` aliases).
Everything else falls back to interpreted Python. Each exclusion is a
construct whose
scalar lowering would be actively wrong on a Series -- `if s:` / `s and t` /
`not s`
raise `ValueError`, `s is None` is always False, `s < 0` treats a missing
value as False
rather than NULL or TypeError, and a body that never references a parameter
returns a
scalar the interpreted path rejects. Comparisons are therefore out of scope;
they need
their own NaN-aware lowering and are left as follow-up work.
Two pandas value rules are reproduced explicitly: `Series.isnull()` lowers to
`isnull(a) OR isnan(a)` (Arrow-to-pandas conversion has already collapsed a
NULL in a
fractional column into NaN), and a numeric body is wrapped in an
isnan-to-NULL
normalization (a numpy-backed result Series is masked with `series.isnull()`
on the way
out). That mask is conditional, which bounds validity: with
`preferIntExtensionDtype=true`
the Series is a pandas masked extension array, the mask is skipped, and
normalizing NaN
to NULL would be backwards -- so pandas transpilation is refused when that
conf is on.
Transpilation is likewise refused for an Arrow-optimized UDF under the
legacy pandas
conversion. Both confs can be flipped after the UDF is built, so
`ConvertToCatalyst`
re-checks them the way it already re-checks the ANSI gate.
The eval-type set is expressed as membership rather than negation, so a
newly added eval
type must be admitted consciously instead of inheriting transpilation by
failing to match
an exclusion; a canary over every `PythonEvalType` fails until a new one is
triaged.
### Why are the changes needed?
Arrow optimization is the default for regular Python UDFs, so before this the
transpilation flag did nothing on default configuration -- the feature was
effectively
unreachable without opting out of the default execution path. Extending to
scalar pandas
UDFs then covers the other commonly used scalar shape, conservatively.
### Does this PR introduce _any_ user-facing change?
Yes, when Python UDF transpilation is enabled (it is off by default).
Arrow-optimized and
scalar pandas UDFs that previously always ran as interpreted Python may now
be rewritten
into Catalyst expressions. Anything outside the supported subset falls back
to interpreted
Python, so behavior is unchanged there.
Known divergences for pandas are numpy's fixed-width arithmetic, documented
in
`transpile.py` and `SQLConf` and pinned by tests with the actual values:
integer overflow
wraps, float32 survives a float literal (NEP 50), an integral column becomes
float64 as
soon as a batch contains a NULL, and `a % 0` yields a missing value (or 0
under the
extension dtype). In each case the transpiled expression is the exact or
ANSI-correct one.
### How was this patch tested?
Targeted suites per eval type, split into differentials over the supported
subset,
refusal tests asserting behavior is unchanged (including that the raising
bodies still
raise), and divergence tests pinning the known differences. Plus an
eval-type canary and
a data-type canary that fail until a new type is triaged, and parity runs of
the existing
Arrow and scalar-pandas UDF mixins with the rewrite active. Both sides of
every
differential use the same eval type, since the pickled, Arrow and pandas
paths coerce
return values differently.
Run locally rebased on current master, JDK 21, all green:
| Suite | Result |
|---|---|
| `arrow.test_arrow_python_udf_transpile` | 47/47 |
| `arrow.test_arrow_python_udf_transpile_type_canary` | 14/14 |
| `arrow.test_arrow_python_udf_transpile_parity` | 307 OK (11 skipped) |
| `pandas.test_pandas_udf_transpile` | 42/42 |
| `pandas.test_pandas_udf_transpile_parity` | 81 OK (1 skipped) |
| `test_udf_transpile_unit` | 83/83 |
| `test_udf_transpile_parity` | 114 OK |
| `test_udf_transpile_hypothesis` | 23/23 (`RUN_HYPOTHESIS=1`, 200 examples)
|
| `test_udf` | 95/95 |
| catalyst `Optimizer*` suites + `ConvertToCatalystSuite` | all passed
(40/40) |
Two suites needed updating for changes that landed on master in the
meantime, in
a separate commit:
- `test_middle_of_arrow_udf_chain_is_not_elided` counted Arrow operators and
expected two. Master now pipelines a UDF chain into one
`ArrowEvalPythonExec`,
so the operator count is 1 whether or not the middle call was lowered. It
now
counts `PythonUDF` expression nodes, which still distinguishes the two,
with a
contrast case asserted alongside so it cannot pass by nothing being
lowered.
- The eval-type/data-type canary went red because it was working: master
added
`TimestampLTZNanosType` and `TimestampNTZNanosType`, and the canary fails
until
a new type is triaged. Both are `NO_TRANSPILE` as input and return type,
like
every other timestamp.
Transpilation is non-Connect only by design, so the parity suites are
guarded with
`is_remote_only()` and are intentionally not inherited into the Spark
Connect parity
tests; there is no Connect-side parity to run for this change.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5
Generated-by: GLM 5.2 (Cursor)
<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
<!-- coderabbit-ai-generated:v1 -->
## Summary by CodeRabbit
* **New Features**
* Experimental transpilation now supports regular Python UDFs, including
Arrow-optimized UDFs, and scalar pandas UDFs that use supported element-wise
operations.
* Eligible UDFs can use transpilation when the required settings are
enabled. Unsupported patterns and incompatible conversion settings continue to
use interpreted Python execution.
* **Documentation**
* Updated configuration guidance describes eligible UDF types, supported
operations, numerical differences, and conditions that use interpreted
execution.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
--
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]
