This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 45eb755f6a6e [SPARK-59959][PYTHON] Migrate SQL_SCALAR_PANDAS_ITER_UDF
and SQL_MAP_PANDAS_ITER_UDF to pandas eval type handlers
45eb755f6a6e is described below
commit 45eb755f6a6e2c02f57d716d439c078a6241ba48
Author: Yicong Huang <[email protected]>
AuthorDate: Wed Oct 7 22:51:31 2026 +0000
[SPARK-59959][PYTHON] Migrate SQL_SCALAR_PANDAS_ITER_UDF and
SQL_MAP_PANDAS_ITER_UDF to pandas eval type handlers
### 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.
Closes #59211 from Yicong-Huang/pandas-scalar-map-iter.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit 37d06912ffea5a1472babe712d72ed04352d53ef)
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/eval_handlers/_pandas.py | 169 ++++++++++++++--
.../tests/test_pandas_eval_type_handlers.py | 163 +++++++++++++++-
python/pyspark/eval_handlers/verification.py | 61 +++++-
python/pyspark/worker.py | 213 +--------------------
4 files changed, 377 insertions(+), 229 deletions(-)
diff --git a/python/pyspark/eval_handlers/_pandas.py
b/python/pyspark/eval_handlers/_pandas.py
index b5c0c7f83b3e..07302fd66f39 100644
--- a/python/pyspark/eval_handlers/_pandas.py
+++ b/python/pyspark/eval_handlers/_pandas.py
@@ -19,8 +19,8 @@
``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.
+pandas and pyarrow are imported lazily (inside ``__init__``/``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.
@@ -33,7 +33,14 @@ 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.eval_handlers.verification import (
+ verify_iter_result_row_count,
+ verify_iterator_exhausted,
+ verify_output_row_limit,
+ verify_pandas_result,
+ verify_result_row_count,
+ verify_return_type,
+)
from pyspark.sql.conversion import ArrowToPandasConversion,
PandasToArrowConversion
from pyspark.sql.pandas.utils import (
require_minimum_pandas_version,
@@ -70,17 +77,16 @@ class
PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
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,
+ timezone=self._runner_conf.timezone,
struct_in_pandas="dict",
ndarray_as_list=False,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
df_for_struct=True,
)
@@ -114,10 +120,151 @@ class
PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
yield PandasToArrowConversion.from_pandas(
results,
self._return_schema,
- timezone=runner_conf.timezone,
- safecheck=runner_conf.safecheck,
+ timezone=self._runner_conf.timezone,
+ safecheck=self._runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=self._runner_conf.use_large_var_types,
+ assign_cols_by_name=self._runner_conf.assign_cols_by_name,
+
int_to_decimal_coercion_enabled=self._runner_conf.int_to_decimal_coercion_enabled,
+ )
+
+
+class PandasScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_SCALAR_PANDAS_ITER_UDF: the single UDF receives an iterator of its
+ argument columns (each a pandas Series, struct columns as DataFrames) and
yields
+ an iterator of pandas Series/DataFrames; enforce the declared element
type, verify
+ the total row count matches the input, and convert each result back to one
+ RecordBatch. The pandas counterpart of ArrowScalarIterUDFHandler, calling
the
+ Arrow<->pandas conversions in ``pyspark.sql.conversion`` directly."""
+
+ eval_type = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf",
eval_conf: "EvalConf"
+ ) -> None:
+ import pandas as pd
+
+ require_minimum_pandas_version()
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ assert len(udfs) == 1, "One SCALAR_PANDAS_ITER UDF expected here."
+ self._udf_func, self._args_offsets, _, self._return_type = udfs[0]
+ self._return_schema = StructType([StructField("_0",
self._return_type)])
+ self._expected_iter_type = (
+ Iterator[pd.DataFrame]
+ if isinstance(self._return_type, StructType)
+ else Iterator[pd.Series]
+ )
+
+ def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") ->
"Iterator[pa.RecordBatch]":
+ num_input_rows = 0
+
+ def extract_args(batch: "pa.RecordBatch") -> Any:
+ nonlocal num_input_rows
+ # Input: Arrow -> pandas Series (struct columns become DataFrames).
+ pandas_columns = ArrowToPandasConversion.to_pandas(
+ batch,
+ timezone=self._runner_conf.timezone,
+ struct_in_pandas="dict",
+ ndarray_as_list=False,
+ prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
+ df_for_struct=True,
+ )
+ args = tuple(pandas_columns[o] for o in self._args_offsets)
+ num_input_rows += batch.num_rows
+ return args[0] if len(args) == 1 else args
+
+ # Extract args from input batches (streaming), then call the UDF and
verify the
+ # result type (iterator of pd.Series / pd.DataFrame).
+ args_iter = map(extract_args, data)
+ verified_iter = verify_return_type(self._udf_func(args_iter),
self._expected_iter_type)
+
+ def convert(result: Any) -> "pa.RecordBatch":
+ verify_pandas_result(
+ result, self._return_type, assign_cols_by_name=True,
truncate_return_schema=True
+ )
+ return PandasToArrowConversion.from_pandas(
+ [result],
+ self._return_schema,
+ timezone=self._runner_conf.timezone,
+ safecheck=self._runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=self._runner_conf.use_large_var_types,
+ assign_cols_by_name=self._runner_conf.assign_cols_by_name,
+
int_to_decimal_coercion_enabled=self._runner_conf.int_to_decimal_coercion_enabled,
+ )
+
+ # Row-limit check (fail-fast) then exact row-count match (final).
+ batches = verify_output_row_limit(map(convert, verified_iter), lambda:
num_input_rows)
+ yield from verify_iter_result_row_count(batches, lambda:
num_input_rows)
+
+ # Verify the input iterator was fully consumed.
+ verify_iterator_exhausted(args_iter)
+
+
+class PandasMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_MAP_PANDAS_ITER_UDF (mapInPandas): the single 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. The pandas counterpart of ArrowMapUDFHandler."""
+
+ eval_type = PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf",
eval_conf: "EvalConf"
+ ) -> None:
+ import pandas as pd
+
+ require_minimum_pandas_version()
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ assert len(udfs) == 1, "One MAP_PANDAS_ITER UDF expected here."
+ self._udf_func, _, _, self._return_type = udfs[0]
+ self._return_schema = StructType([StructField("_0",
self._return_type)])
+ self._expected_iter_type = (
+ Iterator[pd.DataFrame]
+ if isinstance(self._return_type, StructType)
+ else Iterator[pd.Series]
+ )
+
+ def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") ->
"Iterator[pa.RecordBatch]":
+ def dataframe_iter() -> "Iterator[Any]":
+ # Input batches have a single struct column (see
MapInBatchEvaluatorFactory);
+ # convert lazily so peakmem stays bounded by one batch.
+ for batch in data:
+ yield ArrowToPandasConversion.to_pandas(
+ batch,
+ timezone=self._runner_conf.timezone,
+
prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
+ df_for_struct=True,
+ )[0]
+
+ result = self._udf_func(dataframe_iter())
+ # The declared signature is Iterator[...], so a strict iterator is
required by
+ # default. With the legacy flag, accept any object Python can iterate
over -- via
+ # iter(...), which honors both __iter__ and the sequence protocol
(__getitem__) --
+ # by adapting it into an iterator before the shared element-type
verification.
+ if self._runner_conf.map_in_batch_legacy_accept_any_iterable and not
isinstance(
+ result, Iterator
+ ):
+ try:
+ result = iter(result)
+ except TypeError:
+ # Not iterable at all; leave it so verify_return_type below
raises the
+ # standard UDF_RETURN_TYPE error.
+ pass
+
+ for df in verify_return_type(result, self._expected_iter_type):
+ verify_pandas_result(
+ df, self._return_type, assign_cols_by_name=True,
truncate_return_schema=True
+ )
+ yield PandasToArrowConversion.from_pandas(
+ [df],
+ self._return_schema,
+ timezone=self._runner_conf.timezone,
+ safecheck=self._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,
+ prefers_large_types=self._runner_conf.use_large_var_types,
+ assign_cols_by_name=self._runner_conf.assign_cols_by_name,
+
int_to_decimal_coercion_enabled=self._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
index 3f2e4d516be6..befbd13b2284 100644
--- a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
@@ -47,7 +47,12 @@ if have_pandas and have_pyarrow:
import pandas as pd
import pyarrow as pa
- from pyspark.eval_handlers._pandas import PandasScalarUDFHandler
+ from pyspark.eval_handlers._pandas import (
+ PandasMapUDFHandler,
+ PandasScalarIterUDFHandler,
+ PandasScalarUDFHandler,
+ )
+ from pyspark.sql.conversion import ArrowBatchTransformer
_missing_message = pandas_requirement_message or pyarrow_requirement_message
@@ -72,14 +77,34 @@ def _handler(*udfs, runner_conf=None):
)
+def _scalar_iter_handler(udf, return_type=None, args=(0,), runner_conf=None):
+ """Build a PandasScalarIterUDFHandler whose one UDF reads the given arg
offsets."""
+ return PandasScalarIterUDFHandler(
+ udfs=[(udf, list(args), {}, return_type or LongType())],
+ runner_conf=runner_conf or RunnerConf({}),
+ eval_conf=None,
+ )
+
+
+def _map_handler(udf, return_type, runner_conf=None):
+ """Build a PandasMapUDFHandler from its ``(func, None, None,
return_type)`` UDF tuple."""
+ return PandasMapUDFHandler(
+ udfs=[(udf, None, None, return_type)],
+ 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,
- )
+ # Each migrated pandas eval type dispatches to its handler by lookup.
+ for eval_type, handler_cls in (
+ (PythonEvalType.SQL_SCALAR_PANDAS_UDF, PandasScalarUDFHandler),
+ (PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
PandasScalarIterUDFHandler),
+ (PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, PandasMapUDFHandler),
+ ):
+ self.assertIs(get_eval_type_handler(eval_type), handler_cls)
@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
@@ -156,6 +181,132 @@ class PandasScalarUDFHandlerTests(unittest.TestCase):
self.assertTrue(pa.types.is_large_string(out[0].schema.field(0).type))
[email protected](not (have_pandas and have_pyarrow), _missing_message)
+class PandasScalarIterUDFHandlerTests(unittest.TestCase):
+ def test_invokes_udf_over_batch_stream(self):
+ def add_one(series_iter):
+ for s in series_iter:
+ yield s + 1
+
+ handler = _scalar_iter_handler(add_one)
+ out = list(handler.run(0, iter([_batch(a=[1, 2]), _batch(a=[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 yields int32, but the declared return type is LongType
(int64).
+ def add_one(series_iter):
+ for s in series_iter:
+ yield (s + 1).astype("int32")
+
+ handler = _scalar_iter_handler(add_one)
+ 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_rejects_too_many_rows(self):
+ # Emitting more rows than consumed fails fast via the output row-limit
guard.
+ def too_many(series_iter):
+ for s in series_iter:
+ yield pd.Series(list(range(len(s) + 1)))
+
+ handler = _scalar_iter_handler(too_many)
+ with self.assertRaises(PySparkRuntimeError) as cm:
+ list(handler.run(0, iter([_batch(a=[1, 2])])))
+ self.assertEqual(cm.exception.getCondition(),
"OUTPUT_EXCEEDS_INPUT_ROWS")
+
+ def test_rejects_too_few_rows(self):
+ # Emitting fewer rows than consumed fails the final row-count equality
guard.
+ def too_few(series_iter):
+ for s in series_iter:
+ yield s.head(len(s) - 1)
+
+ handler = _scalar_iter_handler(too_few)
+ with self.assertRaises(PySparkRuntimeError) as cm:
+ list(handler.run(0, iter([_batch(a=[1, 2, 3])])))
+ self.assertEqual(cm.exception.getCondition(), "RESULT_ROWS_MISMATCH")
+
+ def test_rejects_unconsumed_input(self):
+ # A UDF that stops reading before the input stream is exhausted must
raise
+ # INPUT_NOT_FULLY_CONSUMED.
+ def first_only(series_iter):
+ yield next(series_iter)
+
+ handler = _scalar_iter_handler(first_only)
+ with self.assertRaises(PySparkRuntimeError) as cm:
+ list(handler.run(0, iter([_batch(a=[1, 2]), _batch(a=[3, 4])])))
+ self.assertEqual(cm.exception.getCondition(),
"INPUT_NOT_FULLY_CONSUMED")
+
+ def test_struct_return_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())])
+
+ def to_struct(series_iter):
+ for s in series_iter:
+ yield pd.DataFrame({"x": s, "y": s * 10})
+
+ handler = _scalar_iter_handler(to_struct, 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}],
+ )
+
+
[email protected](not (have_pandas and have_pyarrow), _missing_message)
+class PandasMapUDFHandlerTests(unittest.TestCase):
+ _return_type = StructType([StructField("v", LongType())])
+
+ def test_maps_batch_stream(self):
+ # mapInPandas expands the wire struct into a DataFrame and re-wraps
the output.
+ def double_v(df_iter):
+ for df in df_iter:
+ yield pd.DataFrame({"v": df["v"] * 2})
+
+ handler = _map_handler(double_v, self._return_type)
+ # mapInPandas sends a single struct column per batch; the handler
expands it to
+ # a DataFrame, so build the input in that wire format here.
+ struct_batch = ArrowBatchTransformer.wrap_struct(_batch(v=[1, 2, 3]))
+ out = list(handler.run(0, iter([struct_batch])))
+ self.assertEqual(out[0].column(0).field("v").to_pylist(), [2, 4, 6])
+
+ def test_legacy_accepts_any_iterable(self):
+ # With the legacy flag on (its default), a UDF may return any iterable
-- here a
+ # plain list of DataFrames -- which the handler adapts via iter()
before verifying.
+ def as_list(df_iter):
+ return [pd.DataFrame({"v": df["v"] * 2}) for df in df_iter]
+
+ handler = _map_handler(as_list, self._return_type)
+ struct_batch = ArrowBatchTransformer.wrap_struct(_batch(v=[1, 2, 3]))
+ out = list(handler.run(0, iter([struct_batch])))
+ self.assertEqual(out[0].column(0).field("v").to_pylist(), [2, 4, 6])
+
+ def test_rejects_non_iterator_result(self):
+ # With the legacy accept-any-iterable flag off, a UDF returning a
DataFrame (not
+ # an iterator of them) is rejected as a non-iterator before the input
stream is
+ # read, so the input batch shape is irrelevant.
+ strict = RunnerConf(
+
{"spark.sql.execution.pythonUDF.mapInBatch.legacy.acceptAnyIterable.enabled":
"false"}
+ )
+ handler = _map_handler(
+ lambda df_iter: pd.DataFrame({"v": [1]}), self._return_type,
runner_conf=strict
+ )
+ with self.assertRaises(PySparkTypeError) as cm:
+ list(handler.run(0, iter([_batch(v=[1, 2])])))
+ # The non-iterator branch reports the actual type ("DataFrame"); the
per-element
+ # branch (test_rejects_wrong_element_type) reports "iterator of ...".
+ self.assertEqual(cm.exception.getMessageParameters()["actual"],
"DataFrame")
+
+ def test_rejects_wrong_element_type(self):
+ # Each yielded element must be a DataFrame for a struct return type.
+ def yields_series(df_iter):
+ for _ in df_iter:
+ yield pd.Series([1, 2]) # a Series, not a DataFrame
+
+ handler = _map_handler(yields_series, self._return_type)
+ with self.assertRaises(PySparkTypeError):
+ list(handler.run(0, iter([_batch(v=[1, 2])])))
+
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git a/python/pyspark/eval_handlers/verification.py
b/python/pyspark/eval_handlers/verification.py
index 3b8399984dee..9f413e8dca0e 100644
--- a/python/pyspark/eval_handlers/verification.py
+++ b/python/pyspark/eval_handlers/verification.py
@@ -21,9 +21,13 @@ Used by the eval type handlers and by the worker's
``read_udfs``.
"""
from collections.abc import Iterator
-from typing import Any, Callable, Type, TypeVar, Union, cast, get_args,
get_origin
+from typing import TYPE_CHECKING, Any, Callable, Type, TypeVar, Union, cast,
get_args, get_origin
from pyspark.errors import PySparkRuntimeError, PySparkTypeError
+from pyspark.sql.types import DataType, StructType
+
+if TYPE_CHECKING:
+ import pandas as pd
T = TypeVar("T")
@@ -151,3 +155,58 @@ def verify_scalar_result(result: Any, num_rows: int) ->
Any:
)
verify_result_row_count(result_length, num_rows)
return result
+
+
+def _verify_column_schema(
+ actual_names: list, expected_names: list, *, assign_cols_by_name: bool
+) -> None:
+ """Check column names (by-name) or count (by-position) match the expected
schema."""
+ if assign_cols_by_name:
+ actual_set = set(actual_names)
+ expected_set = set(expected_names)
+ missing = sorted(expected_set.difference(actual_set))
+ extra = sorted(actual_set.difference(expected_set))
+ if missing or extra:
+ raise PySparkRuntimeError(
+ errorClass="RESULT_COLUMN_NAMES_MISMATCH",
+ messageParameters={
+ "missing": f" Missing: {', '.join(missing)}." if missing
else "",
+ "extra": f" Unexpected: {', '.join(extra)}." if extra else
"",
+ },
+ )
+ elif len(actual_names) != len(expected_names):
+ raise PySparkRuntimeError(
+ errorClass="RESULT_COLUMN_SCHEMA_MISMATCH",
+ messageParameters={
+ "expected": str(len(expected_names)),
+ "actual": str(len(actual_names)),
+ },
+ )
+
+
+def verify_pandas_result(
+ result: Union["pd.DataFrame", "pd.Series"],
+ return_type: DataType,
+ assign_cols_by_name: bool,
+ truncate_return_schema: bool,
+) -> None:
+ import pandas as pd
+
+ if not isinstance(return_type, StructType):
+ verify_return_type(result, pd.Series)
+ return
+
+ verify_return_type(result, pd.DataFrame)
+
+ # Skip schema check on a fully empty result (no rows and no columns).
+ if result.empty and len(result.columns) == 0:
+ return
+
+ field_names = [field.name for field in return_type.fields]
+ actual_names = (
+ list(result.columns[: len(field_names)]) if truncate_return_schema
else list(result.columns)
+ )
+ # By-name mode only applies when the result has string column names;
+ # a numeric RangeIndex falls back to a by-position count check.
+ by_name = assign_cols_by_name and any(isinstance(n, str) for n in
result.columns)
+ _verify_column_schema(actual_names, field_names,
assign_cols_by_name=by_name)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index ffcf12b22c29..48373827a68d 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -56,9 +56,9 @@ from pyspark.errors import PySparkRuntimeError,
PySparkTypeError, PySparkValueEr
from pyspark.eval_handlers._base import get_eval_type_handler
from pyspark.eval_handlers.utils import extract_key_value_indexes
from pyspark.eval_handlers.verification import (
- verify_iter_result_row_count,
+ _verify_column_schema,
verify_iterator_exhausted,
- verify_output_row_limit,
+ verify_pandas_result,
verify_result_row_count,
verify_return_type,
verify_scalar_result,
@@ -152,61 +152,6 @@ def chain(f, g):
return lambda *a: g(f(*a))
-def _verify_column_schema(
- actual_names: list, expected_names: list, *, assign_cols_by_name: bool
-) -> None:
- """Check column names (by-name) or count (by-position) match the expected
schema."""
- if assign_cols_by_name:
- actual_set = set(actual_names)
- expected_set = set(expected_names)
- missing = sorted(expected_set.difference(actual_set))
- extra = sorted(actual_set.difference(expected_set))
- if missing or extra:
- raise PySparkRuntimeError(
- errorClass="RESULT_COLUMN_NAMES_MISMATCH",
- messageParameters={
- "missing": f" Missing: {', '.join(missing)}." if missing
else "",
- "extra": f" Unexpected: {', '.join(extra)}." if extra else
"",
- },
- )
- elif len(actual_names) != len(expected_names):
- raise PySparkRuntimeError(
- errorClass="RESULT_COLUMN_SCHEMA_MISMATCH",
- messageParameters={
- "expected": str(len(expected_names)),
- "actual": str(len(actual_names)),
- },
- )
-
-
-def verify_pandas_result(
- result: Union["pd.DataFrame", "pd.Series"],
- return_type: DataType,
- assign_cols_by_name: bool,
- truncate_return_schema: bool,
-) -> None:
- import pandas as pd
-
- if not isinstance(return_type, StructType):
- verify_return_type(result, pd.Series)
- return
-
- verify_return_type(result, pd.DataFrame)
-
- # Skip schema check on a fully empty result (no rows and no columns).
- if result.empty and len(result.columns) == 0:
- return
-
- field_names = [field.name for field in return_type.fields]
- actual_names = (
- list(result.columns[: len(field_names)]) if truncate_return_schema
else list(result.columns)
- )
- # By-name mode only applies when the result has string column names;
- # a numeric RangeIndex falls back to a by-position count check.
- by_name = assign_cols_by_name and any(isinstance(n, str) for n in
result.columns)
- _verify_column_schema(actual_names, field_names,
assign_cols_by_name=by_name)
-
-
def verify_arrow_result(
result: Union["pa.Table", "pa.RecordBatch"],
assign_cols_by_name: bool,
@@ -1818,8 +1763,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
- PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
@@ -2188,82 +2131,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
return grouped_func, ser
- if eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF:
- import pandas as pd
- import pyarrow as pa
-
- assert num_udfs == 1, "One MAP_PANDAS_ITER UDF expected here."
- map_udf, _, _, return_type = udfs[0]
- output_schema = StructType([StructField("_0", return_type)])
- iter_type_label = (
- "pandas.DataFrame" if isinstance(return_type, StructType) else
"pandas.Series"
- )
- elem_type = pd.DataFrame if isinstance(return_type, StructType) else
pd.Series
-
- def func(
- split_index: int,
- data: Iterator[pa.RecordBatch],
- ) -> Iterator[pa.RecordBatch]:
- """Apply mapInPandas UDF."""
-
- def dataframe_iter():
- # Input batches have a single struct column (see
- # MapInBatchEvaluatorFactory); convert lazily so peakmem stays
- # bounded by one batch.
- for batch in data:
- yield ArrowToPandasConversion.to_pandas(
- batch,
- timezone=runner_conf.timezone,
- prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
- df_for_struct=True,
- )[0]
-
- result = map_udf(dataframe_iter())
- # The declared signature is Iterator[...], so a strict iterator is
required by
- # default. With the legacy flag, accept any object Python can
iterate over -- via
- # iter(...), which honors both __iter__ and the sequence protocol
(__getitem__) --
- # by adapting it into an iterator.
- if runner_conf.map_in_batch_legacy_accept_any_iterable and not
isinstance(
- result, Iterator
- ):
- try:
- result = iter(result)
- except TypeError:
- pass # Not iterable; fall through to the UDF_RETURN_TYPE
error below.
- if not isinstance(result, Iterator):
- raise PySparkTypeError(
- errorClass="UDF_RETURN_TYPE",
- messageParameters={
- "expected": "iterator of {}".format(iter_type_label),
- "actual": type(result).__name__,
- },
- )
-
- for df in result:
- if not isinstance(df, elem_type):
- raise PySparkTypeError(
- errorClass="UDF_RETURN_TYPE",
- messageParameters={
- "expected": "iterator of
{}".format(iter_type_label),
- "actual": "iterator of
{}".format(type(df).__name__),
- },
- )
- verify_pandas_result(
- df, return_type, assign_cols_by_name=True,
truncate_return_schema=True
- )
- yield PandasToArrowConversion.from_pandas(
- [df],
- output_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_COGROUPED_MAP_PANDAS_UDF:
import pandas as pd
import pyarrow as pa
@@ -2945,82 +2812,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
return func, ser
- if eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF:
- import pandas as pd
- import pyarrow as pa
-
- assert num_udfs == 1, "One SCALAR_PANDAS_ITER UDF expected here."
- udf_func, args_offsets, kwargs_offsets, return_type = udfs[0]
-
- # Pre-compute target schema for output coercion
- return_schema = StructType([StructField("_0", return_type)])
- expected_iter_type = (
- Iterator[pd.DataFrame] if isinstance(return_type, StructType) else
Iterator[pd.Series]
- )
-
- def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
- """Apply scalar pandas iterator UDF"""
-
- num_input_rows = 0
-
- def extract_args(batch: pa.RecordBatch):
- nonlocal num_input_rows
- # Input: Arrow -> pandas Series (struct columns become
DataFrames)
- pandas_columns = ArrowToPandasConversion.to_pandas(
- 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,
- )
- args = tuple(pandas_columns[o] for o in args_offsets)
- num_input_rows += batch.num_rows
- return args[0] if len(args) == 1 else args
-
- # Extract args from input batches (streaming)
- args_iter = map(extract_args, data)
-
- # Call UDF and verify result type (iterator of pd.Series /
pd.DataFrame)
- verified_iter = verify_return_type(udf_func(args_iter),
expected_iter_type)
-
- # Process results: verify each element and convert pandas -> Arrow
- def process_results():
- for result in verified_iter:
- verify_pandas_result(
- result, return_type, assign_cols_by_name=True,
truncate_return_schema=True
- )
- yield PandasToArrowConversion.from_pandas(
- [result],
- 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,
- )
-
- # Apply row limit check (fail-fast)
- limited = verify_output_row_limit(
- process_results(),
- lambda: num_input_rows,
- )
-
- # Apply row count match check (final)
- matched = verify_iter_result_row_count(
- limited,
- lambda: num_input_rows,
- )
-
- # Yield batches
- yield from matched
-
- # Verify iterator consumed
- verify_iterator_exhausted(args_iter)
-
- return func, ser
-
if eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF:
import pandas as pd
import pyarrow as pa
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
