You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/10/02 10:44:25 UTC
(spark) branch master updated: [SPARK-59911][PYTHON] Migrate SQL_SCALAR_PANDAS_UDF to a pandas eval type handler
This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 49a83896423a [SPARK-59911][PYTHON] Migrate SQL_SCALAR_PANDAS_UDF to a pandas eval type handler
49a83896423a is described below
commit 49a83896423a2bc6a1ff28f0ede8b94d6f7f6164
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Fri Oct 2 10:44:09 2026 +0000
[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` (`ArrowToPandasConve [...]
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_supp [...]
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 verifie [...]
### 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 <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
dev/sparktestsupport/modules.py | 1 +
python/pyspark/eval_handlers/__init__.py | 6 +-
python/pyspark/eval_handlers/_pandas.py | 123 ++++++++++++++++
.../tests/test_pandas_eval_type_handlers.py | 162 +++++++++++++++++++++
python/pyspark/sql/conversion.py | 7 +-
python/pyspark/worker.py | 72 ---------
6 files changed, 293 insertions(+), 78 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index ff921796e65d..cb964c62d437 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -665,6 +665,7 @@ pyspark_sql = Module(
# unittests
"pyspark.eval_handlers.tests.test_arrow_eval_type_handlers",
"pyspark.eval_handlers.tests.test_base_eval_type_handlers",
+ "pyspark.eval_handlers.tests.test_pandas_eval_type_handlers",
"pyspark.sql.tests.test_artifact",
"pyspark.sql.tests.test_catalog",
"pyspark.sql.tests.test_column",
diff --git a/python/pyspark/eval_handlers/__init__.py b/python/pyspark/eval_handlers/__init__.py
index 2f90d4b68b6a..63ab04aab9fd 100644
--- a/python/pyspark/eval_handlers/__init__.py
+++ b/python/pyspark/eval_handlers/__init__.py
@@ -24,8 +24,8 @@ declares its ``eval_type`` and self-registers at class definition, which
``read_udfs`` looks up via ``get_eval_type_handler``. Importing this package
imports the concrete handler submodules so they register.
-``_arrow`` imports pyarrow lazily, so it registers its handlers without pyarrow
-installed and defers the pyarrow check to when a handler runs.
+``_arrow`` and ``_pandas`` import pyarrow (and pandas) lazily, so they register
+their handlers without those installed and defer the checks to when a handler runs.
"""
-from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import
+from pyspark.eval_handlers import _arrow, _pandas # noqa: F401 # registers handlers on import
diff --git a/python/pyspark/eval_handlers/_pandas.py b/python/pyspark/eval_handlers/_pandas.py
new file mode 100644
index 000000000000..b5c0c7f83b3e
--- /dev/null
+++ b/python/pyspark/eval_handlers/_pandas.py
@@ -0,0 +1,123 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Handlers for the pandas UDF eval types (the UDF exchanges ``pd.Series`` /
+``pd.DataFrame`` values, converted to and from Arrow via the conversion utilities
+in ``pyspark.sql.conversion``, the pandas counterpart of ``ArrowBatchTransformer``).
+
+pandas and pyarrow are imported lazily (inside ``run`` and the type-checking block)
+so the module stays importable and its handlers register without either installed.
+Each handler calls ``require_minimum_pandas_version`` and
+``require_minimum_pyarrow_version`` in ``__init__`` so a missing or too-old
+dependency surfaces a clear error when the handler runs.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
+
+from pyspark.errors import PySparkTypeError, PySparkValueError
+from pyspark.eval_handlers._base import BatchEvalTypeHandler
+from pyspark.eval_handlers.verification import verify_result_row_count
+from pyspark.sql.conversion import ArrowToPandasConversion, PandasToArrowConversion
+from pyspark.sql.pandas.utils import (
+ require_minimum_pandas_version,
+ require_minimum_pyarrow_version,
+)
+from pyspark.sql.types import StructField, StructType
+from pyspark.util import PythonEvalType
+
+if TYPE_CHECKING:
+ import pyarrow as pa
+
+ from pyspark.worker_util import EvalConf, RunnerConf
+
+
+class PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_SCALAR_PANDAS_UDF: convert each input RecordBatch to pandas Series
+ (struct columns become DataFrames), invoke each UDF once per batch, check the
+ row count, and convert the pandas results back to one RecordBatch. The pandas
+ counterpart of ArrowScalarUDFHandler, calling the Arrow<->pandas conversions in
+ ``pyspark.sql.conversion`` directly with the runner_conf-derived parameters."""
+
+ eval_type = PythonEvalType.SQL_SCALAR_PANDAS_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
+ ) -> None:
+ require_minimum_pandas_version()
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ self._return_schema = StructType(
+ [StructField("_%d" % i, rt) for i, (_, _, _, rt) in enumerate(udfs)]
+ )
+
+ def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]":
+ import pandas as pd
+
+ runner_conf = self._runner_conf
+ for input_batch in data:
+ num_rows = input_batch.num_rows
+
+ # Input: Arrow -> pandas Series (struct columns become DataFrames).
+ pandas_columns = ArrowToPandasConversion.to_pandas(
+ input_batch,
+ timezone=runner_conf.timezone,
+ struct_in_pandas="dict",
+ ndarray_as_list=False,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ df_for_struct=True,
+ )
+
+ # Process: evaluate each UDF column-wise on pandas Series.
+ results = []
+ for udf_func, args_offsets, kwargs_offsets, return_type in self._udfs:
+ result = udf_func(
+ *[pandas_columns[o] for o in args_offsets],
+ **{k: pandas_columns[v] for k, v in kwargs_offsets.items()},
+ )
+ if not hasattr(result, "__len__"):
+ pd_type = (
+ "pandas.DataFrame"
+ if isinstance(return_type, StructType)
+ else "pandas.Series"
+ )
+ raise PySparkTypeError(
+ errorClass="UDF_RETURN_TYPE",
+ messageParameters={"expected": pd_type, "actual": type(result).__name__},
+ )
+ verify_result_row_count(len(result), num_rows)
+ # struct_in_pandas="dict": UDF must return a DataFrame for struct types.
+ if isinstance(return_type, StructType) and not isinstance(result, pd.DataFrame):
+ raise PySparkValueError(
+ "Invalid return type. Please make sure that the UDF returns a "
+ "pandas.DataFrame when the specified return type is StructType."
+ )
+ results.append(result)
+
+ # Output: pandas -> Arrow.
+ yield PandasToArrowConversion.from_pandas(
+ results,
+ self._return_schema,
+ timezone=runner_conf.timezone,
+ safecheck=runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=runner_conf.use_large_var_types,
+ assign_cols_by_name=runner_conf.assign_cols_by_name,
+ int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ )
diff --git a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
new file mode 100644
index 000000000000..3f2e4d516be6
--- /dev/null
+++ b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
@@ -0,0 +1,162 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for the pandas eval type handlers (``_pandas``).
+
+The handler tests build handler input in the same wire format the serializers
+produce -- a flat RecordBatch, one column per UDF argument -- and assert on the
+output batches via ``run(0, <input>)``. Handlers are built on a real ``RunnerConf``
+(from a plain conf dict), so the tests exercise the actual conf keys the handler
+reads; the forwarding tests pass a single conf key to confirm it reaches the
+Arrow<->pandas conversion.
+"""
+
+import os
+import unittest
+from unittest.mock import patch
+
+from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
+from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.sql.types import LongType, StringType, StructField, StructType
+from pyspark.testing.utils import (
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
+from pyspark.util import PythonEvalType
+
+with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
+ from pyspark.worker_util import RunnerConf
+
+if have_pandas and have_pyarrow:
+ import pandas as pd
+ import pyarrow as pa
+
+ from pyspark.eval_handlers._pandas import PandasScalarUDFHandler
+
+_missing_message = pandas_requirement_message or pyarrow_requirement_message
+
+
+def _batch(**columns):
+ """A RecordBatch of int64 columns, one per ``name=values`` kwarg."""
+ return pa.RecordBatch.from_arrays(
+ [pa.array(values, type=pa.int64()) for values in columns.values()],
+ list(columns),
+ )
+
+
+def _udf(func, return_type=None, args=(0,), kwargs=None):
+ """A scalar pandas UDF tuple ``(func, args_offsets, kwargs_offsets, return_type)``."""
+ return (func, list(args), dict(kwargs or {}), return_type or LongType())
+
+
+def _handler(*udfs, runner_conf=None):
+ """Build a PandasScalarUDFHandler from one or more ``_udf`` tuples."""
+ return PandasScalarUDFHandler(
+ udfs=list(udfs), runner_conf=runner_conf or RunnerConf({}), eval_conf=None
+ )
+
+
+@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
+class PandasEvalTypeHandlerRegistrationTests(unittest.TestCase):
+ def test_pandas_eval_types_are_registered(self):
+ # The migrated pandas eval type dispatches to its handler by lookup.
+ self.assertIs(
+ get_eval_type_handler(PythonEvalType.SQL_SCALAR_PANDAS_UDF),
+ PandasScalarUDFHandler,
+ )
+
+
+@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
+class PandasScalarUDFHandlerTests(unittest.TestCase):
+ def test_invokes_udf_per_batch(self):
+ handler = _handler(_udf(lambda s: s + 1))
+ out = list(handler.run(0, iter([_batch(a=[1, 2, 3])])))
+ self.assertEqual([b.column(0).to_pylist() for b in out], [[2, 3, 4]])
+
+ def test_coerces_output_to_return_type(self):
+ # The UDF returns int32, but the declared return type is LongType (int64).
+ handler = _handler(_udf(lambda s: (s + 1).astype("int32")))
+ out = list(handler.run(0, iter([_batch(a=[10, 20])])))
+ self.assertEqual(out[0].schema.field(0).type, pa.int64())
+ self.assertEqual(out[0].column(0).to_pylist(), [11, 21])
+
+ def test_multiple_udfs_produce_one_column_each(self):
+ handler = _handler(
+ _udf(lambda s: s + 1, args=(0,)),
+ _udf(lambda s: s * 2, args=(1,)),
+ )
+ out = list(handler.run(0, iter([_batch(a=[1, 2], b=[10, 20])])))
+ self.assertEqual(out[0].column(0).to_pylist(), [2, 3])
+ self.assertEqual(out[0].column(1).to_pylist(), [20, 40])
+
+ def test_passes_arg_by_keyword_offset(self):
+ # Validates the inline args/kwargs offset handling: y is bound by keyword.
+ handler = _handler(_udf(lambda x, y: x - y, args=(0,), kwargs={"y": 1}))
+ out = list(handler.run(0, iter([_batch(a=[10, 20], b=[3, 5])])))
+ self.assertEqual(out[0].column(0).to_pylist(), [7, 15])
+
+ def test_struct_return_type_takes_dataframe(self):
+ # struct_in_pandas="dict" + df_for_struct=True: a struct return is a DataFrame.
+ struct_type = StructType([StructField("x", LongType()), StructField("y", LongType())])
+ handler = _handler(
+ _udf(lambda s: pd.DataFrame({"x": s, "y": s * 10}), return_type=struct_type)
+ )
+ out = list(handler.run(0, iter([_batch(a=[1, 2])])))
+ self.assertEqual(
+ out[0].column(0).to_pylist(),
+ [{"x": 1, "y": 10}, {"x": 2, "y": 20}],
+ )
+
+ def test_rejects_non_sized_result(self):
+ handler = _handler(_udf(lambda s: 42))
+ with self.assertRaises(PySparkTypeError):
+ list(handler.run(0, iter([_batch(a=[1, 2])])))
+
+ def test_rejects_row_count_mismatch(self):
+ handler = _handler(_udf(lambda s: s.head(1)))
+ with self.assertRaises(PySparkRuntimeError):
+ list(handler.run(0, iter([_batch(a=[1, 2, 3])])))
+
+ def test_requires_dataframe_for_struct_return(self):
+ struct_type = StructType([StructField("x", LongType())])
+ handler = _handler(_udf(lambda s: s, return_type=struct_type))
+ with self.assertRaises(PySparkValueError):
+ list(handler.run(0, iter([_batch(a=[1, 2])])))
+
+ def test_forwards_prefer_int_ext_dtype_to_input_conversion(self):
+ # preferIntExtensionDtype flows to the Arrow->pandas conversion: the UDF sees
+ # a nullable "Int64" input instead of the default "int64".
+ dtype_name = _udf(lambda s: pd.Series([str(s.dtype)] * len(s)), return_type=StringType())
+ conf = RunnerConf({"spark.sql.execution.pythonUDF.pandas.preferIntExtensionDtype": "true"})
+ out = list(_handler(dtype_name, runner_conf=conf).run(0, iter([_batch(a=[1, 2])])))
+ self.assertEqual(out[0].column(0).to_pylist(), ["Int64", "Int64"])
+
+ def test_forwards_use_large_var_types_to_output_conversion(self):
+ # useLargeVarTypes flows to the pandas->Arrow conversion: a string result
+ # becomes large_string rather than string.
+ to_str = _udf(lambda s: s.astype(str), return_type=StringType())
+ conf = RunnerConf({"spark.sql.execution.arrow.useLargeVarTypes": "true"})
+ out = list(_handler(to_str, runner_conf=conf).run(0, iter([_batch(a=[1, 2])])))
+ self.assertTrue(pa.types.is_large_string(out[0].schema.field(0).type))
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py
index b0c819a69d7f..0711a110d147 100644
--- a/python/pyspark/sql/conversion.py
+++ b/python/pyspark/sql/conversion.py
@@ -1824,7 +1824,7 @@ class ArrowToPandasConversion:
def to_pandas(
cls,
batch: Union["pa.RecordBatch", "pa.Table"],
- timezone: str,
+ timezone: Optional[str],
schema: Optional["StructType"] = None,
struct_in_pandas: str = "dict",
ndarray_as_list: bool = False,
@@ -1838,8 +1838,9 @@ class ArrowToPandasConversion:
----------
batch : pa.RecordBatch or pa.Table
The Arrow RecordBatch or Table to convert.
- timezone : str
- Timezone for timestamp conversion.
+ timezone : str or None
+ Timezone for timestamp conversion. Required if the data contains
+ timestamp types.
schema : StructType, optional
Spark schema for type conversion. If None, types are inferred from Arrow.
struct_in_pandas : str
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index e92fa5d7afc0..ffcf12b22c29 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -1817,7 +1817,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_UDF,
PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
@@ -2946,77 +2945,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
return func, ser
- if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_UDF:
- import pandas as pd
- import pyarrow as pa
-
- # --- UDF preparation ---
- udf_infos = []
- for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in udfs:
- wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
- udf_func, udf_args_offsets, udf_kwargs_offsets
- )
- udf_infos.append((wrapped_func, args_kwargs_offsets, udf_return_type))
- return_schema = StructType(
- [StructField(f"_{i}", info[2]) for i, info in enumerate(udf_infos)]
- )
-
- def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
- for input_batch in data:
- num_rows = input_batch.num_rows
-
- # --- Input: Arrow -> pandas Series (struct columns become DataFrames) ---
- pandas_columns = ArrowToPandasConversion.to_pandas(
- input_batch,
- timezone=runner_conf.timezone,
- struct_in_pandas="dict",
- ndarray_as_list=False,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- df_for_struct=True,
- )
-
- # --- Process: evaluate each UDF column-wise on pandas Series ---
- results = []
- for udf_func, offsets, udf_return_type in udf_infos:
- result = udf_func(*[pandas_columns[o] for o in offsets])
- if not hasattr(result, "__len__"):
- pd_type = (
- "pandas.DataFrame"
- if isinstance(udf_return_type, StructType)
- else "pandas.Series"
- )
- raise PySparkTypeError(
- errorClass="UDF_RETURN_TYPE",
- messageParameters={
- "expected": pd_type,
- "actual": type(result).__name__,
- },
- )
- verify_result_row_count(len(result), num_rows)
- # struct_in_pandas="dict": UDF must return DataFrame for struct types
- if isinstance(udf_return_type, StructType) and not isinstance(
- result, pd.DataFrame
- ):
- raise PySparkValueError(
- "Invalid return type. Please make sure that the UDF returns a "
- "pandas.DataFrame when the specified return type is StructType."
- )
- results.append(result)
-
- # --- Output: pandas -> Arrow ---
- yield PandasToArrowConversion.from_pandas(
- results,
- return_schema,
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
- arrow_cast=True,
- prefers_large_types=runner_conf.use_large_var_types,
- assign_cols_by_name=runner_conf.assign_cols_by_name,
- int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
- )
-
- return func, ser
-
if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF:
import pandas as pd
import pyarrow as pa
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org