Skip to content

Commit 49a8389

Browse files
Yicong-HuangYicong-Huang
committed
[SPARK-59911][PYTHON] Migrate SQL_SCALAR_PANDAS_UDF to a pandas eval type handler
### 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` adds `PandasScalarUDFHandler(BatchEvalTypeHandler)` for `SQL_SCALAR_PANDAS_UDF`, 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. It calls the Arrow<->pandas conversions in `pyspark.sql.conversion` (`ArrowToPandasConversion.to_pandas` / `PandasToArrowConversion.from_pandas`) directly with the runner_conf-derived parameters, exactly as `ArrowScalarUDFHandler` calls `ArrowBatchTransformer` directly -- no pandas-specific intermediate base and no conversion wrappers, so the 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 ArrowScalarUDFHandler (SQL_SCALAR_ARROW_UDF, shipped) 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). One supporting fix in `pyspark/sql/conversion.py`: `ArrowToPandasConversion.to_pandas` declared `timezone: str`, but every sibling conversion method (and the `_convert_array` it delegates to) types it `Optional[str]`, and the worker passes `runner_conf.timezone` which is `Optional[str]`. The handler surfaces this (it is type-checked, unlike `worker.py`), so the annotation is corrected to `Optional[str]`. Annotation-only; no runtime change. 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 -- it is the previous `read_udfs` mapper body moved into the handler. 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. The handler passes `prefers_large_types=runner_conf.use_large_var_types` inline, matching the current mapper. Note the pre-existing inconsistencies in the other (untouched) pandas branches -- `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 follow this pattern Confirming the pattern generalizes to the Grouped category. `SQL_GROUPED_MAP_PANDAS_UDF` is a grouped type (input `Iterator[Iterator[pa.RecordBatch]]`, `ArrowStreamGroupSerializer`), so it extends the grouped category base directly and calls the same conversions directly -- no pandas-specific base: class PandasGroupedMapUDFHandler(GroupedEvalTypeHandler["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 = ArrowToPandasConversion.to_pandas(table, timezone=rc.timezone, # prefer_int_ext_dtype=rc.prefer_int_ext_dtype) # build the value DataFrame from value_offsets, call the UDF (with the key # tuple when num_udf_args == 2), verify the pandas result, then # yield PandasToArrowConversion.from_pandas([result], output_schema, ...) The current `GROUPED_MAP_PANDAS_UDF` mapper already makes exactly those two calls, so moving it into a handler is the same mechanical move as this PR -- only the category base (`Grouped` vs `Batch`), the per-group loop, key handling, the pandas result verification, and the output-batch resizing wrapper differ. If a second pandas handler ends up repeating the same runner_conf-to-parameter argument list, that is the point to factor a shared helper (as the Arrow side keeps shared verifiers in `verification.py`) -- deferred here since there is a single caller. ### 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. ### 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` (built on a real `RunnerConf`): 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 that the handler forwards `preferIntExtensionDtype` to the input conversion and `useLargeVarTypes` to the output conversion. - 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. Closes #59171 from Yicong-Huang/pandas-handler-base. Authored-by: Yicong Huang <17627829+Yicong-Huang@users.noreply.github.com> Signed-off-by: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com>
1 parent 9ca6037 commit 49a8389

6 files changed

Lines changed: 293 additions & 78 deletions

File tree

‎dev/sparktestsupport/modules.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -665,6 +665,7 @@ def __hash__(self):
665665
# unittests
666666
"pyspark.eval_handlers.tests.test_arrow_eval_type_handlers",
667667
"pyspark.eval_handlers.tests.test_base_eval_type_handlers",
668+
"pyspark.eval_handlers.tests.test_pandas_eval_type_handlers",
668669
"pyspark.sql.tests.test_artifact",
669670
"pyspark.sql.tests.test_catalog",
670671
"pyspark.sql.tests.test_column",

‎python/pyspark/eval_handlers/__init__.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@
2424
``read_udfs`` looks up via ``get_eval_type_handler``. Importing this package
2525
imports the concrete handler submodules so they register.
2626
27-
``_arrow`` imports pyarrow lazily, so it registers its handlers without pyarrow
28-
installed and defers the pyarrow check to when a handler runs.
27+
``_arrow`` and ``_pandas`` import pyarrow (and pandas) lazily, so they register
28+
their handlers without those installed and defer the checks to when a handler runs.
2929
"""
3030

31-
from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import
31+
from pyspark.eval_handlers import _arrow, _pandas # noqa: F401 # registers handlers on import

‎python/pyspark/eval_handlers/_pandas.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one or more
3+
# contributor license agreements. See the NOTICE file distributed with
4+
# this work for additional information regarding copyright ownership.
5+
# The ASF licenses this file to You under the Apache License, Version 2.0
6+
# (the "License"); you may not use this file except in compliance with
7+
# the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
"""Handlers for the pandas UDF eval types (the UDF exchanges ``pd.Series`` /
19+
``pd.DataFrame`` values, converted to and from Arrow via the conversion utilities
20+
in ``pyspark.sql.conversion``, the pandas counterpart of ``ArrowBatchTransformer``).
21+
22+
pandas and pyarrow are imported lazily (inside ``run`` and the type-checking block)
23+
so the module stays importable and its handlers register without either installed.
24+
Each handler calls ``require_minimum_pandas_version`` and
25+
``require_minimum_pyarrow_version`` in ``__init__`` so a missing or too-old
26+
dependency surfaces a clear error when the handler runs.
27+
"""
28+
29+
from __future__ import annotations
30+
31+
from collections.abc import Iterator
32+
from typing import TYPE_CHECKING, Any
33+
34+
from pyspark.errors import PySparkTypeError, PySparkValueError
35+
from pyspark.eval_handlers._base import BatchEvalTypeHandler
36+
from pyspark.eval_handlers.verification import verify_result_row_count
37+
from pyspark.sql.conversion import ArrowToPandasConversion, PandasToArrowConversion
38+
from pyspark.sql.pandas.utils import (
39+
require_minimum_pandas_version,
40+
require_minimum_pyarrow_version,
41+
)
42+
from pyspark.sql.types import StructField, StructType
43+
from pyspark.util import PythonEvalType
44+
45+
if TYPE_CHECKING:
46+
import pyarrow as pa
47+
48+
from pyspark.worker_util import EvalConf, RunnerConf
49+
50+
51+
class PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
52+
"""SQL_SCALAR_PANDAS_UDF: convert each input RecordBatch to pandas Series
53+
(struct columns become DataFrames), invoke each UDF once per batch, check the
54+
row count, and convert the pandas results back to one RecordBatch. The pandas
55+
counterpart of ArrowScalarUDFHandler, calling the Arrow<->pandas conversions in
56+
``pyspark.sql.conversion`` directly with the runner_conf-derived parameters."""
57+
58+
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_UDF
59+
60+
def __init__(
61+
self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
62+
) -> None:
63+
require_minimum_pandas_version()
64+
require_minimum_pyarrow_version()
65+
super().__init__(udfs, runner_conf, eval_conf)
66+
self._return_schema = StructType(
67+
[StructField("_%d" % i, rt) for i, (_, _, _, rt) in enumerate(udfs)]
68+
)
69+
70+
def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]":
71+
import pandas as pd
72+
73+
runner_conf = self._runner_conf
74+
for input_batch in data:
75+
num_rows = input_batch.num_rows
76+
77+
# Input: Arrow -> pandas Series (struct columns become DataFrames).
78+
pandas_columns = ArrowToPandasConversion.to_pandas(
79+
input_batch,
80+
timezone=runner_conf.timezone,
81+
struct_in_pandas="dict",
82+
ndarray_as_list=False,
83+
prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
84+
df_for_struct=True,
85+
)
86+
87+
# Process: evaluate each UDF column-wise on pandas Series.
88+
results = []
89+
for udf_func, args_offsets, kwargs_offsets, return_type in self._udfs:
90+
result = udf_func(
91+
*[pandas_columns[o] for o in args_offsets],
92+
**{k: pandas_columns[v] for k, v in kwargs_offsets.items()},
93+
)
94+
if not hasattr(result, "__len__"):
95+
pd_type = (
96+
"pandas.DataFrame"
97+
if isinstance(return_type, StructType)
98+
else "pandas.Series"
99+
)
100+
raise PySparkTypeError(
101+
errorClass="UDF_RETURN_TYPE",
102+
messageParameters={"expected": pd_type, "actual": type(result).__name__},
103+
)
104+
verify_result_row_count(len(result), num_rows)
105+
# struct_in_pandas="dict": UDF must return a DataFrame for struct types.
106+
if isinstance(return_type, StructType) and not isinstance(result, pd.DataFrame):
107+
raise PySparkValueError(
108+
"Invalid return type. Please make sure that the UDF returns a "
109+
"pandas.DataFrame when the specified return type is StructType."
110+
)
111+
results.append(result)
112+
113+
# Output: pandas -> Arrow.
114+
yield PandasToArrowConversion.from_pandas(
115+
results,
116+
self._return_schema,
117+
timezone=runner_conf.timezone,
118+
safecheck=runner_conf.safecheck,
119+
arrow_cast=True,
120+
prefers_large_types=runner_conf.use_large_var_types,
121+
assign_cols_by_name=runner_conf.assign_cols_by_name,
122+
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
123+
)

‎python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one or more
3+
# contributor license agreements. See the NOTICE file distributed with
4+
# this work for additional information regarding copyright ownership.
5+
# The ASF licenses this file to You under the Apache License, Version 2.0
6+
# (the "License"); you may not use this file except in compliance with
7+
# the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
"""Tests for the pandas eval type handlers (``_pandas``).
19+
20+
The handler tests build handler input in the same wire format the serializers
21+
produce -- a flat RecordBatch, one column per UDF argument -- and assert on the
22+
output batches via ``run(0, <input>)``. Handlers are built on a real ``RunnerConf``
23+
(from a plain conf dict), so the tests exercise the actual conf keys the handler
24+
reads; the forwarding tests pass a single conf key to confirm it reaches the
25+
Arrow<->pandas conversion.
26+
"""
27+
28+
import os
29+
import unittest
30+
from unittest.mock import patch
31+
32+
from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
33+
from pyspark.eval_handlers._base import get_eval_type_handler
34+
from pyspark.sql.types import LongType, StringType, StructField, StructType
35+
from pyspark.testing.utils import (
36+
have_pandas,
37+
have_pyarrow,
38+
pandas_requirement_message,
39+
pyarrow_requirement_message,
40+
)
41+
from pyspark.util import PythonEvalType
42+
43+
with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
44+
from pyspark.worker_util import RunnerConf
45+
46+
if have_pandas and have_pyarrow:
47+
import pandas as pd
48+
import pyarrow as pa
49+
50+
from pyspark.eval_handlers._pandas import PandasScalarUDFHandler
51+
52+
_missing_message = pandas_requirement_message or pyarrow_requirement_message
53+
54+
55+
def _batch(**columns):
56+
"""A RecordBatch of int64 columns, one per ``name=values`` kwarg."""
57+
return pa.RecordBatch.from_arrays(
58+
[pa.array(values, type=pa.int64()) for values in columns.values()],
59+
list(columns),
60+
)
61+
62+
63+
def _udf(func, return_type=None, args=(0,), kwargs=None):
64+
"""A scalar pandas UDF tuple ``(func, args_offsets, kwargs_offsets, return_type)``."""
65+
return (func, list(args), dict(kwargs or {}), return_type or LongType())
66+
67+
68+
def _handler(*udfs, runner_conf=None):
69+
"""Build a PandasScalarUDFHandler from one or more ``_udf`` tuples."""
70+
return PandasScalarUDFHandler(
71+
udfs=list(udfs), runner_conf=runner_conf or RunnerConf({}), eval_conf=None
72+
)
73+
74+
75+
@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
76+
class PandasEvalTypeHandlerRegistrationTests(unittest.TestCase):
77+
def test_pandas_eval_types_are_registered(self):
78+
# The migrated pandas eval type dispatches to its handler by lookup.
79+
self.assertIs(
80+
get_eval_type_handler(PythonEvalType.SQL_SCALAR_PANDAS_UDF),
81+
PandasScalarUDFHandler,
82+
)
83+
84+
85+
@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
86+
class PandasScalarUDFHandlerTests(unittest.TestCase):
87+
def test_invokes_udf_per_batch(self):
88+
handler = _handler(_udf(lambda s: s + 1))
89+
out = list(handler.run(0, iter([_batch(a=[1, 2, 3])])))
90+
self.assertEqual([b.column(0).to_pylist() for b in out], [[2, 3, 4]])
91+
92+
def test_coerces_output_to_return_type(self):
93+
# The UDF returns int32, but the declared return type is LongType (int64).
94+
handler = _handler(_udf(lambda s: (s + 1).astype("int32")))
95+
out = list(handler.run(0, iter([_batch(a=[10, 20])])))
96+
self.assertEqual(out[0].schema.field(0).type, pa.int64())
97+
self.assertEqual(out[0].column(0).to_pylist(), [11, 21])
98+
99+
def test_multiple_udfs_produce_one_column_each(self):
100+
handler = _handler(
101+
_udf(lambda s: s + 1, args=(0,)),
102+
_udf(lambda s: s * 2, args=(1,)),
103+
)
104+
out = list(handler.run(0, iter([_batch(a=[1, 2], b=[10, 20])])))
105+
self.assertEqual(out[0].column(0).to_pylist(), [2, 3])
106+
self.assertEqual(out[0].column(1).to_pylist(), [20, 40])
107+
108+
def test_passes_arg_by_keyword_offset(self):
109+
# Validates the inline args/kwargs offset handling: y is bound by keyword.
110+
handler = _handler(_udf(lambda x, y: x - y, args=(0,), kwargs={"y": 1}))
111+
out = list(handler.run(0, iter([_batch(a=[10, 20], b=[3, 5])])))
112+
self.assertEqual(out[0].column(0).to_pylist(), [7, 15])
113+
114+
def test_struct_return_type_takes_dataframe(self):
115+
# struct_in_pandas="dict" + df_for_struct=True: a struct return is a DataFrame.
116+
struct_type = StructType([StructField("x", LongType()), StructField("y", LongType())])
117+
handler = _handler(
118+
_udf(lambda s: pd.DataFrame({"x": s, "y": s * 10}), return_type=struct_type)
119+
)
120+
out = list(handler.run(0, iter([_batch(a=[1, 2])])))
121+
self.assertEqual(
122+
out[0].column(0).to_pylist(),
123+
[{"x": 1, "y": 10}, {"x": 2, "y": 20}],
124+
)
125+
126+
def test_rejects_non_sized_result(self):
127+
handler = _handler(_udf(lambda s: 42))
128+
with self.assertRaises(PySparkTypeError):
129+
list(handler.run(0, iter([_batch(a=[1, 2])])))
130+
131+
def test_rejects_row_count_mismatch(self):
132+
handler = _handler(_udf(lambda s: s.head(1)))
133+
with self.assertRaises(PySparkRuntimeError):
134+
list(handler.run(0, iter([_batch(a=[1, 2, 3])])))
135+
136+
def test_requires_dataframe_for_struct_return(self):
137+
struct_type = StructType([StructField("x", LongType())])
138+
handler = _handler(_udf(lambda s: s, return_type=struct_type))
139+
with self.assertRaises(PySparkValueError):
140+
list(handler.run(0, iter([_batch(a=[1, 2])])))
141+
142+
def test_forwards_prefer_int_ext_dtype_to_input_conversion(self):
143+
# preferIntExtensionDtype flows to the Arrow->pandas conversion: the UDF sees
144+
# a nullable "Int64" input instead of the default "int64".
145+
dtype_name = _udf(lambda s: pd.Series([str(s.dtype)] * len(s)), return_type=StringType())
146+
conf = RunnerConf({"spark.sql.execution.pythonUDF.pandas.preferIntExtensionDtype": "true"})
147+
out = list(_handler(dtype_name, runner_conf=conf).run(0, iter([_batch(a=[1, 2])])))
148+
self.assertEqual(out[0].column(0).to_pylist(), ["Int64", "Int64"])
149+
150+
def test_forwards_use_large_var_types_to_output_conversion(self):
151+
# useLargeVarTypes flows to the pandas->Arrow conversion: a string result
152+
# becomes large_string rather than string.
153+
to_str = _udf(lambda s: s.astype(str), return_type=StringType())
154+
conf = RunnerConf({"spark.sql.execution.arrow.useLargeVarTypes": "true"})
155+
out = list(_handler(to_str, runner_conf=conf).run(0, iter([_batch(a=[1, 2])])))
156+
self.assertTrue(pa.types.is_large_string(out[0].schema.field(0).type))
157+
158+
159+
if __name__ == "__main__":
160+
from pyspark.testing import main
161+
162+
main()

‎python/pyspark/sql/conversion.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1824,7 +1824,7 @@ class ArrowToPandasConversion:
18241824
def to_pandas(
18251825
cls,
18261826
batch: Union["pa.RecordBatch", "pa.Table"],
1827-
timezone: str,
1827+
timezone: Optional[str],
18281828
schema: Optional["StructType"] = None,
18291829
struct_in_pandas: str = "dict",
18301830
ndarray_as_list: bool = False,
@@ -1838,8 +1838,9 @@ def to_pandas(
18381838
----------
18391839
batch : pa.RecordBatch or pa.Table
18401840
The Arrow RecordBatch or Table to convert.
1841-
timezone : str
1842-
Timezone for timestamp conversion.
1841+
timezone : str or None
1842+
Timezone for timestamp conversion. Required if the data contains
1843+
timestamp types.
18431844
schema : StructType, optional
18441845
Spark schema for type conversion. If None, types are inferred from Arrow.
18451846
struct_in_pandas : str

0 commit comments

Comments
 (0)