Yicong-Huang opened a new pull request, #59171:
URL: https://github.com/apache/spark/pull/59171
### What changes were proposed in this pull request?
This is the first step of migrating the pandas UDF eval types onto the
eval-handler framework (umbrella SPARK-59415), following the Arrow handler work
already shipped in `pyspark/eval_handlers/_arrow.py`.
New file `python/pyspark/eval_handlers/_pandas.py`:
- Two module-level functions bind the repeated Arrow<->pandas conversion
parameters from `runner_conf` once, so every pandas category base can reuse
them without duplicating the runner_conf-to-parameter mapping:
- `to_pandas(batch, runner_conf, *, schema, struct_in_pandas,
ndarray_as_list, df_for_struct)` binds `timezone` and `prefer_int_ext_dtype`;
the caller chooses the per-branch knobs.
- `from_pandas(data, schema, runner_conf, *, prefers_large_types=None)`
binds `timezone`, `safecheck`, `arrow_cast=True`, `assign_cols_by_name` and
int-to-decimal coercion; `prefers_large_types` defaults to
`runner_conf.use_large_var_types` but the caller may override it.
- `PandasBatchEvalTypeHandler(BatchEvalTypeHandler)` is a linear category
base (not a mixin, per the design decision) whose `_to_pandas` / `_from_pandas`
methods just delegate to the two module functions with `self._runner_conf`
bound.
- `PandasScalarUDFHandler(PandasBatchEvalTypeHandler)` handles
`SQL_SCALAR_PANDAS_UDF`, reading as the pandas counterpart of
`ArrowScalarUDFHandler`: convert each input RecordBatch to pandas Series
(struct columns become DataFrames via `df_for_struct=True`), invoke each UDF
once per batch, check the row count, and convert the pandas results back to one
RecordBatch.
The class hierarchy after this PR:
EvalTypeHandler
BatchEvalTypeHandler
ArrowScalarUDFHandler (SQL_SCALAR_ARROW_UDF, shipped)
PandasBatchEvalTypeHandler (new base)
PandasScalarUDFHandler (SQL_SCALAR_PANDAS_UDF, new)
The `SQL_SCALAR_PANDAS_UDF` branch is removed from `worker.py::read_udfs`
(both the serializer-selection entry and the mapper block); `read_udfs` now
dispatches it through `get_eval_type_handler`, exactly as the migrated Arrow
types already do. All other pandas branches are left untouched.
`read_single_udf` still handles `SQL_SCALAR_PANDAS_UDF` (the handler path calls
it).
This is a pure refactor with no behavior change: the handler issues the same
`ArrowToPandasConversion.to_pandas` / `PandasToArrowConversion.from_pandas`
calls with the same parameters, in the same validation order, raising the same
error types. The one structural change -- inlining the args/kwargs offsets
(`f(*pos, **kw)`) instead of pre-combining them through `wrap_kwargs_support`
-- is equivalent (`wrap_kwargs_support` builds the same call) and matches how
`ArrowScalarUDFHandler` already does it. The handler serializer is
`ArrowStreamSerializer(write_start_stream=True)`, the same serializer the
legacy path selected for this eval type.
Two pre-existing inconsistencies across the pandas branches were noted but
deliberately preserved (a base that binds `prefers_large_types` from
`runner_conf.use_large_var_types` with a caller override reproduces all of
them): `GROUPED_AGG_PANDAS_ITER` pins `prefers_large_types=False`, and the
stateful branches omit it (use the default). If these turn out to be bugs they
should be fixed in a separate ticket, not here.
#### Sketch: how SQL_GROUPED_MAP_PANDAS_UDF(201) would subclass this base
Confirming the base generalizes to the Grouped category before we commit to
the pattern. `SQL_GROUPED_MAP_PANDAS_UDF` is a grouped type (input
`Iterator[Iterator[pa.RecordBatch]]`, `ArrowStreamGroupSerializer`), so it
needs the grouped counterpart of `PandasBatchEvalTypeHandler`:
class PandasGroupedEvalTypeHandler(GroupedEvalTypeHandler[OutputBatch]):
# identical two-line _to_pandas / _from_pandas delegates as the
batch base,
# reusing the SAME module-level to_pandas / from_pandas functions.
class
PandasGroupedMapUDFHandler(PandasGroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF
# __init__: parse key/value offsets (extract_key_value_indexes),
build
# output_schema = StructType([StructField("_0", return_type)]).
# run(split_index, data): for each group, materialize its batches
into one
# pa.Table, then
# all_series = self._to_pandas(table) # defaults:
df_for_struct=False
# build the value DataFrame from value_offsets, call the UDF (with
the key
# tuple when num_udf_args == 2), verify_pandas_result, then
# yield self._from_pandas([result], self._output_schema) #
default prefers_large_types
The current `GROUPED_MAP_PANDAS_UDF` mapper already calls
`ArrowToPandasConversion.to_pandas(table, timezone=...,
prefer_int_ext_dtype=...)` and `PandasToArrowConversion.from_pandas([result],
output_schema, timezone=..., safecheck=..., arrow_cast=True,
prefers_large_types=use_large_var_types, assign_cols_by_name=...,
int_to_decimal_coercion_enabled=...)` -- exactly the parameter set the two
module functions bind, and it uses `to_pandas`'s defaults
(`df_for_struct=False`) plus the default `prefers_large_types`. So the same two
functions serve both eval types with no duplication; only the category base
(`Grouped` vs `Batch`), the per-group loop, key handling,
`verify_pandas_result`, and the output-batch resizing wrapper differ. The
delegate methods are re-declared per category base rather than shared through a
mixin, per the linear-base design decision.
### Why are the changes needed?
Consolidating each Arrow/Pandas UDF eval type behind an `EvalTypeHandler`
keeps the per-type logic together and shrinks the long `read_udfs` dispatch.
This PR extends the pattern from the Arrow types to the first pandas type and
establishes the pandas conversion-binding base the remaining pandas types will
build on.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- New unit tests
`python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py`:
registration, per-batch invocation, output type coercion, multiple UDFs,
keyword-offset binding, struct return (DataFrame), the three error paths
(non-sized result, row-count mismatch, struct return that is not a DataFrame),
and the module functions' runner_conf binding plus the `prefers_large_types`
override.
- The existing per-eval-type suite
`pyspark.sql.tests.pandas.test_pandas_udf_scalar` (the behavior gate for
`SQL_SCALAR_PANDAS_UDF`) and the ASV per-eval-type benchmark (regression gate).
### Was this patch authored or co-authored using generative AI tooling?
Yes, authored with Claude (Opus 4.8) via Claude Code.
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]
