viirya opened a new pull request, #56940:
URL: https://github.com/apache/spark/pull/56940
### What changes were proposed in this pull request?
In the Arrow-optimized (non-legacy) regular Python UDF path, the UDF output
is converted from Python objects to an Arrow array. Previously, when the return
type needed a converter (e.g. `array<string>`), a per-element Python converter
(`LocalDataToArrowConversion`) was run over every element of every row before
building the Arrow array, even when the elements already matched the declared
type and PyArrow could build the array directly.
This PR adds a fast path: for return types with no value-transforming
coercion, it first tries `pa.array(results, type=arrow_return_type)` on the raw
UDF results and only falls back to the per-element converter if PyArrow rejects
them (an element genuinely needs coercion). A static per-UDF predicate,
`_output_fast_path_safe`, gates the fast path and excludes types whose
converter transforms a value that `pa.array` would also accept (silent
divergence) or where `pa.array` accepts input the converter is meant to reject:
- Timestamp / TimestampNTZ: converter truncates to the session time zone.
- Decimal: `pa.array` coerces int -> decimal, bypassing the
`intToDecimalCoercionEnabled` gate and rescaling / `Decimal('NaN')` handling.
- UDT / Variant / Geometry / Geography: converter serializes to a storage
form.
- Null.
Numeric / bool / string / binary and nested containers of those are safe.
### Why are the changes needed?
The per-element output converter is pure overhead when results already match
the declared type. For an `array<string>` UDF (the `reverse_array` shape), the
converter accounts for a large fraction of output time; a microbenchmark on
`select max(reverse_array(array(a, b)))` over 4M rows (best-of-6) shows the
Arrow path improving from 2.48s to 1.86s (~25% faster), narrowing the gap to
the legacy pandas path from 1.77x to 1.31x.
### Does this PR introduce _any_ user-facing change?
No. Output values are unchanged; the fast path is only taken when it
produces results identical to the existing converter path.
### How was this patch tested?
New tests in `test_arrow_python_udf.py` covering `array<string>` output
(already correct types), `array<string>` output requiring int->string coercion
(fallback path), and `array<timestamp>` output verifying session-timezone
handling (an excluded type). They run under both legacy and non-legacy
conversion configs. The full `test_arrow_python_udf` module (301 tests) passes.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
--
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]
