Yicong-Huang opened a new pull request, #59211:
URL: https://github.com/apache/spark/pull/59211
### What changes were proposed in this pull request?
This continues migrating the pandas UDF eval types onto the eval-handler
framework (umbrella SPARK-59415), following `PandasScalarUDFHandler`
(SPARK-59911) and the Arrow handler work. It moves two Batch-category iterator
eval types off the `worker.py::read_udfs` if/elif chain onto handler classes in
`python/pyspark/eval_handlers/_pandas.py`:
- `PandasScalarIterUDFHandler` for `SQL_SCALAR_PANDAS_ITER_UDF` -- the
pandas counterpart of `ArrowScalarIterUDFHandler`: the UDF receives an iterator
of its argument columns (struct columns as DataFrames), yields an iterator of
pandas Series/DataFrames, and the handler enforces the declared element type,
verifies the total row count matches the input, and converts each result back
to one RecordBatch.
- `PandasMapUDFHandler` for `SQL_MAP_PANDAS_ITER_UDF` (mapInPandas) -- the
pandas counterpart of `ArrowMapUDFHandler`: the UDF receives the input
RecordBatch stream as an iterator of pandas DataFrames (one per batch, the
single struct column expanded) and yields pandas DataFrames/Series, converted
back to one RecordBatch each.
Both extend `BatchEvalTypeHandler` directly and call the Arrow<->pandas
conversions in `pyspark.sql.conversion` (`ArrowToPandasConversion.to_pandas` /
`PandasToArrowConversion.from_pandas`) inline with the runner_conf-derived
parameters, exactly as `PandasScalarUDFHandler` does -- no pandas-specific
intermediate base and no conversion wrappers, so each handler reads as a
straight pandas analog of the Arrow one.
The class hierarchy after this PR (pandas mirrors Arrow -- both extend the
category base directly):
```
EvalTypeHandler
BatchEvalTypeHandler
ArrowScalarIterUDFHandler (SQL_SCALAR_ARROW_ITER_UDF, shipped)
ArrowMapUDFHandler (SQL_MAP_ARROW_ITER_UDF, shipped)
PandasScalarUDFHandler (SQL_SCALAR_PANDAS_UDF, shipped)
PandasScalarIterUDFHandler (SQL_SCALAR_PANDAS_ITER_UDF, new)
PandasMapUDFHandler (SQL_MAP_PANDAS_ITER_UDF, new)
```
The `SQL_SCALAR_PANDAS_ITER_UDF` and `SQL_MAP_PANDAS_ITER_UDF` branches are
removed from `read_udfs` (both the serializer-selection entries and the mapper
blocks); `read_udfs` now dispatches them through `get_eval_type_handler`,
exactly as the migrated Arrow types already do. All other pandas branches are
left untouched, and `read_single_udf` still handles both types (the handler
path calls it).
`verify_pandas_result` and its helper `_verify_column_schema` move from
`worker.py` into `pyspark/eval_handlers/verification.py`, where the other
result-verification helpers already live, so the handlers can call them without
importing `worker.py`. `worker.py` imports them back, so the remaining pandas
branches and `verify_arrow_result` are unchanged.
### Why are the changes needed?
This is a step in the umbrella SPARK-59415 effort to replace the long
`read_udfs` if/elif chain with a per-eval-type handler registry, so each eval
type's execution logic lives in one cohesive, directly-testable class rather
than a `worker.py` branch. Migrating these two iterator types keeps the pandas
side moving toward parity with the already-migrated Arrow handlers.
### Does this PR introduce _any_ user-facing change?
No. This is an internal refactor of the PySpark worker's dispatch; the UDF
API, the on-the-wire Arrow format, and all observable behavior (including the
iterator row-count / exhaustion / output-row-limit checks and the pandas result
verification) are preserved.
### How was this patch tested?
- New unit tests in
`python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py` for both
handlers (registration, batch-stream mapping, output coercion, row-count
mismatch, struct-return handling, and non-iterator / wrong-element-type
rejection), driving `run()` against the serializer wire format.
- Existing per-eval-type integration suites pass:
`pyspark.sql.tests.pandas.test_pandas_map` and
`pyspark.sql.tests.pandas.test_pandas_udf_scalar`.
- The ASV `bench_eval_type` benchmarks for both types
(`ScalarPandasIterUDF*Bench`, `MapPandasIterUDF*Bench`) continue to run the
full worker pipeline through the new handlers.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.8
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]
