gaogaotiantian opened a new pull request, #59845: URL: https://github.com/apache/spark/pull/59845
### What changes were proposed in this pull request? `ArrowBatchedUDFArrowHandler` (Arrow-optimized `@udf`, `SQL_ARROW_BATCHED_UDF`) used to evaluate each UDF by zipping its argument columns into a list of row tuples and then calling `udf_func(*row)` for each row: ```python rows = list(zip(*[columns[o] for o in offsets])) ... return [udf_func(*row) for row in rows] ``` This PR passes the argument columns straight to `map()`. It takes the i-th element of each column and passes them positionally, so no tuple is built per row: ```python args = [columns[o] for o in offsets] ... return list(map(udf_func, *args)) ``` The `arrow.concurrency.level` thread-pool path changes the same way (`pool.map(udf_func, *args)`). Zero-arg UDFs keep a dedicated branch, because `map()` needs at least one iterable. This PR also adds a 3-arg `concat3_udf` case to `bench_eval_type.ArrowBatchedUDFTimeBench`; all existing cases there were single-arg. It adds a unit test for the thread-pool path with multi-arg and zero-arg UDFs. ### Why are the changes needed? Performance. The old code allocated one tuple per row, per UDF, per batch, only to unpack it right away. Tuples are tracked by the cyclic GC, so allocating thousands of them per batch also triggers extra GC collections, which scan the live Python column lists. Most of the measured gain comes from that: with `gc.disable()`, the difference drops to within noise. Results from `ArrowBatchedUDFTimeBench` (in-process `worker.main`, no JVM). These are min wall times; the baseline and this PR alternated across 5 rounds of 7 runs each. Apple Silicon, CPython 3.12, PyArrow 23.0.1: | Scenario | `identity_udf` (1 arg) | `concat3_udf` (3 args) | |---|---|---| | `sm_batch_few_col` | 17.44 -> 17.01 ms (-2.5%) | 19.51 -> 19.43 ms (-0.4%) | | `lg_batch_few_col` | 43.68 -> 41.32 ms (-5.4%) | 49.93 -> 47.36 ms (-5.2%) | | `pure_ints` | 63.50 -> 61.71 ms (-2.8%) | 71.38 -> 67.38 ms (-5.6%) | | `pure_floats` | 61.99 -> 59.46 ms (-4.1%) | 106.77 -> 103.33 ms (-3.2%) | | `pure_strings` | 115.03 -> 111.98 ms (-2.7%) | 117.20 -> 114.11 ms (-2.6%) | | `mixed_types` | 89.10 -> 86.47 ms (-3.0%) | 96.01 -> 92.49 ms (-3.7%) | | `sm_batch_many_col` | 43.78 -> 43.25 ms (-1.2%) | 44.09 -> 43.96 ms (-0.3%) | | `lg_batch_many_col` | 174.05 -> 171.57 ms (-1.4%) | 175.78 -> 172.38 ms (-1.9%) | This PR was faster in all 5 paired rounds for nearly every case. Run-to-run spread is around 5-15%, so the individual percentages are approximate. The `*_many_col` scenarios send 50 columns to a UDF that reads at most 3, so their time is dominated by converting input columns that the UDF never reads. In real queries, `EvalPythonEvaluatorFactory` projects only the UDF argument expressions, so the few-column numbers are closer to what real workloads get. ### Does this PR introduce _any_ user-facing change? No. Results, their order, and exception / `StopIteration` handling are unchanged. I checked that the decoded output of every `ArrowBatchedUDFTimeBench` scenario is identical before and after this change. ### How was this patch tested? - Existing tests in `pyspark.eval_handlers.tests.test_arrow_eval_type_handlers` (all 45 pass), including the zero-arg, keyword-arg, multi-arg, concurrency, and `StopIteration` cases. - New `test_concurrency_level_multi_and_zero_arg`, which covers the thread-pool path with multi-arg and zero-arg UDFs. It passes both before and after this change. - Benchmarked with `ArrowBatchedUDFTimeBench` as described above. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 5.5) This pull request and its description were written by Isaac. -- 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]
