Yicong-Huang opened a new pull request, #58729: URL: https://github.com/apache/spark/pull/58729
### What changes were proposed in this pull request? This PR introduces an extensible, typed execution model for Arrow/Pandas UDF eval types in the PySpark worker, and migrates the first eval type (`SQL_SCALAR_ARROW_UDF`) onto it. It is the foundation for the umbrella refactor tracked in [SPARK-59415](https://issues.apache.org/jira/browse/SPARK-59415) (design: "1DD: Python UDF Eval Handler API Design", Design A). Every Arrow/Pandas UDF eval type runs the same three-stage pipeline in the worker: prepare the input, invoke user code, then validate and normalize the result. Historically each eval type spelled this out inline in `read_udfs`, which has grown into a single ~2,500-line central `if/elif` dispatcher. This PR makes that lifecycle explicit: - `EvalTypeHandler[InputBatch, OutputBatch]`: a generic interface with three abstract stages -- `pre_process` (argument extraction / conversions), `process` (user-code invocation only), `post_process` (schema enforcement, row-count checks) -- plus a `run(split_index, data)` runner that chains them lazily as generators. `run` matches the `func(split_index, data)` shape `read_udfs` already returns. - Three typed category bases that fix the input-stream element type and the default serializer: `BatchEvalTypeHandler` (`Iterator[pa.RecordBatch]`, `ArrowStreamSerializer`), `GroupedEvalTypeHandler` (`Iterator[GroupedBatch]`, `ArrowStreamGroupSerializer`), and `CoGroupedEvalTypeHandler` (`Iterator[CoGroupedBatch]`, `ArrowStreamCoGroupSerializer`). - Automatic registration: a concrete handler declares its `eval_type`, and `EvalTypeHandler.__init_subclass__` registers it in `_EVAL_TYPE_HANDLERS`. `read_udfs` looks the eval type up first and delegates to the handler, so a later eval type attaches as a self-contained subclass without editing any central branch. Duplicate registrations are rejected. - `ArrowScalarUDFHandler` migrates `SQL_SCALAR_ARROW_UDF` as the first handler; its old inline branch and serializer-tuple entry in `read_udfs` are removed. The remaining 28 Arrow/Pandas eval types are left on the existing `if/elif` path and will be migrated incrementally in follow-up PRs, each independently revertible. Class hierarchy: ``` EvalTypeHandler[InputBatch, OutputBatch] (ABC: pre_process / process / post_process / run) |- BatchEvalTypeHandler (Iterator[pa.RecordBatch]) | |- ArrowScalarUDFHandler (SQL_SCALAR_ARROW_UDF) |- GroupedEvalTypeHandler (Iterator[GroupedBatch]) |- CoGroupedEvalTypeHandler (Iterator[CoGroupedBatch]) ``` ### Why are the changes needed? `read_udfs` is a single central dispatcher whose implicit per-eval-type lifecycle is hard to read, extend, and test. Making the three stages explicit and letting eval types register themselves removes the central `if/elif` chain (for both function construction and serializer selection), keeps a strong per-category type contract, and lets each subsequent eval type be migrated and reverted on its own. ### Does this PR introduce _any_ user-facing change? No. This is an internal worker refactor. There is no change to any user-facing UDF API, the on-the-wire format, or JVM-side code, and the migrated `SQL_SCALAR_ARROW_UDF` path is behavior-identical. ### How was this patch tested? - New `pyspark.sql.tests.test_eval_type_handlers` unit suite (no JVM required): handler registration, category-base abstractness, non-registration of the abstract bases, per-category default serializers, `run` stage-chaining, duplicate-registration rejection, and end-to-end `ArrowScalarUDFHandler` output including schema coercion. - Existing `pyspark.sql.tests.arrow.test_arrow_udf_scalar` (55 tests + 20 subtests) and `pyspark.sql.tests.arrow.test_arrow_udf_typehints` (15 tests) pass unchanged against a local `-Phive` build, confirming zero behavior change on the migrated path. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Isaac 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]
