FelixYBW opened a new pull request, #59062: URL: https://github.com/apache/spark/pull/59062
### What changes were proposed in this pull request? This is a follow-up of SPARK-56350 (#55120) for Arrow-optimized Python UDTFs. When the child of `ArrowEvalPythonUDTFExec` supports columnar output and all the UDTF arguments are columns of the child, the operator now reads the columnar child directly. The argument columns are serialized to the Python worker from their Arrow vectors through `ColumnarArrowPythonInput`, which skips the `ColumnarToRowExec` -> `ArrowWriter` round-trip. Batches that are not Arrow-backed, or whose vectors are not compatible with the declared argument schema, fall back to row-by-row writing, as for UDFs. A UDTF returns any number of rows per input row, plus the rows of `terminate()`, so its output is still joined row by row with the input rows. Only the child columns in the output (`requiredChildOutput`) are buffered. - `ArrowEvalPythonUDTFExec` declares `supportsColumnar = child.supportsColumnar` (when enabled and all the arguments are child columns) and `supportsRowBased = true`, like `ArrowEvalPythonExec`. `doExecuteColumnar()` converts the joined rows with `RowToColumnarEvaluatorFactory`. - `ArrowPythonUDTFRunner` now extends a generic `BaseArrowPythonUDTFRunner[IN]`; its constructor is unchanged. The new `ColumnarArrowPythonUDTFRunner` mixes in `ColumnarArrowPythonInput`. - `EvalPythonUDTFExec`: the argument flattening and the output/input join are extracted into helpers (`flattenArguments`, `argumentSchema`, `joinWithInput`, `UDTFInputBuffer`) shared by the row-based and the columnar input paths. There is no behavior change for the row-based path. - UDTF arguments that are not child columns (literals, expressions) keep using the row-based path. - New conf `spark.sql.execution.arrow.pythonUDTF.columnarInput.enabled` (default `true`). ### Why are the changes needed? When an operator produces Arrow-backed `ColumnarBatch` (e.g., DataSource V2 connectors reading columnar formats into Arrow vectors), Arrow Python UDTFs convert the data to rows with `ColumnarToRowExec` and then back to Arrow with `ArrowWriter` to send it to the Python worker. Since the data is already in Arrow format, this round-trip is unnecessary. SPARK-56350 removed it for Arrow Python UDFs; this PR does the same for Arrow Python UDTFs. ### Does this PR introduce _any_ user-facing change? No, except for the new configuration `spark.sql.execution.arrow.pythonUDTF.columnarInput.enabled`, which can be used to disable the new path. ### How was this patch tested? New `ArrowColumnarPythonUDTFSuite`, reading from the Arrow-backed test DataSource V2 source of SPARK-56350: - There is no `ColumnarToRowExec` between the scan and `ArrowEvalPythonUDTFExec`. - The results are the same as the row-based path (conf disabled), with 0 to 4 output rows per input row, null values, `terminate()` rows and multiple partitions. - A pruned child output (child columns only used as UDTF arguments) gives the same results. - UDTF arguments that are not child columns use the row-based path. Existing `PythonUDTFSuite` and `ArrowColumnarPythonUDFSuite` pass. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (claude-opus-5-5) 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- 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]
