From 948a183952920b23ed24baa0c794587805880a58 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:20:58 +0000 Subject: [PATCH 01/20] refactor: migrate Arrow map/iter UDF eval types to eval handlers Co-authored-by: Isaac --- python/pyspark/eval_handlers/_arrow.py | 333 +++++++++++++- python/pyspark/eval_handlers/verification.py | 93 +++- .../pyspark/tests/test_eval_type_handlers.py | 194 +++++++- python/pyspark/worker.py | 414 +----------------- python/pyspark/worker_util.py | 30 ++ 5 files changed, 645 insertions(+), 419 deletions(-) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 9f5df996ad22d..682bc317dae7a 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -18,22 +18,49 @@ """Handlers for the Arrow-native UDF eval types (the UDF exchanges ``pa.Array`` / ``pa.RecordBatch`` values directly, without a pandas conversion).""" +import itertools from collections.abc import Iterator from typing import TYPE_CHECKING, Any -from pyspark.eval_handlers._base import BatchEvalTypeHandler -from pyspark.eval_handlers.verification import verify_scalar_result +from pyspark.eval_handlers._base import ( + BatchEvalTypeHandler, + CoGroupedEvalTypeHandler, + GroupedEvalTypeHandler, +) +from pyspark.eval_handlers.verification import ( + verify_iter_result_row_count, + verify_iterator_exhausted, + verify_output_row_limit, + verify_return_type, + verify_scalar_result, +) from pyspark.sql.conversion import ArrowBatchTransformer -from pyspark.sql.pandas.types import to_arrow_schema +from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type from pyspark.sql.types import StructField, StructType from pyspark.util import PythonEvalType +from pyspark.worker_util import extract_key_value_indexes if TYPE_CHECKING: import pyarrow as pa + from pyspark.eval_handlers._typing import CoGroupedBatch, GroupedBatch from pyspark.worker_util import EvalConf, RunnerConf +def _arrow_return_schema(return_type: Any, use_large_var_types: bool) -> "pa.Schema": + """Arrow schema for a grouped/cogrouped map UDF's declared struct return type. + + The return type is a StructType, so ``to_arrow_type`` yields a struct type whose + fields are the output columns; the group's output batches carry those fields flat. + """ + import pyarrow as pa + + arrow_return_type = to_arrow_type( + return_type, timezone="UTC", prefers_large_types=use_large_var_types + ) + return pa.schema(list(arrow_return_type)) + + class ArrowScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): """SQL_SCALAR_ARROW_UDF: invoke each UDF once per input RecordBatch, coerce the result to the declared schema, and check the row count.""" @@ -70,3 +97,303 @@ def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[p ) verify_scalar_result(output_batch, batch.num_rows) yield output_batch + + +class ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): + """SQL_SCALAR_ARROW_ITER_UDF: the UDF receives an iterator of the argument + columns and yields an iterator of pa.Array; enforce the declared type on each + result and verify the total row count matches the input.""" + + eval_type = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + assert len(udfs) == 1, "One SCALAR_ARROW_ITER UDF expected here." + self._udf_func, self._args_offsets, _, return_type = udfs[0] + self._arrow_return_type = to_arrow_type( + return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types + ) + + def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]": + import pyarrow as pa + + args_offsets = self._args_offsets + num_input_rows = 0 + + def extract_args(batch: "pa.RecordBatch"): + nonlocal num_input_rows + args = tuple(batch.column(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 pa.Array) + verified_iter = verify_return_type( + self._udf_func(args_iter), + Iterator[pa.Array], # type: ignore[type-abstract] + ) + + # Process results: enforce schema and assemble into RecordBatch + target_schema = pa.schema([pa.field("_0", self._arrow_return_type)]) + + def process_results(): + for result in verified_iter: + batch = pa.RecordBatch.from_arrays([result], ["_0"]) + yield ArrowBatchTransformer.enforce_schema(batch, target_schema, safecheck=True) + + # 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) + + +class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): + """SQL_MAP_ARROW_ITER_UDF (mapInArrow): the single UDF receives the input + RecordBatch stream and yields a RecordBatch stream, exchanged as flattened + columns on the wire and wrapped back into a single struct column.""" + + eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + assert len(udfs) == 1, "One MAP_ARROW_ITER UDF expected here." + self._udf_func = udfs[0][0] + + def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]": + import pyarrow as pa + + # Pre-processing + input_batches: "Iterator[pa.RecordBatch]" = map(ArrowBatchTransformer.flatten_struct, data) + + # invoke the UDF + output_batches = self._udf_func(input_batches) + + # 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( + output_batches, Iterator + ): + try: + output_batches = iter(output_batches) + except TypeError: + # Not iterable at all; leave it so verify_return_type below raises the + # standard UDF_RETURN_TYPE error. + pass + + # Post-processing + verified_iter = verify_return_type( + output_batches, + Iterator[pa.RecordBatch], # type: ignore[type-abstract] + ) + yield from map(ArrowBatchTransformer.wrap_struct, verified_iter) + + +class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): + """SQL_GROUPED_MAP_ARROW_UDF (applyInArrow): the single UDF receives each + group as one pa.Table and returns one pa.Table, coerced to the declared + schema.""" + + eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + assert len(udfs) == 1, "One GROUPED_MAP_ARROW UDF expected here." + self._grouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] + parsed_offsets = extract_key_value_indexes(arg_offsets) + assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_ARROW UDF." + self._key_offsets = parsed_offsets[0][0] + self._value_offsets = parsed_offsets[0][1] + self._arrow_return_schema = _arrow_return_schema( + return_type, runner_conf.use_large_var_types + ) + + def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]": + """Apply groupBy Arrow UDF (non-iterator variant).""" + import pyarrow as pa + + key_offsets = self._key_offsets + value_offsets = self._value_offsets + for group in data: + # Flatten struct column into separate columns + flattened = map(ArrowBatchTransformer.flatten_struct, group) + + # Materialize first batch to get keys + first_batch = next(flattened) + keys = pa.RecordBatch.from_arrays( + [first_batch.columns[o] for o in key_offsets], + [first_batch.schema.names[o] for o in key_offsets], + ) + value_batches = ( + pa.RecordBatch.from_arrays( + [b.columns[o] for o in value_offsets], + [b.schema.names[o] for o in value_offsets], + ) + for b in itertools.chain((first_batch,), flattened) + ) + + # Call UDF + value_table = pa.Table.from_batches(value_batches) + if self._num_udf_args == 1: + result = self._grouped_udf(value_table) + else: + key = tuple(c[0] for c in keys.columns) + result = self._grouped_udf(key, value_table) + + verify_return_type(result, pa.Table) + # Verify types (and reorder by name when configured). + result = ArrowBatchTransformer.enforce_schema( + result, + self._arrow_return_schema, + arrow_cast=False, + reorder_by_name=self._runner_conf.assign_cols_by_name, + ) + + for batch in result.to_batches(): + yield ArrowBatchTransformer.wrap_struct(batch) + + +class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): + """SQL_GROUPED_MAP_ARROW_ITER_UDF: the single UDF receives each group as an + iterator of RecordBatches and returns an iterator of RecordBatches, coerced + to the declared schema.""" + + eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + assert len(udfs) == 1, "One GROUPED_MAP_ARROW_ITER UDF expected here." + self._grouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] + parsed_offsets = extract_key_value_indexes(arg_offsets) + assert len(parsed_offsets) == 1, ( + "Expected one pair of offsets for GROUPED_MAP_ARROW_ITER UDF." + ) + self._key_offsets = parsed_offsets[0][0] + self._value_offsets = parsed_offsets[0][1] + self._arrow_return_schema = _arrow_return_schema( + return_type, runner_conf.use_large_var_types + ) + + def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]": + """Apply groupBy Arrow UDF (iterator variant).""" + import pyarrow as pa + + key_offsets = self._key_offsets + value_offsets = self._value_offsets + for group in data: + # Flatten struct column into separate columns + flattened_iter = map(ArrowBatchTransformer.flatten_struct, group) + + # Materialize first batch to get keys + first_batch = next(flattened_iter) + keys = pa.RecordBatch.from_arrays( + [first_batch.columns[o] for o in key_offsets], + [first_batch.schema.names[o] for o in key_offsets], + ) + value_batches = ( + pa.RecordBatch.from_arrays( + [b.columns[o] for o in value_offsets], + [b.schema.names[o] for o in value_offsets], + ) + for b in itertools.chain((first_batch,), flattened_iter) + ) + + # Call UDF with iterator of batches + if self._num_udf_args == 1: + result = self._grouped_udf(value_batches) + else: + key = tuple(c[0] for c in keys.columns) + result = self._grouped_udf(key, value_batches) + + # Verify (and reorder by name when configured) each output batch + for batch in verify_return_type(result, Iterator[pa.RecordBatch]): + batch = ArrowBatchTransformer.enforce_schema( + batch, + self._arrow_return_schema, + arrow_cast=False, + reorder_by_name=self._runner_conf.assign_cols_by_name, + ) + yield ArrowBatchTransformer.wrap_struct(batch) + + # Drain remaining input batches to maintain stream position + for _ in value_batches: + pass + + +class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]): + """SQL_COGROUPED_MAP_ARROW_UDF (applyInArrow on a cogroup): the single UDF + receives the two sides' value tables and returns one pa.Table, coerced to the + declared schema.""" + + eval_type = PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + assert len(udfs) == 1, "One COGROUPED_MAP_ARROW UDF expected here." + self._cogrouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] + parsed_offsets = extract_key_value_indexes(arg_offsets) + self._left_key_cols, self._left_val_cols = parsed_offsets[0] + self._right_key_cols, self._right_val_cols = parsed_offsets[1] + self._arrow_return_schema = _arrow_return_schema( + return_type, runner_conf.use_large_var_types + ) + + def run(self, split_index: int, data: "Iterator[CoGroupedBatch]") -> "Iterator[pa.RecordBatch]": + """Apply cogroupBy Arrow UDF.""" + import pyarrow as pa + + select_columns = ArrowBatchTransformer.select_columns + + def table_from_batches(batches, cols): + return pa.Table.from_batches([select_columns(b, cols) for b in batches]) + + for left_batches, right_batches in data: + left_keys = table_from_batches(left_batches, self._left_key_cols) + left_values = table_from_batches(left_batches, self._left_val_cols) + right_keys = table_from_batches(right_batches, self._right_key_cols) + right_values = table_from_batches(right_batches, self._right_val_cols) + + if self._num_udf_args == 2: + result = self._cogrouped_udf(left_values, right_values) + else: + key_table = left_keys if left_keys.num_rows > 0 else right_keys + key = tuple(c[0] for c in key_table.columns) + result = self._cogrouped_udf(key, left_values, right_values) + + verify_return_type(result, pa.Table) + # Verify types (and reorder by name when configured). + result = ArrowBatchTransformer.enforce_schema( + result, + self._arrow_return_schema, + arrow_cast=False, + reorder_by_name=self._runner_conf.assign_cols_by_name, + ) + + for batch in result.to_batches(): + yield ArrowBatchTransformer.wrap_struct(batch) diff --git a/python/pyspark/eval_handlers/verification.py b/python/pyspark/eval_handlers/verification.py index f63097fae7804..00c94b8f33ff5 100644 --- a/python/pyspark/eval_handlers/verification.py +++ b/python/pyspark/eval_handlers/verification.py @@ -20,10 +20,13 @@ Used by the eval type handlers and by the worker's ``read_udfs``. """ -from typing import Any +from collections.abc import Iterator +from typing import Any, Callable, Type, TypeVar, Union, get_args, get_origin from pyspark.errors import PySparkRuntimeError, PySparkTypeError +T = TypeVar("T") + def verify_result_row_count(result_length: int, expected: int) -> None: """Raise if the result row count doesn't match the expected input row count.""" @@ -60,3 +63,91 @@ def verify_scalar_result(result: Any, num_rows: int) -> Any: ) verify_result_row_count(result_length, num_rows) return result + + +def _top_level_package(t: type) -> str: + """Return the top-level package of ``t`` (``pandas`` for ``pd.DataFrame``).""" + return (t.__module__ or "").split(".", 1)[0] + + +def verify_return_type(result: T, expected_type: Type[T]) -> T: + """ + Verify a UDF return value against an expected type. + + Returns ``result`` unchanged if ``isinstance(result, expected_type)``. + For ``Iterator[T]``, returns a lazy iterator that checks each element + against ``T`` on consumption. Raises ``PySparkTypeError`` on mismatch. + """ + if get_origin(expected_type) is Iterator: + (element_type,) = get_args(expected_type) + label = f"iterator of {_top_level_package(element_type)}.{element_type.__name__}" + + if not isinstance(result, Iterator): + raise PySparkTypeError( + errorClass="UDF_RETURN_TYPE", + messageParameters={"expected": label, "actual": type(result).__name__}, + ) + + def check_element(element: T) -> T: + if not isinstance(element, element_type): + raise PySparkTypeError( + errorClass="UDF_RETURN_TYPE", + messageParameters={ + "expected": label, + "actual": f"iterator of {type(element).__name__}", + }, + ) + return element + + return map(check_element, result) # type: ignore[return-value] + + if not isinstance(result, expected_type): + raise PySparkTypeError( + errorClass="UDF_RETURN_TYPE", + messageParameters={ + "expected": f"{_top_level_package(expected_type)}.{expected_type.__name__}", + "actual": type(result).__name__, + }, + ) + return result + + +def verify_iterator_exhausted(iterator: Iterator) -> None: + """Verify that an iterator has been fully consumed.""" + try: + next(iterator) + except StopIteration: + pass + else: + raise PySparkRuntimeError(errorClass="INPUT_NOT_FULLY_CONSUMED", messageParameters={}) + + +def verify_output_row_limit( + iterator: Iterator, + max_rows: Union[int, Callable[[], int]], +) -> Iterator: + """Yield elements while verifying total rows do not exceed a limit (fail-fast).""" + total_rows = 0 + for element in iterator: + total_rows += len(element) + if total_rows > (max_rows() if callable(max_rows) else max_rows): + raise PySparkRuntimeError(errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={}) + yield element + + +def verify_iter_result_row_count( + iterator: Iterator, + expected_rows: Callable[[], int], +) -> Iterator: + """Yield elements and verify final row count matches expected exactly. + + ``expected_rows`` is a callable because the expected count is only known once + the iterator is fully consumed (input rows are counted lazily as a side effect + of pulling batches), so it must be read after this generator is exhausted. + """ + actual_rows = 0 + for element in iterator: + actual_rows += len(element) + yield element + + verify_result_row_count(actual_rows, expected_rows()) diff --git a/python/pyspark/tests/test_eval_type_handlers.py b/python/pyspark/tests/test_eval_type_handlers.py index 4f7cd96a2d899..b8629a8e52f32 100644 --- a/python/pyspark/tests/test_eval_type_handlers.py +++ b/python/pyspark/tests/test_eval_type_handlers.py @@ -17,7 +17,14 @@ import unittest -from pyspark.eval_handlers._arrow import ArrowScalarUDFHandler +from pyspark.eval_handlers._arrow import ( + ArrowCoGroupedMapUDFHandler, + ArrowGroupedMapIterUDFHandler, + ArrowGroupedMapUDFHandler, + ArrowMapUDFHandler, + ArrowScalarIterUDFHandler, + ArrowScalarUDFHandler, +) from pyspark.eval_handlers._base import ( BatchEvalTypeHandler, CoGroupedEvalTypeHandler, @@ -31,7 +38,7 @@ ArrowStreamGroupSerializer, ArrowStreamSerializer, ) -from pyspark.sql.types import LongType +from pyspark.sql.types import LongType, StructField, StructType from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message from pyspark.util import PythonEvalType @@ -41,14 +48,22 @@ class _RunnerConf: the handlers under test read.""" use_large_var_types = False + assign_cols_by_name = True + map_in_batch_legacy_accept_any_iterable = False class EvalTypeHandlerTests(unittest.TestCase): - def test_scalar_arrow_udf_is_registered(self): - self.assertIs( - get_eval_type_handler(PythonEvalType.SQL_SCALAR_ARROW_UDF), - ArrowScalarUDFHandler, - ) + def test_arrow_eval_types_are_registered(self): + # Every migrated Arrow map/iter eval type dispatches to its handler by lookup. + for eval_type, handler_cls in ( + (PythonEvalType.SQL_SCALAR_ARROW_UDF, ArrowScalarUDFHandler), + (PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF, ArrowScalarIterUDFHandler), + (PythonEvalType.SQL_MAP_ARROW_ITER_UDF, ArrowMapUDFHandler), + (PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF, ArrowGroupedMapUDFHandler), + (PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF, ArrowGroupedMapIterUDFHandler), + (PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, ArrowCoGroupedMapUDFHandler), + ): + self.assertIs(get_eval_type_handler(eval_type), handler_cls) def test_category_bases_are_abstract(self): # The interface and the three category bases must not be instantiable: @@ -204,6 +219,171 @@ def arrow_bytes(batches): self.assertEqual([b.num_rows for b in right_side], [1]) +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowScalarIterUDFHandlerTests(unittest.TestCase): + def test_end_to_end_output(self): + import pyarrow as pa + + # The UDF receives an iterator of the single argument column and yields + # an iterator of pa.Array; the handler assembles each into a RecordBatch. + def add_one(col_iter): + for col in col_iter: + yield pa.array([v.as_py() + 1 for v in col], type=pa.int64()) + + udfs = [(add_one, [0], {}, LongType())] + handler = ArrowScalarIterUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) + + batches = [ + pa.RecordBatch.from_arrays([pa.array([1, 2], type=pa.int64())], ["_0"]), + pa.RecordBatch.from_arrays([pa.array([3], type=pa.int64())], ["_0"]), + ] + out = list(handler.run(0, iter(batches))) + self.assertEqual([b.column(0).to_pylist() for b in out], [[2, 3], [4]]) + + def test_row_count_mismatch_is_rejected(self): + import pyarrow as pa + + from pyspark.errors import PySparkRuntimeError + + # Emitting more rows than were consumed must fail (fail-fast row limit). + def too_many(col_iter): + for col in col_iter: + yield pa.array(list(range(len(col) + 1)), type=pa.int64()) + + udfs = [(too_many, [0], {}, LongType())] + handler = ArrowScalarIterUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) + batch = pa.RecordBatch.from_arrays([pa.array([1, 2], type=pa.int64())], ["_0"]) + with self.assertRaises(PySparkRuntimeError): + list(handler.run(0, iter([batch]))) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowMapUDFHandlerTests(unittest.TestCase): + def test_end_to_end_output(self): + import pyarrow as pa + + from pyspark.sql.conversion import ArrowBatchTransformer + + # mapInArrow exchanges a single struct column on the wire; the handler + # flattens it for the UDF and re-wraps the UDF's output. + def double_a(batch_iter): + for batch in batch_iter: + doubled = pa.array([v.as_py() * 2 for v in batch.column(0)], type=pa.int64()) + yield pa.RecordBatch.from_arrays([doubled], ["a"]) + + inner = pa.RecordBatch.from_arrays([pa.array([1, 2, 3], type=pa.int64())], ["a"]) + wrapped = ArrowBatchTransformer.wrap_struct(inner) + + udfs = [(double_a, None, None, None)] + handler = ArrowMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) + out = list(handler.run(0, iter([wrapped]))) + + self.assertEqual(len(out), 1) + # Output is a single struct column; its "a" field carries the doubled values. + self.assertEqual(out[0].num_columns, 1) + self.assertEqual(out[0].column(0).field("a").to_pylist(), [2, 4, 6]) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowGroupedMapUDFHandlerTests(unittest.TestCase): + # arg_offsets encoding for one DataFrame with key column 0 and value column 1: + # [group_len=3, num_keys=1, key_offset=0, value_offset=1] + _ARG_OFFSETS = [3, 1, 0, 1] + + def _grouped_input(self): + import pyarrow as pa + + from pyspark.sql.conversion import ArrowBatchTransformer + + inner = pa.RecordBatch.from_arrays( + [pa.array([7, 7], type=pa.int64()), pa.array([1, 2], type=pa.int64())], ["k", "v"] + ) + wrapped = ArrowBatchTransformer.wrap_struct(inner) + # One group, whose batches arrive as an iterator (matching the group serializer). + return iter([iter([wrapped])]) + + def test_values_only(self): + import pyarrow as pa + + return_type = StructType([StructField("v", LongType())]) + + def grouped_udf(value_table): + return pa.table({"v": pa.array([c.as_py() * 10 for c in value_table.column("v")])}) + + udfs = [(grouped_udf, self._ARG_OFFSETS, return_type, 1)] + handler = ArrowGroupedMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) + out = list(handler.run(0, self._grouped_input())) + self.assertEqual(out[0].column(0).field("v").to_pylist(), [10, 20]) + + def test_key_and_values(self): + import pyarrow as pa + + return_type = StructType([StructField("v", LongType())]) + + def grouped_udf(key, value_table): + # key is the grouping-key tuple; add it to every value. + k = key[0].as_py() + return pa.table({"v": pa.array([c.as_py() + k for c in value_table.column("v")])}) + + udfs = [(grouped_udf, self._ARG_OFFSETS, return_type, 2)] + handler = ArrowGroupedMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) + out = list(handler.run(0, self._grouped_input())) + self.assertEqual(out[0].column(0).field("v").to_pylist(), [8, 9]) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowGroupedMapIterUDFHandlerTests(unittest.TestCase): + def test_end_to_end_output(self): + import pyarrow as pa + + from pyspark.sql.conversion import ArrowBatchTransformer + + return_type = StructType([StructField("v", LongType())]) + + def grouped_udf(value_batches): + for batch in value_batches: + yield pa.RecordBatch.from_arrays( + [pa.array([c.as_py() + 1 for c in batch.column("v")], type=pa.int64())], ["v"] + ) + + inner = pa.RecordBatch.from_arrays( + [pa.array([7], type=pa.int64()), pa.array([41], type=pa.int64())], ["k", "v"] + ) + wrapped = ArrowBatchTransformer.wrap_struct(inner) + udfs = [(grouped_udf, [3, 1, 0, 1], return_type, 1)] + handler = ArrowGroupedMapIterUDFHandler( + udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None + ) + out = list(handler.run(0, iter([iter([wrapped])]))) + self.assertEqual(out[0].column(0).field("v").to_pylist(), [42]) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowCoGroupedMapUDFHandlerTests(unittest.TestCase): + def test_end_to_end_output(self): + import pyarrow as pa + + return_type = StructType([StructField("out", LongType())]) + + def cogrouped_udf(left_values, right_values): + total = left_values.column("lv")[0].as_py() + right_values.column("rv")[0].as_py() + return pa.table({"out": pa.array([total], type=pa.int64())}) + + # A co-group deserializes to a pair of lists of (non-struct-wrapped) batches. + left = pa.RecordBatch.from_arrays( + [pa.array([5], type=pa.int64()), pa.array([10], type=pa.int64())], ["k", "lv"] + ) + right = pa.RecordBatch.from_arrays( + [pa.array([5], type=pa.int64()), pa.array([20], type=pa.int64())], ["k", "rv"] + ) + # Two DataFrames, each key column 0 and value column 1. + arg_offsets = [3, 1, 0, 1, 3, 1, 0, 1] + udfs = [(cogrouped_udf, arg_offsets, return_type, 2)] + handler = ArrowCoGroupedMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) + out = list(handler.run(0, iter([([left], [right])]))) + self.assertEqual(out[0].column(0).field("out").to_pylist(), [30]) + + if __name__ == "__main__": from pyspark.testing import main diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 28d2ae348ec40..7d05f2a246af7 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -36,15 +36,9 @@ Iterable, Optional, Tuple, - Type, - TypeVar, Union, - get_args, - get_origin, ) -T = TypeVar("T") - if TYPE_CHECKING: import pandas as pd import pyarrow as pa @@ -60,7 +54,11 @@ from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.eval_handlers._base import get_eval_type_handler from pyspark.eval_handlers.verification import ( + verify_iter_result_row_count, + verify_iterator_exhausted, + verify_output_row_limit, verify_result_row_count, + verify_return_type, verify_scalar_result, ) from pyspark.logger.worker_io import capture_outputs @@ -118,6 +116,7 @@ EvalConf, RunnerConf, check_python_version, + extract_key_value_indexes, get_sock_file_to_executor, pickleSer, read_command, @@ -177,94 +176,6 @@ def canon(v: Any) -> Any: return object() -def verify_return_type(result: T, expected_type: Type[T]) -> T: - """ - Verify a UDF return value against an expected type. - - Returns ``result`` unchanged if ``isinstance(result, expected_type)``. - For ``Iterator[T]``, returns a lazy iterator that checks each element - against ``T`` on consumption. Raises ``PySparkTypeError`` on mismatch. - """ - if get_origin(expected_type) is Iterator: - (element_type,) = get_args(expected_type) - label = f"iterator of {_top_level_package(element_type)}.{element_type.__name__}" - - if not isinstance(result, Iterator): - raise PySparkTypeError( - errorClass="UDF_RETURN_TYPE", - messageParameters={"expected": label, "actual": type(result).__name__}, - ) - - def check_element(element: T) -> T: - if not isinstance(element, element_type): - raise PySparkTypeError( - errorClass="UDF_RETURN_TYPE", - messageParameters={ - "expected": label, - "actual": f"iterator of {type(element).__name__}", - }, - ) - return element - - return map(check_element, result) # type: ignore[return-value] - - if not isinstance(result, expected_type): - raise PySparkTypeError( - errorClass="UDF_RETURN_TYPE", - messageParameters={ - "expected": f"{_top_level_package(expected_type)}.{expected_type.__name__}", - "actual": type(result).__name__, - }, - ) - return result - - -def _top_level_package(t: type) -> str: - """Return the top-level package of ``t`` (``pandas`` for ``pd.DataFrame``).""" - return (t.__module__ or "").split(".", 1)[0] - - -def verify_iterator_exhausted(iterator: Iterator) -> None: - """Verify that an iterator has been fully consumed.""" - try: - next(iterator) - except StopIteration: - pass - else: - raise PySparkRuntimeError(errorClass="INPUT_NOT_FULLY_CONSUMED", messageParameters={}) - - -def verify_output_row_limit( - iterator: Iterator, - max_rows: Union[int, Callable[[], int]], -) -> Iterator: - """Yield elements while verifying total rows do not exceed a limit (fail-fast).""" - total_rows = 0 - for element in iterator: - total_rows += len(element) - if total_rows > (max_rows() if callable(max_rows) else max_rows): - raise PySparkRuntimeError(errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={}) - yield element - - -def verify_iter_result_row_count( - iterator: Iterator, - expected_rows: Callable[[], int], -) -> Iterator: - """Yield elements and verify final row count matches expected exactly. - - ``expected_rows`` is a callable because the expected count is only known once - the iterator is fully consumed (input rows are counted lazily as a side effect - of pulling batches), so it must be read after this generator is exhausted. - """ - actual_rows = 0 - for element in iterator: - actual_rows += len(element) - yield element - - verify_result_row_count(actual_rows, expected_rows()) - - def _verify_column_schema( actual_names: list, expected_names: list, *, assign_cols_by_name: bool ) -> None: @@ -1902,9 +1813,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): PythonEvalType.SQL_SCALAR_PANDAS_UDF, PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF, - PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, - PythonEvalType.SQL_MAP_ARROW_ITER_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, @@ -1914,9 +1823,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, - PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF, - PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF, - PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF, PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF, PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF, @@ -1935,8 +1841,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): # inside the worker, so it uses the plain (non-grouped) stream serializer below. Only # the post-shuffle FINAL stage receives one Arrow stream per group. PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, - PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF, - PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, @@ -1944,10 +1848,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF, ): ser = ArrowStreamGroupSerializer(write_start_stream=True) - elif eval_type in ( - PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, - PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF, - ): + elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF: ser = ArrowStreamCoGroupSerializer(write_start_stream=True) else: ser = ArrowStreamSerializer(write_start_stream=True) @@ -1962,134 +1863,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): num_udfs = len(udfs) - def extract_key_value_indexes(grouped_arg_offsets): - """ - Helper function to extract the key and value indexes from arg_offsets for the grouped and - cogrouped pandas udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. - - Parameters - ---------- - grouped_arg_offsets: list - List containing the key and value indexes of columns of the - DataFrames to be passed to the udf. It consists of n repeating groups where n is the - number of DataFrames. Each group has the following format: - group[0]: length of group - group[1]: length of key indexes - group[2.. group[1] +2]: key attributes - group[group[1] +3 group[0]]: value attributes - """ - parsed = [] - idx = 0 - while idx < len(grouped_arg_offsets): - offsets_len = grouped_arg_offsets[idx] - idx += 1 - offsets = grouped_arg_offsets[idx : idx + offsets_len] - split_index = offsets[0] + 1 - offset_keys = offsets[1:split_index] - offset_values = offsets[split_index:] - parsed.append([offset_keys, offset_values]) - idx += offsets_len - return parsed - - if eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF: - import pyarrow as pa - - assert num_udfs == 1, "One MAP_ARROW_ITER UDF expected here." - udf_func: Callable[[Iterator[pa.RecordBatch]], Iterator[pa.RecordBatch]] = udfs[0][0] - - def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: - """Apply mapInArrow UDF""" - - # Pre-processing - input_batches: Iterator[pa.RecordBatch] = map( - ArrowBatchTransformer.flatten_struct, data - ) - - # invoke the UDF - output_batches = udf_func(input_batches) - - # 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 runner_conf.map_in_batch_legacy_accept_any_iterable and not isinstance( - output_batches, Iterator - ): - try: - output_batches = iter(output_batches) - except TypeError: - # Not iterable at all; leave it so verify_return_type below raises the - # standard UDF_RETURN_TYPE error. - pass - - # Post-processing - verified_iter = verify_return_type( - output_batches, - Iterator[pa.RecordBatch], # type: ignore[type-abstract] - ) - yield from map(ArrowBatchTransformer.wrap_struct, verified_iter) - - return func, ser - - if eval_type == PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF: - import pyarrow as pa - - assert num_udfs == 1, "One SCALAR_ARROW_ITER UDF expected here." - udf_func, args_offsets, kwargs_offsets, return_type = udfs[0] - - # Pre-compute target Arrow type for output coercion - arrow_return_type = to_arrow_type( - return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types - ) - - def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: - """Apply scalar Arrow iterator UDF""" - - num_input_rows = 0 - - def extract_args(batch: pa.RecordBatch): - nonlocal num_input_rows - args = tuple(batch.column(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 pa.Array) - verified_iter = verify_return_type( - udf_func(args_iter), - Iterator[pa.Array], # type: ignore[type-abstract] - ) - - # Process results: enforce schema and assemble into RecordBatch - target_schema = pa.schema([pa.field("_0", arrow_return_type)]) - - def process_results(): - for result in verified_iter: - batch = pa.RecordBatch.from_arrays([result], ["_0"]) - yield ArrowBatchTransformer.enforce_schema(batch, target_schema, safecheck=True) - - # 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_GROUPED_AGG_ARROW_UDF: import pyarrow as pa @@ -2659,129 +2432,6 @@ def grouped_func( return grouped_func, ser - if eval_type == PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF: - import pyarrow as pa - - assert num_udfs == 1, "One GROUPED_MAP_ARROW UDF expected here." - grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] - parsed_offsets = extract_key_value_indexes(arg_offsets) - assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_ARROW UDF." - - arrow_return_type = to_arrow_type( - return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types - ) - arrow_return_schema = pa.schema(list(arrow_return_type)) - - key_offsets = parsed_offsets[0][0] - value_offsets = parsed_offsets[0][1] - - def grouped_func( - split_index: int, data: Iterator["GroupedBatch"] - ) -> Iterator[pa.RecordBatch]: - """Apply groupBy Arrow UDF (non-iterator variant).""" - for group in data: - # Flatten struct column into separate columns - flattened = map(ArrowBatchTransformer.flatten_struct, group) - - # Materialize first batch to get keys - first_batch = next(flattened) - keys = pa.RecordBatch.from_arrays( - [first_batch.columns[o] for o in key_offsets], - [first_batch.schema.names[o] for o in key_offsets], - ) - value_batches = ( - pa.RecordBatch.from_arrays( - [b.columns[o] for o in value_offsets], - [b.schema.names[o] for o in value_offsets], - ) - for b in itertools.chain((first_batch,), flattened) - ) - - # Call UDF - value_table = pa.Table.from_batches(value_batches) - if num_udf_args == 1: - result = grouped_udf(value_table) - else: - key = tuple(c[0] for c in keys.columns) - result = grouped_udf(key, value_table) - - verify_return_type(result, pa.Table) - # Verify types (and reorder by name when configured). - result = ArrowBatchTransformer.enforce_schema( - result, - arrow_return_schema, - arrow_cast=False, - reorder_by_name=runner_conf.assign_cols_by_name, - ) - - for batch in result.to_batches(): - yield ArrowBatchTransformer.wrap_struct(batch) - - return grouped_func, ser - - if eval_type == PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF: - import pyarrow as pa - - assert num_udfs == 1, "One GROUPED_MAP_ARROW_ITER UDF expected here." - grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] - parsed_offsets = extract_key_value_indexes(arg_offsets) - assert len(parsed_offsets) == 1, ( - "Expected one pair of offsets for GROUPED_MAP_ARROW_ITER UDF." - ) - - arrow_return_type = to_arrow_type( - return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types - ) - arrow_return_schema = pa.schema(list(arrow_return_type)) - - key_offsets = parsed_offsets[0][0] - value_offsets = parsed_offsets[0][1] - - def grouped_func( - split_index: int, data: Iterator["GroupedBatch"] - ) -> Iterator[pa.RecordBatch]: - """Apply groupBy Arrow UDF (iterator variant).""" - for group in data: - # Flatten struct column into separate columns - flattened_iter = map(ArrowBatchTransformer.flatten_struct, group) - - # Materialize first batch to get keys - first_batch = next(flattened_iter) - keys = pa.RecordBatch.from_arrays( - [first_batch.columns[o] for o in key_offsets], - [first_batch.schema.names[o] for o in key_offsets], - ) - value_batches = ( - pa.RecordBatch.from_arrays( - [b.columns[o] for o in value_offsets], - [b.schema.names[o] for o in value_offsets], - ) - for b in itertools.chain((first_batch,), flattened_iter) - ) - - # Call UDF with iterator of batches - if num_udf_args == 1: - result = grouped_udf(value_batches) - else: - key = tuple(c[0] for c in keys.columns) - result = grouped_udf(key, value_batches) - - # Verify (and reorder by name when configured) each output batch - for batch in verify_return_type(result, Iterator[pa.RecordBatch]): - batch = ArrowBatchTransformer.enforce_schema( - batch, - arrow_return_schema, - arrow_cast=False, - reorder_by_name=runner_conf.assign_cols_by_name, - ) - yield ArrowBatchTransformer.wrap_struct(batch) - - # Drain remaining input batches to maintain stream position - for _ in value_batches: - pass - - return grouped_func, ser - if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF: import pandas as pd import pyarrow as pa @@ -2941,58 +2591,6 @@ def dataframe_iter(): return grouped_func, ser - if eval_type == PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF: - import pyarrow as pa - - assert num_udfs == 1, "One COGROUPED_MAP_ARROW UDF expected here." - cogrouped_udf, arg_offsets, return_type, num_udf_args = udfs[0] - - parsed_offsets = extract_key_value_indexes(arg_offsets) - - arrow_return_type = to_arrow_type( - return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types - ) - arrow_return_schema = pa.schema(list(arrow_return_type)) - - select_columns = ArrowBatchTransformer.select_columns - left_key_cols, left_val_cols = parsed_offsets[0] - right_key_cols, right_val_cols = parsed_offsets[1] - - def table_from_batches(batches, cols): - return pa.Table.from_batches([select_columns(b, cols) for b in batches]) - - def cogrouped_func( - split_index: int, - data: Iterator[Tuple[list[pa.RecordBatch], list[pa.RecordBatch]]], - ) -> Iterator[pa.RecordBatch]: - """Apply cogroupBy Arrow UDF.""" - for left_batches, right_batches in data: - left_keys = table_from_batches(left_batches, left_key_cols) - left_values = table_from_batches(left_batches, left_val_cols) - right_keys = table_from_batches(right_batches, right_key_cols) - right_values = table_from_batches(right_batches, right_val_cols) - - if num_udf_args == 2: - result = cogrouped_udf(left_values, right_values) - else: - key_table = left_keys if left_keys.num_rows > 0 else right_keys - key = tuple(c[0] for c in key_table.columns) - result = cogrouped_udf(key, left_values, right_values) - - verify_return_type(result, pa.Table) - # Verify types (and reorder by name when configured). - result = ArrowBatchTransformer.enforce_schema( - result, - arrow_return_schema, - arrow_cast=False, - reorder_by_name=runner_conf.assign_cols_by_name, - ) - - for batch in result.to_batches(): - yield ArrowBatchTransformer.wrap_struct(batch) - - return cogrouped_func, ser - if eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF: import pandas as pd import pyarrow as pa diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py index 3a24bde38116b..9f2f9f861ba5e 100644 --- a/python/pyspark/worker_util.py +++ b/python/pyspark/worker_util.py @@ -265,6 +265,36 @@ def send_accumulator_updates(outfile: IO) -> None: pickleSer._write_with_length((aid, accum._value), outfile) +def extract_key_value_indexes(grouped_arg_offsets: list) -> list: + """ + Helper function to extract the key and value indexes from arg_offsets for the grouped and + cogrouped grouped-map udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. + + Parameters + ---------- + grouped_arg_offsets: list + List containing the key and value indexes of columns of the + DataFrames to be passed to the udf. It consists of n repeating groups where n is the + number of DataFrames. Each group has the following format: + group[0]: length of group + group[1]: length of key indexes + group[2.. group[1] +2]: key attributes + group[group[1] +3 group[0]]: value attributes + """ + parsed = [] + idx = 0 + while idx < len(grouped_arg_offsets): + offsets_len = grouped_arg_offsets[idx] + idx += 1 + offsets = grouped_arg_offsets[idx : idx + offsets_len] + split_index = offsets[0] + 1 + offset_keys = offsets[1:split_index] + offset_values = offsets[split_index:] + parsed.append([offset_keys, offset_values]) + idx += offsets_len + return parsed + + class Conf: def __init__(self, infile_or_dict: Optional[Union[dict[str, str], IO]] = None) -> None: self._conf: dict[str, Any] = {} From ea5295d2b53a2deb908fda351170512080140d93 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:27:50 +0000 Subject: [PATCH 02/20] fix: address review - top-level pyarrow, no type-ignore, non-string annotations Import pyarrow at the top of _arrow.py and guard its import in the package __init__ behind pyarrow availability. Move extract_key_value_indexes to the driver-safe _base.py (worker_util is worker-only under SPARK_TESTING). Use from __future__ import annotations so annotations are real names, not strings, and drop the type: ignore in verification.py via cast. Co-authored-by: Isaac --- python/pyspark/eval_handlers/__init__.py | 12 ++- python/pyspark/eval_handlers/_arrow.py | 99 ++++++++----------- python/pyspark/eval_handlers/_base.py | 30 ++++++ python/pyspark/eval_handlers/verification.py | 4 +- .../pyspark/tests/test_eval_type_handlers.py | 20 ++-- python/pyspark/worker.py | 5 +- python/pyspark/worker_util.py | 30 ------ 7 files changed, 97 insertions(+), 103 deletions(-) diff --git a/python/pyspark/eval_handlers/__init__.py b/python/pyspark/eval_handlers/__init__.py index 4c51849909861..e2490c83bae27 100644 --- a/python/pyspark/eval_handlers/__init__.py +++ b/python/pyspark/eval_handlers/__init__.py @@ -22,7 +22,15 @@ Each eval type handled here is an ``EvalTypeHandler`` subclass (in ``_base``) that 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 (``_arrow``) so they register. +imports the concrete handler submodules so they register. + +``_arrow`` requires pyarrow and imports it at module top, so it is only imported +when pyarrow is available; the Arrow eval types it serves cannot run without it. """ -from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import +try: + import pyarrow # noqa: F401 +except ImportError: + pass +else: + from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 682bc317dae7a..952945a16d21c 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -16,16 +16,25 @@ # """Handlers for the Arrow-native UDF eval types (the UDF exchanges ``pa.Array`` / -``pa.RecordBatch`` values directly, without a pandas conversion).""" +``pa.RecordBatch`` values directly, without a pandas conversion). + +This module imports ``pyarrow`` at the top level, so the package ``__init__`` only +imports it when pyarrow is available; callers must do the same. +""" + +from __future__ import annotations import itertools from collections.abc import Iterator from typing import TYPE_CHECKING, Any +import pyarrow as pa + from pyspark.eval_handlers._base import ( BatchEvalTypeHandler, CoGroupedEvalTypeHandler, GroupedEvalTypeHandler, + extract_key_value_indexes, ) from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, @@ -36,39 +45,37 @@ ) from pyspark.sql.conversion import ArrowBatchTransformer from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type -from pyspark.sql.types import StructField, StructType +from pyspark.sql.types import DataType, StructField, StructType from pyspark.util import PythonEvalType -from pyspark.worker_util import extract_key_value_indexes if TYPE_CHECKING: - import pyarrow as pa - + # Annotation-only: these live in worker-only or forward-ref modules, so they are + # not imported at runtime. ``from __future__ import annotations`` keeps every + # annotation below an unevaluated name rather than a string literal. from pyspark.eval_handlers._typing import CoGroupedBatch, GroupedBatch from pyspark.worker_util import EvalConf, RunnerConf -def _arrow_return_schema(return_type: Any, use_large_var_types: bool) -> "pa.Schema": +def _arrow_return_schema(return_type: DataType, use_large_var_types: bool) -> pa.Schema: """Arrow schema for a grouped/cogrouped map UDF's declared struct return type. The return type is a StructType, so ``to_arrow_type`` yields a struct type whose fields are the output columns; the group's output batches carry those fields flat. """ - import pyarrow as pa - arrow_return_type = to_arrow_type( return_type, timezone="UTC", prefers_large_types=use_large_var_types ) return pa.schema(list(arrow_return_type)) -class ArrowScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): +class ArrowScalarUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): """SQL_SCALAR_ARROW_UDF: invoke each UDF once per input RecordBatch, coerce the result to the declared schema, and check the row count.""" eval_type = PythonEvalType.SQL_SCALAR_ARROW_UDF def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) self._col_names = ["_%d" % i for i in range(len(udfs))] @@ -78,9 +85,7 @@ def __init__( prefers_large_types=runner_conf.use_large_var_types, ) - def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]": - import pyarrow as pa - + def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: for batch in data: output_batch = pa.RecordBatch.from_arrays( [ @@ -99,7 +104,7 @@ def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[p yield output_batch -class ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): +class ArrowScalarIterUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): """SQL_SCALAR_ARROW_ITER_UDF: the UDF receives an iterator of the argument columns and yields an iterator of pa.Array; enforce the declared type on each result and verify the total row count matches the input.""" @@ -107,7 +112,7 @@ class ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): eval_type = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One SCALAR_ARROW_ITER UDF expected here." @@ -116,13 +121,11 @@ def __init__( return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types ) - def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]": - import pyarrow as pa - + def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: args_offsets = self._args_offsets num_input_rows = 0 - def extract_args(batch: "pa.RecordBatch"): + def extract_args(batch: pa.RecordBatch) -> Any: nonlocal num_input_rows args = tuple(batch.column(o) for o in args_offsets) num_input_rows += batch.num_rows @@ -132,30 +135,21 @@ def extract_args(batch: "pa.RecordBatch"): args_iter = map(extract_args, data) # Call UDF and verify result type (iterator of pa.Array) - verified_iter = verify_return_type( - self._udf_func(args_iter), - Iterator[pa.Array], # type: ignore[type-abstract] - ) + verified_iter = verify_return_type(self._udf_func(args_iter), Iterator[pa.Array]) # Process results: enforce schema and assemble into RecordBatch target_schema = pa.schema([pa.field("_0", self._arrow_return_type)]) - def process_results(): + def process_results() -> Iterator[pa.RecordBatch]: for result in verified_iter: batch = pa.RecordBatch.from_arrays([result], ["_0"]) yield ArrowBatchTransformer.enforce_schema(batch, target_schema, safecheck=True) # Apply row limit check (fail-fast) - limited = verify_output_row_limit( - process_results(), - lambda: num_input_rows, - ) + 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, - ) + matched = verify_iter_result_row_count(limited, lambda: num_input_rows) # Yield batches yield from matched @@ -164,7 +158,7 @@ def process_results(): verify_iterator_exhausted(args_iter) -class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): +class ArrowMapUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): """SQL_MAP_ARROW_ITER_UDF (mapInArrow): the single UDF receives the input RecordBatch stream and yields a RecordBatch stream, exchanged as flattened columns on the wire and wrapped back into a single struct column.""" @@ -172,17 +166,15 @@ class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One MAP_ARROW_ITER UDF expected here." self._udf_func = udfs[0][0] - def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]": - import pyarrow as pa - + def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: # Pre-processing - input_batches: "Iterator[pa.RecordBatch]" = map(ArrowBatchTransformer.flatten_struct, data) + input_batches = map(ArrowBatchTransformer.flatten_struct, data) # invoke the UDF output_batches = self._udf_func(input_batches) @@ -202,14 +194,11 @@ def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[p pass # Post-processing - verified_iter = verify_return_type( - output_batches, - Iterator[pa.RecordBatch], # type: ignore[type-abstract] - ) + verified_iter = verify_return_type(output_batches, Iterator[pa.RecordBatch]) yield from map(ArrowBatchTransformer.wrap_struct, verified_iter) -class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): +class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): """SQL_GROUPED_MAP_ARROW_UDF (applyInArrow): the single UDF receives each group as one pa.Table and returns one pa.Table, coerced to the declared schema.""" @@ -217,7 +206,7 @@ class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One GROUPED_MAP_ARROW UDF expected here." @@ -230,10 +219,8 @@ def __init__( return_type, runner_conf.use_large_var_types ) - def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]": + def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: """Apply groupBy Arrow UDF (non-iterator variant).""" - import pyarrow as pa - key_offsets = self._key_offsets value_offsets = self._value_offsets for group in data: @@ -275,7 +262,7 @@ def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa. yield ArrowBatchTransformer.wrap_struct(batch) -class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): +class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): """SQL_GROUPED_MAP_ARROW_ITER_UDF: the single UDF receives each group as an iterator of RecordBatches and returns an iterator of RecordBatches, coerced to the declared schema.""" @@ -283,7 +270,7 @@ class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One GROUPED_MAP_ARROW_ITER UDF expected here." @@ -298,10 +285,8 @@ def __init__( return_type, runner_conf.use_large_var_types ) - def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]": + def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: """Apply groupBy Arrow UDF (iterator variant).""" - import pyarrow as pa - key_offsets = self._key_offsets value_offsets = self._value_offsets for group in data: @@ -344,7 +329,7 @@ def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa. pass -class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]): +class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler[pa.RecordBatch]): """SQL_COGROUPED_MAP_ARROW_UDF (applyInArrow on a cogroup): the single UDF receives the two sides' value tables and returns one pa.Table, coerced to the declared schema.""" @@ -352,7 +337,7 @@ class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]): eval_type = PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf" + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One COGROUPED_MAP_ARROW UDF expected here." @@ -364,13 +349,11 @@ def __init__( return_type, runner_conf.use_large_var_types ) - def run(self, split_index: int, data: "Iterator[CoGroupedBatch]") -> "Iterator[pa.RecordBatch]": + def run(self, split_index: int, data: Iterator[CoGroupedBatch]) -> Iterator[pa.RecordBatch]: """Apply cogroupBy Arrow UDF.""" - import pyarrow as pa - select_columns = ArrowBatchTransformer.select_columns - def table_from_batches(batches, cols): + def table_from_batches(batches: list[pa.RecordBatch], cols: list[int]) -> pa.Table: return pa.Table.from_batches([select_columns(b, cols) for b in batches]) for left_batches, right_batches in data: diff --git a/python/pyspark/eval_handlers/_base.py b/python/pyspark/eval_handlers/_base.py index ce8003d7aabeb..c83b848996e58 100644 --- a/python/pyspark/eval_handlers/_base.py +++ b/python/pyspark/eval_handlers/_base.py @@ -52,6 +52,36 @@ def get_eval_type_handler(eval_type: int) -> "Optional[type[EvalTypeHandler]]": return _eval_type_handlers.get(eval_type) +def extract_key_value_indexes(grouped_arg_offsets: list) -> list: + """ + Extract the key and value indexes from arg_offsets for the grouped and cogrouped grouped-map + udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. + + Parameters + ---------- + grouped_arg_offsets: list + List containing the key and value indexes of columns of the + DataFrames to be passed to the udf. It consists of n repeating groups where n is the + number of DataFrames. Each group has the following format: + group[0]: length of group + group[1]: length of key indexes + group[2.. group[1] +2]: key attributes + group[group[1] +3 group[0]]: value attributes + """ + parsed = [] + idx = 0 + while idx < len(grouped_arg_offsets): + offsets_len = grouped_arg_offsets[idx] + idx += 1 + offsets = grouped_arg_offsets[idx : idx + offsets_len] + split_index = offsets[0] + 1 + offset_keys = offsets[1:split_index] + offset_values = offsets[split_index:] + parsed.append([offset_keys, offset_values]) + idx += offsets_len + return parsed + + class _EvalTypeHandlerMeta(ABCMeta): """Registers a concrete handler under its ``eval_type`` at class definition. diff --git a/python/pyspark/eval_handlers/verification.py b/python/pyspark/eval_handlers/verification.py index 00c94b8f33ff5..c77096c622d9a 100644 --- a/python/pyspark/eval_handlers/verification.py +++ b/python/pyspark/eval_handlers/verification.py @@ -21,7 +21,7 @@ """ from collections.abc import Iterator -from typing import Any, Callable, Type, TypeVar, Union, get_args, get_origin +from typing import Any, Callable, Type, TypeVar, Union, cast, get_args, get_origin from pyspark.errors import PySparkRuntimeError, PySparkTypeError @@ -99,7 +99,7 @@ def check_element(element: T) -> T: ) return element - return map(check_element, result) # type: ignore[return-value] + return cast(T, map(check_element, result)) if not isinstance(result, expected_type): raise PySparkTypeError( diff --git a/python/pyspark/tests/test_eval_type_handlers.py b/python/pyspark/tests/test_eval_type_handlers.py index b8629a8e52f32..ff0b3f821f16d 100644 --- a/python/pyspark/tests/test_eval_type_handlers.py +++ b/python/pyspark/tests/test_eval_type_handlers.py @@ -17,14 +17,6 @@ import unittest -from pyspark.eval_handlers._arrow import ( - ArrowCoGroupedMapUDFHandler, - ArrowGroupedMapIterUDFHandler, - ArrowGroupedMapUDFHandler, - ArrowMapUDFHandler, - ArrowScalarIterUDFHandler, - ArrowScalarUDFHandler, -) from pyspark.eval_handlers._base import ( BatchEvalTypeHandler, CoGroupedEvalTypeHandler, @@ -42,6 +34,17 @@ from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message from pyspark.util import PythonEvalType +if have_pyarrow: + # The handlers live in ``_arrow``, which imports pyarrow at module top. + from pyspark.eval_handlers._arrow import ( + ArrowCoGroupedMapUDFHandler, + ArrowGroupedMapIterUDFHandler, + ArrowGroupedMapUDFHandler, + ArrowMapUDFHandler, + ArrowScalarIterUDFHandler, + ArrowScalarUDFHandler, + ) + class _RunnerConf: """Minimal stand-in for the worker's RunnerConf, exposing only the fields @@ -52,6 +55,7 @@ class _RunnerConf: map_in_batch_legacy_accept_any_iterable = False +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class EvalTypeHandlerTests(unittest.TestCase): def test_arrow_eval_types_are_registered(self): # Every migrated Arrow map/iter eval type dispatches to its handler by lookup. diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 7d05f2a246af7..2cc24176f2917 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -52,7 +52,7 @@ _deserialize_accumulator, ) from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError -from pyspark.eval_handlers._base import get_eval_type_handler +from pyspark.eval_handlers._base import extract_key_value_indexes, get_eval_type_handler from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, @@ -116,7 +116,6 @@ EvalConf, RunnerConf, check_python_version, - extract_key_value_indexes, get_sock_file_to_executor, pickleSer, read_command, @@ -3243,7 +3242,7 @@ def extract_flat(batch: pa.RecordBatch): if not is_pandas: verified_iter = verify_return_type( udf_func(flat_args_iter), - Iterator[pa.Array], # type: ignore[type-abstract] + Iterator[pa.Array], ) else: pandas_iter_type = ( diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py index 9f2f9f861ba5e..3a24bde38116b 100644 --- a/python/pyspark/worker_util.py +++ b/python/pyspark/worker_util.py @@ -265,36 +265,6 @@ def send_accumulator_updates(outfile: IO) -> None: pickleSer._write_with_length((aid, accum._value), outfile) -def extract_key_value_indexes(grouped_arg_offsets: list) -> list: - """ - Helper function to extract the key and value indexes from arg_offsets for the grouped and - cogrouped grouped-map udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. - - Parameters - ---------- - grouped_arg_offsets: list - List containing the key and value indexes of columns of the - DataFrames to be passed to the udf. It consists of n repeating groups where n is the - number of DataFrames. Each group has the following format: - group[0]: length of group - group[1]: length of key indexes - group[2.. group[1] +2]: key attributes - group[group[1] +3 group[0]]: value attributes - """ - parsed = [] - idx = 0 - while idx < len(grouped_arg_offsets): - offsets_len = grouped_arg_offsets[idx] - idx += 1 - offsets = grouped_arg_offsets[idx : idx + offsets_len] - split_index = offsets[0] + 1 - offset_keys = offsets[1:split_index] - offset_values = offsets[split_index:] - parsed.append([offset_keys, offset_values]) - idx += offsets_len - return parsed - - class Conf: def __init__(self, infile_or_dict: Optional[Union[dict[str, str], IO]] = None) -> None: self._conf: dict[str, Any] = {} From 9e246fc97455639cee07dc2f5472424adf8d4c47 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:36:12 +0000 Subject: [PATCH 03/20] refactor: keep extract_key_value_indexes in worker_util, relax its import guard worker_util backs both the python worker and the eval handlers, so relax the SPARK_TESTING import guard to allow the handler import path (imported on the driver during test collection) instead of relocating the helper to _base. Co-authored-by: Isaac --- python/pyspark/eval_handlers/_arrow.py | 8 ++--- python/pyspark/eval_handlers/_base.py | 30 ------------------ python/pyspark/worker.py | 3 +- python/pyspark/worker_util.py | 44 ++++++++++++++++++++++++-- 4 files changed, 47 insertions(+), 38 deletions(-) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 952945a16d21c..e1d3c82851fef 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -34,7 +34,6 @@ BatchEvalTypeHandler, CoGroupedEvalTypeHandler, GroupedEvalTypeHandler, - extract_key_value_indexes, ) from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, @@ -47,11 +46,12 @@ from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type from pyspark.sql.types import DataType, StructField, StructType from pyspark.util import PythonEvalType +from pyspark.worker_util import extract_key_value_indexes if TYPE_CHECKING: - # Annotation-only: these live in worker-only or forward-ref modules, so they are - # not imported at runtime. ``from __future__ import annotations`` keeps every - # annotation below an unevaluated name rather than a string literal. + # Annotation-only, so they are not imported at runtime. ``from __future__ import + # annotations`` keeps every annotation below an unevaluated name rather than a + # string literal. from pyspark.eval_handlers._typing import CoGroupedBatch, GroupedBatch from pyspark.worker_util import EvalConf, RunnerConf diff --git a/python/pyspark/eval_handlers/_base.py b/python/pyspark/eval_handlers/_base.py index c83b848996e58..ce8003d7aabeb 100644 --- a/python/pyspark/eval_handlers/_base.py +++ b/python/pyspark/eval_handlers/_base.py @@ -52,36 +52,6 @@ def get_eval_type_handler(eval_type: int) -> "Optional[type[EvalTypeHandler]]": return _eval_type_handlers.get(eval_type) -def extract_key_value_indexes(grouped_arg_offsets: list) -> list: - """ - Extract the key and value indexes from arg_offsets for the grouped and cogrouped grouped-map - udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. - - Parameters - ---------- - grouped_arg_offsets: list - List containing the key and value indexes of columns of the - DataFrames to be passed to the udf. It consists of n repeating groups where n is the - number of DataFrames. Each group has the following format: - group[0]: length of group - group[1]: length of key indexes - group[2.. group[1] +2]: key attributes - group[group[1] +3 group[0]]: value attributes - """ - parsed = [] - idx = 0 - while idx < len(grouped_arg_offsets): - offsets_len = grouped_arg_offsets[idx] - idx += 1 - offsets = grouped_arg_offsets[idx : idx + offsets_len] - split_index = offsets[0] + 1 - offset_keys = offsets[1:split_index] - offset_values = offsets[split_index:] - parsed.append([offset_keys, offset_values]) - idx += offsets_len - return parsed - - class _EvalTypeHandlerMeta(ABCMeta): """Registers a concrete handler under its ``eval_type`` at class definition. diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 2cc24176f2917..2b5dcc18571bc 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -52,7 +52,7 @@ _deserialize_accumulator, ) from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError -from pyspark.eval_handlers._base import extract_key_value_indexes, get_eval_type_handler +from pyspark.eval_handlers._base import get_eval_type_handler from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, @@ -116,6 +116,7 @@ EvalConf, RunnerConf, check_python_version, + extract_key_value_indexes, get_sock_file_to_executor, pickleSer, read_command, diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py index 3a24bde38116b..79503b8e32b09 100644 --- a/python/pyspark/worker_util.py +++ b/python/pyspark/worker_util.py @@ -25,14 +25,22 @@ import sys import warnings from contextlib import contextmanager -from inspect import currentframe, getframeinfo +from inspect import currentframe, getframeinfo, stack from typing import IO, Any, Generator, Optional, Union, overload from pyspark.messages import ZeroCopyByteStream if "SPARK_TESTING" in os.environ: - assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER", ( - "This module can only be imported in python woker" + # worker_util backs both the python worker and the eval handlers, so allow it to be + # imported from either. The handler package is imported on the driver too (e.g. when a + # test collects it), where SPARK_PYTHON_RUNTIME is not set; other driver-side imports + # of this worker module during tests are still caught. + _imported_from_handler = any( + frame.frame.f_globals.get("__name__", "").startswith("pyspark.eval_handlers") + for frame in stack() + ) + assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER" or _imported_from_handler, ( + "This module can only be imported in a python worker or an eval handler" ) # 'resource' is a Unix specific module. @@ -265,6 +273,36 @@ def send_accumulator_updates(outfile: IO) -> None: pickleSer._write_with_length((aid, accum._value), outfile) +def extract_key_value_indexes(grouped_arg_offsets: list) -> list: + """ + Extract the key and value indexes from arg_offsets for the grouped and cogrouped grouped-map + udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. + + Parameters + ---------- + grouped_arg_offsets: list + List containing the key and value indexes of columns of the + DataFrames to be passed to the udf. It consists of n repeating groups where n is the + number of DataFrames. Each group has the following format: + group[0]: length of group + group[1]: length of key indexes + group[2.. group[1] +2]: key attributes + group[group[1] +3 group[0]]: value attributes + """ + parsed = [] + idx = 0 + while idx < len(grouped_arg_offsets): + offsets_len = grouped_arg_offsets[idx] + idx += 1 + offsets = grouped_arg_offsets[idx : idx + offsets_len] + split_index = offsets[0] + 1 + offset_keys = offsets[1:split_index] + offset_values = offsets[split_index:] + parsed.append([offset_keys, offset_values]) + idx += offsets_len + return parsed + + class Conf: def __init__(self, infile_or_dict: Optional[Union[dict[str, str], IO]] = None) -> None: self._conf: dict[str, Any] = {} From ed98d6099bc164e21d803b0bfb6f074f95e617e6 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:20:07 +0000 Subject: [PATCH 04/20] refactor: use have_pyarrow to guard the _arrow handler import Co-authored-by: Isaac --- python/pyspark/eval_handlers/__init__.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/python/pyspark/eval_handlers/__init__.py b/python/pyspark/eval_handlers/__init__.py index e2490c83bae27..466051f5579ad 100644 --- a/python/pyspark/eval_handlers/__init__.py +++ b/python/pyspark/eval_handlers/__init__.py @@ -28,9 +28,7 @@ when pyarrow is available; the Arrow eval types it serves cannot run without it. """ -try: - import pyarrow # noqa: F401 -except ImportError: - pass -else: +from pyspark.testing.utils import have_pyarrow + +if have_pyarrow: from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import From 6c21aefcbd99b47d6fa64f189671f30702eb7e4a Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:26:54 +0000 Subject: [PATCH 05/20] refactor: drop _arrow_return_schema, reuse to_arrow_schema for grouped-map output The declared struct return type maps to the same pa.Schema via the existing to_arrow_schema util, so the local helper is unnecessary. Co-authored-by: Isaac --- python/pyspark/eval_handlers/_arrow.py | 26 +++++++------------------- 1 file changed, 7 insertions(+), 19 deletions(-) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index e1d3c82851fef..1cb71ef786f12 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -44,7 +44,7 @@ ) from pyspark.sql.conversion import ArrowBatchTransformer from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type -from pyspark.sql.types import DataType, StructField, StructType +from pyspark.sql.types import StructField, StructType from pyspark.util import PythonEvalType from pyspark.worker_util import extract_key_value_indexes @@ -56,18 +56,6 @@ from pyspark.worker_util import EvalConf, RunnerConf -def _arrow_return_schema(return_type: DataType, use_large_var_types: bool) -> pa.Schema: - """Arrow schema for a grouped/cogrouped map UDF's declared struct return type. - - The return type is a StructType, so ``to_arrow_type`` yields a struct type whose - fields are the output columns; the group's output batches carry those fields flat. - """ - arrow_return_type = to_arrow_type( - return_type, timezone="UTC", prefers_large_types=use_large_var_types - ) - return pa.schema(list(arrow_return_type)) - - class ArrowScalarUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): """SQL_SCALAR_ARROW_UDF: invoke each UDF once per input RecordBatch, coerce the result to the declared schema, and check the row count.""" @@ -215,8 +203,8 @@ def __init__( assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_ARROW UDF." self._key_offsets = parsed_offsets[0][0] self._value_offsets = parsed_offsets[0][1] - self._arrow_return_schema = _arrow_return_schema( - return_type, runner_conf.use_large_var_types + self._arrow_return_schema = to_arrow_schema( + return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types ) def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: @@ -281,8 +269,8 @@ def __init__( ) self._key_offsets = parsed_offsets[0][0] self._value_offsets = parsed_offsets[0][1] - self._arrow_return_schema = _arrow_return_schema( - return_type, runner_conf.use_large_var_types + self._arrow_return_schema = to_arrow_schema( + return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types ) def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: @@ -345,8 +333,8 @@ def __init__( parsed_offsets = extract_key_value_indexes(arg_offsets) self._left_key_cols, self._left_val_cols = parsed_offsets[0] self._right_key_cols, self._right_val_cols = parsed_offsets[1] - self._arrow_return_schema = _arrow_return_schema( - return_type, runner_conf.use_large_var_types + self._arrow_return_schema = to_arrow_schema( + return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types ) def run(self, split_index: int, data: Iterator[CoGroupedBatch]) -> Iterator[pa.RecordBatch]: From b96e1dff8d6200ab3efd24a6f30a360e3b6b5c83 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:31:51 +0000 Subject: [PATCH 06/20] refactor: order verification helpers and _arrow handlers alphabetically Move the private _top_level_package to the top of verification.py and order the verify_* functions alphabetically; order the handler classes in _arrow.py the same way. Co-authored-by: Isaac --- python/pyspark/eval_handlers/_arrow.py | 382 +++++++++---------- python/pyspark/eval_handlers/verification.py | 134 +++---- 2 files changed, 258 insertions(+), 258 deletions(-) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 1cb71ef786f12..22838c5d66c67 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -56,186 +56,45 @@ from pyspark.worker_util import EvalConf, RunnerConf -class ArrowScalarUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): - """SQL_SCALAR_ARROW_UDF: invoke each UDF once per input RecordBatch, coerce - the result to the declared schema, and check the row count.""" - - eval_type = PythonEvalType.SQL_SCALAR_ARROW_UDF - - def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf - ) -> None: - super().__init__(udfs, runner_conf, eval_conf) - self._col_names = ["_%d" % i for i in range(len(udfs))] - self._combined_arrow_schema = to_arrow_schema( - StructType([StructField(n, rt) for n, (_, _, _, rt) in zip(self._col_names, udfs)]), - timezone="UTC", - prefers_large_types=runner_conf.use_large_var_types, - ) - - def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: - for batch in data: - output_batch = pa.RecordBatch.from_arrays( - [ - udf_func( - *[batch.column(o) for o in args_offsets], - **{k: batch.column(v) for k, v in kwargs_offsets.items()}, - ) - for udf_func, args_offsets, kwargs_offsets, _ in self._udfs - ], - self._col_names, - ) - output_batch = ArrowBatchTransformer.enforce_schema( - output_batch, self._combined_arrow_schema - ) - verify_scalar_result(output_batch, batch.num_rows) - yield output_batch - - -class ArrowScalarIterUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): - """SQL_SCALAR_ARROW_ITER_UDF: the UDF receives an iterator of the argument - columns and yields an iterator of pa.Array; enforce the declared type on each - result and verify the total row count matches the input.""" - - eval_type = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF - - def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf - ) -> None: - super().__init__(udfs, runner_conf, eval_conf) - assert len(udfs) == 1, "One SCALAR_ARROW_ITER UDF expected here." - self._udf_func, self._args_offsets, _, return_type = udfs[0] - self._arrow_return_type = to_arrow_type( - return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types - ) - - def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: - args_offsets = self._args_offsets - num_input_rows = 0 - - def extract_args(batch: pa.RecordBatch) -> Any: - nonlocal num_input_rows - args = tuple(batch.column(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 pa.Array) - verified_iter = verify_return_type(self._udf_func(args_iter), Iterator[pa.Array]) - - # Process results: enforce schema and assemble into RecordBatch - target_schema = pa.schema([pa.field("_0", self._arrow_return_type)]) - - def process_results() -> Iterator[pa.RecordBatch]: - for result in verified_iter: - batch = pa.RecordBatch.from_arrays([result], ["_0"]) - yield ArrowBatchTransformer.enforce_schema(batch, target_schema, safecheck=True) - - # 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) - - -class ArrowMapUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): - """SQL_MAP_ARROW_ITER_UDF (mapInArrow): the single UDF receives the input - RecordBatch stream and yields a RecordBatch stream, exchanged as flattened - columns on the wire and wrapped back into a single struct column.""" - - eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF - - def __init__( - self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf - ) -> None: - super().__init__(udfs, runner_conf, eval_conf) - assert len(udfs) == 1, "One MAP_ARROW_ITER UDF expected here." - self._udf_func = udfs[0][0] - - def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: - # Pre-processing - input_batches = map(ArrowBatchTransformer.flatten_struct, data) - - # invoke the UDF - output_batches = self._udf_func(input_batches) - - # 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( - output_batches, Iterator - ): - try: - output_batches = iter(output_batches) - except TypeError: - # Not iterable at all; leave it so verify_return_type below raises the - # standard UDF_RETURN_TYPE error. - pass - - # Post-processing - verified_iter = verify_return_type(output_batches, Iterator[pa.RecordBatch]) - yield from map(ArrowBatchTransformer.wrap_struct, verified_iter) - - -class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): - """SQL_GROUPED_MAP_ARROW_UDF (applyInArrow): the single UDF receives each - group as one pa.Table and returns one pa.Table, coerced to the declared - schema.""" +class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler[pa.RecordBatch]): + """SQL_COGROUPED_MAP_ARROW_UDF (applyInArrow on a cogroup): the single UDF + receives the two sides' value tables and returns one pa.Table, coerced to the + declared schema.""" - eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF + eval_type = PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) - assert len(udfs) == 1, "One GROUPED_MAP_ARROW UDF expected here." - self._grouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] + assert len(udfs) == 1, "One COGROUPED_MAP_ARROW UDF expected here." + self._cogrouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] parsed_offsets = extract_key_value_indexes(arg_offsets) - assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_ARROW UDF." - self._key_offsets = parsed_offsets[0][0] - self._value_offsets = parsed_offsets[0][1] + self._left_key_cols, self._left_val_cols = parsed_offsets[0] + self._right_key_cols, self._right_val_cols = parsed_offsets[1] self._arrow_return_schema = to_arrow_schema( return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types ) - def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: - """Apply groupBy Arrow UDF (non-iterator variant).""" - key_offsets = self._key_offsets - value_offsets = self._value_offsets - for group in data: - # Flatten struct column into separate columns - flattened = map(ArrowBatchTransformer.flatten_struct, group) + def run(self, split_index: int, data: Iterator[CoGroupedBatch]) -> Iterator[pa.RecordBatch]: + """Apply cogroupBy Arrow UDF.""" + select_columns = ArrowBatchTransformer.select_columns - # Materialize first batch to get keys - first_batch = next(flattened) - keys = pa.RecordBatch.from_arrays( - [first_batch.columns[o] for o in key_offsets], - [first_batch.schema.names[o] for o in key_offsets], - ) - value_batches = ( - pa.RecordBatch.from_arrays( - [b.columns[o] for o in value_offsets], - [b.schema.names[o] for o in value_offsets], - ) - for b in itertools.chain((first_batch,), flattened) - ) + def table_from_batches(batches: list[pa.RecordBatch], cols: list[int]) -> pa.Table: + return pa.Table.from_batches([select_columns(b, cols) for b in batches]) - # Call UDF - value_table = pa.Table.from_batches(value_batches) - if self._num_udf_args == 1: - result = self._grouped_udf(value_table) + for left_batches, right_batches in data: + left_keys = table_from_batches(left_batches, self._left_key_cols) + left_values = table_from_batches(left_batches, self._left_val_cols) + right_keys = table_from_batches(right_batches, self._right_key_cols) + right_values = table_from_batches(right_batches, self._right_val_cols) + + if self._num_udf_args == 2: + result = self._cogrouped_udf(left_values, right_values) else: - key = tuple(c[0] for c in keys.columns) - result = self._grouped_udf(key, value_table) + key_table = left_keys if left_keys.num_rows > 0 else right_keys + key = tuple(c[0] for c in key_table.columns) + result = self._cogrouped_udf(key, left_values, right_values) verify_return_type(result, pa.Table) # Verify types (and reorder by name when configured). @@ -317,45 +176,56 @@ def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.Rec pass -class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler[pa.RecordBatch]): - """SQL_COGROUPED_MAP_ARROW_UDF (applyInArrow on a cogroup): the single UDF - receives the two sides' value tables and returns one pa.Table, coerced to the - declared schema.""" +class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): + """SQL_GROUPED_MAP_ARROW_UDF (applyInArrow): the single UDF receives each + group as one pa.Table and returns one pa.Table, coerced to the declared + schema.""" - eval_type = PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF + eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: super().__init__(udfs, runner_conf, eval_conf) - assert len(udfs) == 1, "One COGROUPED_MAP_ARROW UDF expected here." - self._cogrouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] + assert len(udfs) == 1, "One GROUPED_MAP_ARROW UDF expected here." + self._grouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] parsed_offsets = extract_key_value_indexes(arg_offsets) - self._left_key_cols, self._left_val_cols = parsed_offsets[0] - self._right_key_cols, self._right_val_cols = parsed_offsets[1] + assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_ARROW UDF." + self._key_offsets = parsed_offsets[0][0] + self._value_offsets = parsed_offsets[0][1] self._arrow_return_schema = to_arrow_schema( return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types ) - def run(self, split_index: int, data: Iterator[CoGroupedBatch]) -> Iterator[pa.RecordBatch]: - """Apply cogroupBy Arrow UDF.""" - select_columns = ArrowBatchTransformer.select_columns - - def table_from_batches(batches: list[pa.RecordBatch], cols: list[int]) -> pa.Table: - return pa.Table.from_batches([select_columns(b, cols) for b in batches]) + def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: + """Apply groupBy Arrow UDF (non-iterator variant).""" + key_offsets = self._key_offsets + value_offsets = self._value_offsets + for group in data: + # Flatten struct column into separate columns + flattened = map(ArrowBatchTransformer.flatten_struct, group) - for left_batches, right_batches in data: - left_keys = table_from_batches(left_batches, self._left_key_cols) - left_values = table_from_batches(left_batches, self._left_val_cols) - right_keys = table_from_batches(right_batches, self._right_key_cols) - right_values = table_from_batches(right_batches, self._right_val_cols) + # Materialize first batch to get keys + first_batch = next(flattened) + keys = pa.RecordBatch.from_arrays( + [first_batch.columns[o] for o in key_offsets], + [first_batch.schema.names[o] for o in key_offsets], + ) + value_batches = ( + pa.RecordBatch.from_arrays( + [b.columns[o] for o in value_offsets], + [b.schema.names[o] for o in value_offsets], + ) + for b in itertools.chain((first_batch,), flattened) + ) - if self._num_udf_args == 2: - result = self._cogrouped_udf(left_values, right_values) + # Call UDF + value_table = pa.Table.from_batches(value_batches) + if self._num_udf_args == 1: + result = self._grouped_udf(value_table) else: - key_table = left_keys if left_keys.num_rows > 0 else right_keys - key = tuple(c[0] for c in key_table.columns) - result = self._cogrouped_udf(key, left_values, right_values) + key = tuple(c[0] for c in keys.columns) + result = self._grouped_udf(key, value_table) verify_return_type(result, pa.Table) # Verify types (and reorder by name when configured). @@ -368,3 +238,133 @@ def table_from_batches(batches: list[pa.RecordBatch], cols: list[int]) -> pa.Tab for batch in result.to_batches(): yield ArrowBatchTransformer.wrap_struct(batch) + + +class ArrowMapUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): + """SQL_MAP_ARROW_ITER_UDF (mapInArrow): the single UDF receives the input + RecordBatch stream and yields a RecordBatch stream, exchanged as flattened + columns on the wire and wrapped back into a single struct column.""" + + eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + assert len(udfs) == 1, "One MAP_ARROW_ITER UDF expected here." + self._udf_func = udfs[0][0] + + def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + # Pre-processing + input_batches = map(ArrowBatchTransformer.flatten_struct, data) + + # invoke the UDF + output_batches = self._udf_func(input_batches) + + # 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( + output_batches, Iterator + ): + try: + output_batches = iter(output_batches) + except TypeError: + # Not iterable at all; leave it so verify_return_type below raises the + # standard UDF_RETURN_TYPE error. + pass + + # Post-processing + verified_iter = verify_return_type(output_batches, Iterator[pa.RecordBatch]) + yield from map(ArrowBatchTransformer.wrap_struct, verified_iter) + + +class ArrowScalarIterUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): + """SQL_SCALAR_ARROW_ITER_UDF: the UDF receives an iterator of the argument + columns and yields an iterator of pa.Array; enforce the declared type on each + result and verify the total row count matches the input.""" + + eval_type = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + assert len(udfs) == 1, "One SCALAR_ARROW_ITER UDF expected here." + self._udf_func, self._args_offsets, _, return_type = udfs[0] + self._arrow_return_type = to_arrow_type( + return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types + ) + + def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + args_offsets = self._args_offsets + num_input_rows = 0 + + def extract_args(batch: pa.RecordBatch) -> Any: + nonlocal num_input_rows + args = tuple(batch.column(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 pa.Array) + verified_iter = verify_return_type(self._udf_func(args_iter), Iterator[pa.Array]) + + # Process results: enforce schema and assemble into RecordBatch + target_schema = pa.schema([pa.field("_0", self._arrow_return_type)]) + + def process_results() -> Iterator[pa.RecordBatch]: + for result in verified_iter: + batch = pa.RecordBatch.from_arrays([result], ["_0"]) + yield ArrowBatchTransformer.enforce_schema(batch, target_schema, safecheck=True) + + # 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) + + +class ArrowScalarUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): + """SQL_SCALAR_ARROW_UDF: invoke each UDF once per input RecordBatch, coerce + the result to the declared schema, and check the row count.""" + + eval_type = PythonEvalType.SQL_SCALAR_ARROW_UDF + + def __init__( + self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf + ) -> None: + super().__init__(udfs, runner_conf, eval_conf) + self._col_names = ["_%d" % i for i in range(len(udfs))] + self._combined_arrow_schema = to_arrow_schema( + StructType([StructField(n, rt) for n, (_, _, _, rt) in zip(self._col_names, udfs)]), + timezone="UTC", + prefers_large_types=runner_conf.use_large_var_types, + ) + + def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + for batch in data: + output_batch = pa.RecordBatch.from_arrays( + [ + udf_func( + *[batch.column(o) for o in args_offsets], + **{k: batch.column(v) for k, v in kwargs_offsets.items()}, + ) + for udf_func, args_offsets, kwargs_offsets, _ in self._udfs + ], + self._col_names, + ) + output_batch = ArrowBatchTransformer.enforce_schema( + output_batch, self._combined_arrow_schema + ) + verify_scalar_result(output_batch, batch.num_rows) + yield output_batch diff --git a/python/pyspark/eval_handlers/verification.py b/python/pyspark/eval_handlers/verification.py index c77096c622d9a..3b8399984dee7 100644 --- a/python/pyspark/eval_handlers/verification.py +++ b/python/pyspark/eval_handlers/verification.py @@ -28,6 +28,52 @@ T = TypeVar("T") +def _top_level_package(t: type) -> str: + """Return the top-level package of ``t`` (``pandas`` for ``pd.DataFrame``).""" + return (t.__module__ or "").split(".", 1)[0] + + +def verify_iter_result_row_count( + iterator: Iterator, + expected_rows: Callable[[], int], +) -> Iterator: + """Yield elements and verify final row count matches expected exactly. + + ``expected_rows`` is a callable because the expected count is only known once + the iterator is fully consumed (input rows are counted lazily as a side effect + of pulling batches), so it must be read after this generator is exhausted. + """ + actual_rows = 0 + for element in iterator: + actual_rows += len(element) + yield element + + verify_result_row_count(actual_rows, expected_rows()) + + +def verify_iterator_exhausted(iterator: Iterator) -> None: + """Verify that an iterator has been fully consumed.""" + try: + next(iterator) + except StopIteration: + pass + else: + raise PySparkRuntimeError(errorClass="INPUT_NOT_FULLY_CONSUMED", messageParameters={}) + + +def verify_output_row_limit( + iterator: Iterator, + max_rows: Union[int, Callable[[], int]], +) -> Iterator: + """Yield elements while verifying total rows do not exceed a limit (fail-fast).""" + total_rows = 0 + for element in iterator: + total_rows += len(element) + if total_rows > (max_rows() if callable(max_rows) else max_rows): + raise PySparkRuntimeError(errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={}) + yield element + + def verify_result_row_count(result_length: int, expected: int) -> None: """Raise if the result row count doesn't match the expected input row count.""" if result_length != expected: @@ -40,36 +86,6 @@ def verify_result_row_count(result_length: int, expected: int) -> None: ) -def verify_scalar_result(result: Any, num_rows: int) -> Any: - """ - Verify a scalar UDF result is array-like and has the expected number of rows. - - Parameters - ---------- - result : Any - The UDF result to verify. - num_rows : int - Expected number of rows (must match input batch size). - """ - try: - result_length = len(result) - except TypeError: - raise PySparkTypeError( - errorClass="UDF_RETURN_TYPE", - messageParameters={ - "expected": "array-like object", - "actual": type(result).__name__, - }, - ) - verify_result_row_count(result_length, num_rows) - return result - - -def _top_level_package(t: type) -> str: - """Return the top-level package of ``t`` (``pandas`` for ``pd.DataFrame``).""" - return (t.__module__ or "").split(".", 1)[0] - - def verify_return_type(result: T, expected_type: Type[T]) -> T: """ Verify a UDF return value against an expected type. @@ -112,42 +128,26 @@ def check_element(element: T) -> T: return result -def verify_iterator_exhausted(iterator: Iterator) -> None: - """Verify that an iterator has been fully consumed.""" - try: - next(iterator) - except StopIteration: - pass - else: - raise PySparkRuntimeError(errorClass="INPUT_NOT_FULLY_CONSUMED", messageParameters={}) - - -def verify_output_row_limit( - iterator: Iterator, - max_rows: Union[int, Callable[[], int]], -) -> Iterator: - """Yield elements while verifying total rows do not exceed a limit (fail-fast).""" - total_rows = 0 - for element in iterator: - total_rows += len(element) - if total_rows > (max_rows() if callable(max_rows) else max_rows): - raise PySparkRuntimeError(errorClass="OUTPUT_EXCEEDS_INPUT_ROWS", messageParameters={}) - yield element - - -def verify_iter_result_row_count( - iterator: Iterator, - expected_rows: Callable[[], int], -) -> Iterator: - """Yield elements and verify final row count matches expected exactly. - - ``expected_rows`` is a callable because the expected count is only known once - the iterator is fully consumed (input rows are counted lazily as a side effect - of pulling batches), so it must be read after this generator is exhausted. +def verify_scalar_result(result: Any, num_rows: int) -> Any: """ - actual_rows = 0 - for element in iterator: - actual_rows += len(element) - yield element + Verify a scalar UDF result is array-like and has the expected number of rows. - verify_result_row_count(actual_rows, expected_rows()) + Parameters + ---------- + result : Any + The UDF result to verify. + num_rows : int + Expected number of rows (must match input batch size). + """ + try: + result_length = len(result) + except TypeError: + raise PySparkTypeError( + errorClass="UDF_RETURN_TYPE", + messageParameters={ + "expected": "array-like object", + "actual": type(result).__name__, + }, + ) + verify_result_row_count(result_length, num_rows) + return result From cec36aa5a94480489244f208e361c123842bd1bb Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:39:35 +0000 Subject: [PATCH 07/20] refactor: move extract_key_value_indexes to eval_handlers/_util, revert worker_util guard _util is a leaf module (no pyspark imports), safe to import from both the worker and the handlers on driver and executor, so the handlers no longer import the worker-only worker_util and its SPARK_TESTING guard returns to worker-only. Co-authored-by: Isaac --- python/pyspark/eval_handlers/_arrow.py | 2 +- python/pyspark/eval_handlers/_util.py | 52 ++++++++++++++++++++++++++ python/pyspark/worker.py | 2 +- python/pyspark/worker_util.py | 44 ++-------------------- 4 files changed, 57 insertions(+), 43 deletions(-) create mode 100644 python/pyspark/eval_handlers/_util.py diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 22838c5d66c67..14c2d2ce967ef 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -35,6 +35,7 @@ CoGroupedEvalTypeHandler, GroupedEvalTypeHandler, ) +from pyspark.eval_handlers._util import extract_key_value_indexes from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, @@ -46,7 +47,6 @@ from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type from pyspark.sql.types import StructField, StructType from pyspark.util import PythonEvalType -from pyspark.worker_util import extract_key_value_indexes if TYPE_CHECKING: # Annotation-only, so they are not imported at runtime. ``from __future__ import diff --git a/python/pyspark/eval_handlers/_util.py b/python/pyspark/eval_handlers/_util.py new file mode 100644 index 0000000000000..48a7570203374 --- /dev/null +++ b/python/pyspark/eval_handlers/_util.py @@ -0,0 +1,52 @@ +# +# 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. +# + +"""Shared helpers for the eval type handlers. + +A leaf module (no pyspark imports), so it is safe to import from both the worker +and the handlers, on the driver and the executor alike. +""" + + +def extract_key_value_indexes(grouped_arg_offsets: list) -> list: + """ + Extract the key and value indexes from arg_offsets for the grouped and cogrouped grouped-map + udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. + + Parameters + ---------- + grouped_arg_offsets: list + List containing the key and value indexes of columns of the + DataFrames to be passed to the udf. It consists of n repeating groups where n is the + number of DataFrames. Each group has the following format: + group[0]: length of group + group[1]: length of key indexes + group[2.. group[1] +2]: key attributes + group[group[1] +3 group[0]]: value attributes + """ + parsed = [] + idx = 0 + while idx < len(grouped_arg_offsets): + offsets_len = grouped_arg_offsets[idx] + idx += 1 + offsets = grouped_arg_offsets[idx : idx + offsets_len] + split_index = offsets[0] + 1 + offset_keys = offsets[1:split_index] + offset_values = offsets[split_index:] + parsed.append([offset_keys, offset_values]) + idx += offsets_len + return parsed diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 2b5dcc18571bc..486c42e691754 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -53,6 +53,7 @@ ) from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.eval_handlers._base import get_eval_type_handler +from pyspark.eval_handlers._util import extract_key_value_indexes from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, @@ -116,7 +117,6 @@ EvalConf, RunnerConf, check_python_version, - extract_key_value_indexes, get_sock_file_to_executor, pickleSer, read_command, diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py index 79503b8e32b09..3a24bde38116b 100644 --- a/python/pyspark/worker_util.py +++ b/python/pyspark/worker_util.py @@ -25,22 +25,14 @@ import sys import warnings from contextlib import contextmanager -from inspect import currentframe, getframeinfo, stack +from inspect import currentframe, getframeinfo from typing import IO, Any, Generator, Optional, Union, overload from pyspark.messages import ZeroCopyByteStream if "SPARK_TESTING" in os.environ: - # worker_util backs both the python worker and the eval handlers, so allow it to be - # imported from either. The handler package is imported on the driver too (e.g. when a - # test collects it), where SPARK_PYTHON_RUNTIME is not set; other driver-side imports - # of this worker module during tests are still caught. - _imported_from_handler = any( - frame.frame.f_globals.get("__name__", "").startswith("pyspark.eval_handlers") - for frame in stack() - ) - assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER" or _imported_from_handler, ( - "This module can only be imported in a python worker or an eval handler" + assert os.environ.get("SPARK_PYTHON_RUNTIME") == "PYTHON_WORKER", ( + "This module can only be imported in python woker" ) # 'resource' is a Unix specific module. @@ -273,36 +265,6 @@ def send_accumulator_updates(outfile: IO) -> None: pickleSer._write_with_length((aid, accum._value), outfile) -def extract_key_value_indexes(grouped_arg_offsets: list) -> list: - """ - Extract the key and value indexes from arg_offsets for the grouped and cogrouped grouped-map - udfs. See BasePandasGroupExec.resolveArgOffsets for equivalent scala code. - - Parameters - ---------- - grouped_arg_offsets: list - List containing the key and value indexes of columns of the - DataFrames to be passed to the udf. It consists of n repeating groups where n is the - number of DataFrames. Each group has the following format: - group[0]: length of group - group[1]: length of key indexes - group[2.. group[1] +2]: key attributes - group[group[1] +3 group[0]]: value attributes - """ - parsed = [] - idx = 0 - while idx < len(grouped_arg_offsets): - offsets_len = grouped_arg_offsets[idx] - idx += 1 - offsets = grouped_arg_offsets[idx : idx + offsets_len] - split_index = offsets[0] + 1 - offset_keys = offsets[1:split_index] - offset_values = offsets[split_index:] - parsed.append([offset_keys, offset_values]) - idx += offsets_len - return parsed - - class Conf: def __init__(self, infile_or_dict: Optional[Union[dict[str, str], IO]] = None) -> None: self._conf: dict[str, Any] = {} From e7c02b488d1c35a8beaefbc5f1e355b88dd367a5 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:45:21 +0000 Subject: [PATCH 08/20] test: split handler tests into test_base_ and test_arrow_ eval type handlers Mirror the source layout: framework (_base) tests in test_base_eval_type_handlers and Arrow-flavor tests in test_arrow_eval_type_handlers, leaving room for future flavors (e.g. pandas). The framework tests no longer depend on the Arrow handlers. Co-authored-by: Isaac --- dev/sparktestsupport/modules.py | 3 +- ...rs.py => test_arrow_eval_type_handlers.py} | 170 +++++------------- .../tests/test_base_eval_type_handlers.py | 145 +++++++++++++++ 3 files changed, 187 insertions(+), 131 deletions(-) rename python/pyspark/tests/{test_eval_type_handlers.py => test_arrow_eval_type_handlers.py} (78%) create mode 100644 python/pyspark/tests/test_base_eval_type_handlers.py diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 61f3e4c155705..74a9aee355a18 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -612,7 +612,8 @@ def __hash__(self): "pyspark.tests.test_taskcontext", "pyspark.tests.test_util", "pyspark.tests.test_worker", - "pyspark.tests.test_eval_type_handlers", + "pyspark.tests.test_base_eval_type_handlers", + "pyspark.tests.test_arrow_eval_type_handlers", "pyspark.tests.test_stage_sched", "pyspark.tests.test_zero_copy_byte_stream", # unittests for upstream projects diff --git a/python/pyspark/tests/test_eval_type_handlers.py b/python/pyspark/tests/test_arrow_eval_type_handlers.py similarity index 78% rename from python/pyspark/tests/test_eval_type_handlers.py rename to python/pyspark/tests/test_arrow_eval_type_handlers.py index ff0b3f821f16d..807b45bae6518 100644 --- a/python/pyspark/tests/test_eval_type_handlers.py +++ b/python/pyspark/tests/test_arrow_eval_type_handlers.py @@ -15,21 +15,12 @@ # limitations under the License. # +"""Tests for the Arrow eval type handlers (``_arrow``).""" + import unittest -from pyspark.eval_handlers._base import ( - BatchEvalTypeHandler, - CoGroupedEvalTypeHandler, - EvalTypeHandler, - GroupedEvalTypeHandler, - _eval_type_handlers, - get_eval_type_handler, -) -from pyspark.sql.pandas.serializers import ( - ArrowStreamCoGroupSerializer, - ArrowStreamGroupSerializer, - ArrowStreamSerializer, -) +from pyspark.eval_handlers._base import get_eval_type_handler +from pyspark.sql.pandas.serializers import ArrowStreamCoGroupSerializer, ArrowStreamSerializer from pyspark.sql.types import LongType, StructField, StructType from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message from pyspark.util import PythonEvalType @@ -56,7 +47,7 @@ class _RunnerConf: @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) -class EvalTypeHandlerTests(unittest.TestCase): +class ArrowEvalTypeHandlerRegistrationTests(unittest.TestCase): def test_arrow_eval_types_are_registered(self): # Every migrated Arrow map/iter eval type dispatches to its handler by lookup. for eval_type, handler_cls in ( @@ -69,87 +60,6 @@ def test_arrow_eval_types_are_registered(self): ): self.assertIs(get_eval_type_handler(eval_type), handler_cls) - def test_category_bases_are_abstract(self): - # The interface and the three category bases must not be instantiable: - # they leave ``run`` abstract. - for base in ( - EvalTypeHandler, - BatchEvalTypeHandler, - GroupedEvalTypeHandler, - CoGroupedEvalTypeHandler, - ): - with self.assertRaises(TypeError): - base([], _RunnerConf(), None) - - def test_category_bases_are_not_registered(self): - # Only concrete subclasses that declare an eval type are registered. - registered = set(_eval_type_handlers.values()) - for base in ( - EvalTypeHandler, - BatchEvalTypeHandler, - GroupedEvalTypeHandler, - CoGroupedEvalTypeHandler, - ): - self.assertNotIn(base, registered) - - def test_default_serializer_per_category(self): - class _Batch(BatchEvalTypeHandler["pa.RecordBatch"]): - def run(self, split_index, data): - return data - - class _Grouped(GroupedEvalTypeHandler["pa.RecordBatch"]): - def run(self, split_index, data): - return data - - class _CoGrouped(CoGroupedEvalTypeHandler["pa.RecordBatch"]): - def run(self, split_index, data): - return data - - self.assertIsInstance(_Batch([], _RunnerConf(), None).serializer, ArrowStreamSerializer) - self.assertIsInstance( - _Grouped([], _RunnerConf(), None).serializer, ArrowStreamGroupSerializer - ) - self.assertIsInstance( - _CoGrouped([], _RunnerConf(), None).serializer, ArrowStreamCoGroupSerializer - ) - - def test_run_produces_output(self): - class _Doubler(BatchEvalTypeHandler["pa.RecordBatch"]): - def run(self, split_index, data): - for item in data: - yield item * 2 - - handler = _Doubler([], _RunnerConf(), None) - self.assertEqual(list(handler.run(0, iter([1, 2, 3]))), [2, 4, 6]) - - def test_duplicate_eval_type_rejected(self): - def _define_duplicate(): - class _Dup(BatchEvalTypeHandler["pa.RecordBatch"]): - eval_type = PythonEvalType.SQL_SCALAR_ARROW_UDF - - def run(self, split_index, data): - return data - - self.assertRaises(AssertionError, _define_duplicate) - # The failed definition must not clobber the existing registration. - self.assertIs( - get_eval_type_handler(PythonEvalType.SQL_SCALAR_ARROW_UDF), - ArrowScalarUDFHandler, - ) - - def test_abstract_handler_with_eval_type_rejected(self): - # A subclass that declares an eval_type but leaves run abstract must be - # rejected at class definition. - unused_eval_type = -1 - - def _define_abstract(): - class _Abstract(BatchEvalTypeHandler["pa.RecordBatch"]): - eval_type = unused_eval_type - # run left abstract on purpose - - self.assertRaises(AssertionError, _define_abstract) - self.assertNotIn(unused_eval_type, _eval_type_handlers) - @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowScalarUDFHandlerTests(unittest.TestCase): @@ -188,41 +98,6 @@ def add_one(col): self.assertEqual(out[0].column(0).to_pylist(), [11, 21]) -@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) -class CoGroupedBatchTests(unittest.TestCase): - def test_deserialized_co_group_is_a_pair_of_lists(self): - # CoGroupedBatch must match what ArrowStreamCoGroupSerializer yields: the - # serializer eagerly materializes each side as a list, not an iterator. - import io - - import pyarrow as pa - - from pyspark.serializers import write_int - - def arrow_bytes(batches): - buf = io.BytesIO() - ArrowStreamSerializer().dump_stream(iter(batches), buf) - return buf.getvalue() - - left = pa.RecordBatch.from_arrays([pa.array([1, 2])], ["_0"]) - right = pa.RecordBatch.from_arrays([pa.array([9])], ["_0"]) - - stream = io.BytesIO() - write_int(2, stream) # two DataFrames in the co-group - stream.write(arrow_bytes([left])) - stream.write(arrow_bytes([right])) - write_int(0, stream) # end of stream - stream.seek(0) - - groups = list(ArrowStreamCoGroupSerializer().load_stream(stream)) - self.assertEqual(len(groups), 1) - left_side, right_side = groups[0] - self.assertIsInstance(left_side, list) - self.assertIsInstance(right_side, list) - self.assertEqual([b.num_rows for b in left_side], [2]) - self.assertEqual([b.num_rows for b in right_side], [1]) - - @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowScalarIterUDFHandlerTests(unittest.TestCase): def test_end_to_end_output(self): @@ -388,6 +263,41 @@ def cogrouped_udf(left_values, right_values): self.assertEqual(out[0].column(0).field("out").to_pylist(), [30]) +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class CoGroupedBatchTests(unittest.TestCase): + def test_deserialized_co_group_is_a_pair_of_lists(self): + # CoGroupedBatch must match what ArrowStreamCoGroupSerializer yields: the + # serializer eagerly materializes each side as a list, not an iterator. + import io + + import pyarrow as pa + + from pyspark.serializers import write_int + + def arrow_bytes(batches): + buf = io.BytesIO() + ArrowStreamSerializer().dump_stream(iter(batches), buf) + return buf.getvalue() + + left = pa.RecordBatch.from_arrays([pa.array([1, 2])], ["_0"]) + right = pa.RecordBatch.from_arrays([pa.array([9])], ["_0"]) + + stream = io.BytesIO() + write_int(2, stream) # two DataFrames in the co-group + stream.write(arrow_bytes([left])) + stream.write(arrow_bytes([right])) + write_int(0, stream) # end of stream + stream.seek(0) + + groups = list(ArrowStreamCoGroupSerializer().load_stream(stream)) + self.assertEqual(len(groups), 1) + left_side, right_side = groups[0] + self.assertIsInstance(left_side, list) + self.assertIsInstance(right_side, list) + self.assertEqual([b.num_rows for b in left_side], [2]) + self.assertEqual([b.num_rows for b in right_side], [1]) + + if __name__ == "__main__": from pyspark.testing import main diff --git a/python/pyspark/tests/test_base_eval_type_handlers.py b/python/pyspark/tests/test_base_eval_type_handlers.py new file mode 100644 index 0000000000000..c5fff82312ac5 --- /dev/null +++ b/python/pyspark/tests/test_base_eval_type_handlers.py @@ -0,0 +1,145 @@ +# +# 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. +# + +"""Framework-level tests for the eval type handler pipeline (``_base``). + +Handler-flavor tests live alongside their module, e.g. the Arrow handlers in +``test_arrow_eval_type_handlers``. +""" + +import unittest + +from pyspark.eval_handlers._base import ( + BatchEvalTypeHandler, + CoGroupedEvalTypeHandler, + EvalTypeHandler, + GroupedEvalTypeHandler, + _eval_type_handlers, + get_eval_type_handler, +) +from pyspark.sql.pandas.serializers import ( + ArrowStreamCoGroupSerializer, + ArrowStreamGroupSerializer, + ArrowStreamSerializer, +) + + +class _RunnerConf: + """Minimal stand-in for the worker's RunnerConf, exposing only the fields + the handlers under test read.""" + + use_large_var_types = False + assign_cols_by_name = True + map_in_batch_legacy_accept_any_iterable = False + + +class EvalTypeHandlerTests(unittest.TestCase): + def test_category_bases_are_abstract(self): + # The interface and the three category bases must not be instantiable: + # they leave ``run`` abstract. + for base in ( + EvalTypeHandler, + BatchEvalTypeHandler, + GroupedEvalTypeHandler, + CoGroupedEvalTypeHandler, + ): + with self.assertRaises(TypeError): + base([], _RunnerConf(), None) + + def test_category_bases_are_not_registered(self): + # Only concrete subclasses that declare an eval type are registered. + registered = set(_eval_type_handlers.values()) + for base in ( + EvalTypeHandler, + BatchEvalTypeHandler, + GroupedEvalTypeHandler, + CoGroupedEvalTypeHandler, + ): + self.assertNotIn(base, registered) + + def test_default_serializer_per_category(self): + class _Batch(BatchEvalTypeHandler["pa.RecordBatch"]): + def run(self, split_index, data): + return data + + class _Grouped(GroupedEvalTypeHandler["pa.RecordBatch"]): + def run(self, split_index, data): + return data + + class _CoGrouped(CoGroupedEvalTypeHandler["pa.RecordBatch"]): + def run(self, split_index, data): + return data + + self.assertIsInstance(_Batch([], _RunnerConf(), None).serializer, ArrowStreamSerializer) + self.assertIsInstance( + _Grouped([], _RunnerConf(), None).serializer, ArrowStreamGroupSerializer + ) + self.assertIsInstance( + _CoGrouped([], _RunnerConf(), None).serializer, ArrowStreamCoGroupSerializer + ) + + def test_run_produces_output(self): + class _Doubler(BatchEvalTypeHandler["pa.RecordBatch"]): + def run(self, split_index, data): + for item in data: + yield item * 2 + + handler = _Doubler([], _RunnerConf(), None) + self.assertEqual(list(handler.run(0, iter([1, 2, 3]))), [2, 4, 6]) + + def test_duplicate_eval_type_rejected(self): + unused_eval_type = -2 + + class _First(BatchEvalTypeHandler["pa.RecordBatch"]): + eval_type = unused_eval_type + + def run(self, split_index, data): + return data + + try: + + def _define_duplicate(): + class _Dup(BatchEvalTypeHandler["pa.RecordBatch"]): + eval_type = unused_eval_type + + def run(self, split_index, data): + return data + + self.assertRaises(AssertionError, _define_duplicate) + # The failed definition must not clobber the existing registration. + self.assertIs(get_eval_type_handler(unused_eval_type), _First) + finally: + _eval_type_handlers.pop(unused_eval_type, None) + + def test_abstract_handler_with_eval_type_rejected(self): + # A subclass that declares an eval_type but leaves run abstract must be + # rejected at class definition. + unused_eval_type = -1 + + def _define_abstract(): + class _Abstract(BatchEvalTypeHandler["pa.RecordBatch"]): + eval_type = unused_eval_type + # run left abstract on purpose + + self.assertRaises(AssertionError, _define_abstract) + self.assertNotIn(unused_eval_type, _eval_type_handlers) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() From 897d22d846c10240f8e2d25ad80adafe7c2291f2 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:47:52 +0000 Subject: [PATCH 09/20] test: hoist pyarrow import to file level in test_arrow_eval_type_handlers The whole module is Arrow-only and gated on have_pyarrow, so import pyarrow once under that guard instead of in every test method. Co-authored-by: Isaac --- .../tests/test_arrow_eval_type_handlers.py | 24 ++----------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/python/pyspark/tests/test_arrow_eval_type_handlers.py b/python/pyspark/tests/test_arrow_eval_type_handlers.py index 807b45bae6518..1efca8f059665 100644 --- a/python/pyspark/tests/test_arrow_eval_type_handlers.py +++ b/python/pyspark/tests/test_arrow_eval_type_handlers.py @@ -26,6 +26,8 @@ from pyspark.util import PythonEvalType if have_pyarrow: + import pyarrow as pa + # The handlers live in ``_arrow``, which imports pyarrow at module top. from pyspark.eval_handlers._arrow import ( ArrowCoGroupedMapUDFHandler, @@ -64,8 +66,6 @@ def test_arrow_eval_types_are_registered(self): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowScalarUDFHandlerTests(unittest.TestCase): def test_end_to_end_output(self): - import pyarrow as pa - # One UDF reading column 0 (a pa.Array) and returning column + 1. def add_one(col): return pa.array([v.as_py() + 1 for v in col], type=pa.int64()) @@ -81,8 +81,6 @@ def add_one(col): self.assertEqual(out[0].column(0).to_pylist(), [2, 3, 4]) def test_output_schema_enforced(self): - import pyarrow as pa - # The UDF returns int32, but the declared return type is LongType (int64). # run must enforce the declared schema onto the output batch. def add_one(col): @@ -101,8 +99,6 @@ def add_one(col): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowScalarIterUDFHandlerTests(unittest.TestCase): def test_end_to_end_output(self): - import pyarrow as pa - # The UDF receives an iterator of the single argument column and yields # an iterator of pa.Array; the handler assembles each into a RecordBatch. def add_one(col_iter): @@ -120,8 +116,6 @@ def add_one(col_iter): self.assertEqual([b.column(0).to_pylist() for b in out], [[2, 3], [4]]) def test_row_count_mismatch_is_rejected(self): - import pyarrow as pa - from pyspark.errors import PySparkRuntimeError # Emitting more rows than were consumed must fail (fail-fast row limit). @@ -139,8 +133,6 @@ def too_many(col_iter): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowMapUDFHandlerTests(unittest.TestCase): def test_end_to_end_output(self): - import pyarrow as pa - from pyspark.sql.conversion import ArrowBatchTransformer # mapInArrow exchanges a single struct column on the wire; the handler @@ -170,8 +162,6 @@ class ArrowGroupedMapUDFHandlerTests(unittest.TestCase): _ARG_OFFSETS = [3, 1, 0, 1] def _grouped_input(self): - import pyarrow as pa - from pyspark.sql.conversion import ArrowBatchTransformer inner = pa.RecordBatch.from_arrays( @@ -182,8 +172,6 @@ def _grouped_input(self): return iter([iter([wrapped])]) def test_values_only(self): - import pyarrow as pa - return_type = StructType([StructField("v", LongType())]) def grouped_udf(value_table): @@ -195,8 +183,6 @@ def grouped_udf(value_table): self.assertEqual(out[0].column(0).field("v").to_pylist(), [10, 20]) def test_key_and_values(self): - import pyarrow as pa - return_type = StructType([StructField("v", LongType())]) def grouped_udf(key, value_table): @@ -213,8 +199,6 @@ def grouped_udf(key, value_table): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowGroupedMapIterUDFHandlerTests(unittest.TestCase): def test_end_to_end_output(self): - import pyarrow as pa - from pyspark.sql.conversion import ArrowBatchTransformer return_type = StructType([StructField("v", LongType())]) @@ -240,8 +224,6 @@ def grouped_udf(value_batches): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowCoGroupedMapUDFHandlerTests(unittest.TestCase): def test_end_to_end_output(self): - import pyarrow as pa - return_type = StructType([StructField("out", LongType())]) def cogrouped_udf(left_values, right_values): @@ -270,8 +252,6 @@ def test_deserialized_co_group_is_a_pair_of_lists(self): # serializer eagerly materializes each side as a list, not an iterator. import io - import pyarrow as pa - from pyspark.serializers import write_int def arrow_bytes(batches): From ae5466c25d0bec3dfdf02824f61f744bc6dd092d Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 02:55:00 +0000 Subject: [PATCH 10/20] refactor: guard _arrow import with require_minimum_pyarrow_version, not the test helper have_pyarrow lives in pyspark.testing.utils; use the production require_minimum_pyarrow_version so the package init does not pull the test module into worker startup. Tests keep using have_pyarrow. Co-authored-by: Isaac --- python/pyspark/eval_handlers/__init__.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/pyspark/eval_handlers/__init__.py b/python/pyspark/eval_handlers/__init__.py index 466051f5579ad..ff5439a652cf2 100644 --- a/python/pyspark/eval_handlers/__init__.py +++ b/python/pyspark/eval_handlers/__init__.py @@ -28,7 +28,12 @@ when pyarrow is available; the Arrow eval types it serves cannot run without it. """ -from pyspark.testing.utils import have_pyarrow +try: + from pyspark.sql.pandas.utils import require_minimum_pyarrow_version -if have_pyarrow: + require_minimum_pyarrow_version() +except Exception: + # pyarrow is missing or too old; the Arrow eval types _arrow serves cannot run anyway. + pass +else: from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import From 760a083bffcb834fec5aba3203072f0d300db3b6 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 04:20:08 +0000 Subject: [PATCH 11/20] refactor: rename eval_handlers/_util.py to _utils.py Co-authored-by: Isaac --- python/pyspark/eval_handlers/_arrow.py | 2 +- python/pyspark/eval_handlers/{_util.py => _utils.py} | 0 python/pyspark/worker.py | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) rename python/pyspark/eval_handlers/{_util.py => _utils.py} (100%) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 14c2d2ce967ef..74bf31e4dfee5 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -35,7 +35,7 @@ CoGroupedEvalTypeHandler, GroupedEvalTypeHandler, ) -from pyspark.eval_handlers._util import extract_key_value_indexes +from pyspark.eval_handlers._utils import extract_key_value_indexes from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, diff --git a/python/pyspark/eval_handlers/_util.py b/python/pyspark/eval_handlers/_utils.py similarity index 100% rename from python/pyspark/eval_handlers/_util.py rename to python/pyspark/eval_handlers/_utils.py diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 486c42e691754..e00e1994c84d8 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -53,7 +53,7 @@ ) from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.eval_handlers._base import get_eval_type_handler -from pyspark.eval_handlers._util import extract_key_value_indexes +from pyspark.eval_handlers._utils import extract_key_value_indexes from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, From 81080c894f30853a58f673ba08da993a6b36a801 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 04:21:31 +0000 Subject: [PATCH 12/20] refactor: name the helper module eval_handlers/utils.py (no leading underscore) Co-authored-by: Isaac --- python/pyspark/eval_handlers/_arrow.py | 2 +- python/pyspark/eval_handlers/{_utils.py => utils.py} | 0 python/pyspark/worker.py | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) rename python/pyspark/eval_handlers/{_utils.py => utils.py} (100%) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 74bf31e4dfee5..54e50d391bd98 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -35,7 +35,7 @@ CoGroupedEvalTypeHandler, GroupedEvalTypeHandler, ) -from pyspark.eval_handlers._utils import extract_key_value_indexes +from pyspark.eval_handlers.utils import extract_key_value_indexes from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, diff --git a/python/pyspark/eval_handlers/_utils.py b/python/pyspark/eval_handlers/utils.py similarity index 100% rename from python/pyspark/eval_handlers/_utils.py rename to python/pyspark/eval_handlers/utils.py diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index e00e1994c84d8..b3a99ee9ff8a9 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -53,7 +53,7 @@ ) from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.eval_handlers._base import get_eval_type_handler -from pyspark.eval_handlers._utils import extract_key_value_indexes +from pyspark.eval_handlers.utils import extract_key_value_indexes from pyspark.eval_handlers.verification import ( verify_iter_result_row_count, verify_iterator_exhausted, From 6b53587e2e03a73bedb06bc57c0d70f324a1e938 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 05:27:01 +0000 Subject: [PATCH 13/20] test: unify the Arrow handler tests Build handler input through shared helpers (_batch, _struct_batch, _one_group, _one_cogroup, _grouped_arg_offsets) and construct handlers via _scalar_handler / _grouped_handler, so every test follows the same shape and drops the ad hoc arg_offsets literals. Cover both the values-only and with-key branches uniformly for the grouped, grouped-iter, and co-grouped handlers. Co-authored-by: Isaac --- .../tests/test_arrow_eval_type_handlers.py | 257 +++++++++--------- 1 file changed, 132 insertions(+), 125 deletions(-) diff --git a/python/pyspark/tests/test_arrow_eval_type_handlers.py b/python/pyspark/tests/test_arrow_eval_type_handlers.py index 1efca8f059665..79e876ed0df6e 100644 --- a/python/pyspark/tests/test_arrow_eval_type_handlers.py +++ b/python/pyspark/tests/test_arrow_eval_type_handlers.py @@ -15,10 +15,16 @@ # limitations under the License. # -"""Tests for the Arrow eval type handlers (``_arrow``).""" +"""Tests for the Arrow eval type handlers (``_arrow``). + +Each handler has one test class. The helpers below build handler input in the +same wire format the serializers produce, so every test constructs its input the +same way: ``run(0, )`` and assert on the output batches. +""" import unittest +from pyspark.errors import PySparkRuntimeError from pyspark.eval_handlers._base import get_eval_type_handler from pyspark.sql.pandas.serializers import ArrowStreamCoGroupSerializer, ArrowStreamSerializer from pyspark.sql.types import LongType, StructField, StructType @@ -28,7 +34,6 @@ if have_pyarrow: import pyarrow as pa - # The handlers live in ``_arrow``, which imports pyarrow at module top. from pyspark.eval_handlers._arrow import ( ArrowCoGroupedMapUDFHandler, ArrowGroupedMapIterUDFHandler, @@ -37,6 +42,7 @@ ArrowScalarIterUDFHandler, ArrowScalarUDFHandler, ) + from pyspark.sql.conversion import ArrowBatchTransformer class _RunnerConf: @@ -48,6 +54,69 @@ class _RunnerConf: map_in_batch_legacy_accept_any_iterable = False +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 _struct_batch(**columns): + """``_batch`` wrapped into a single struct column (the grouped/map wire format).""" + return ArrowBatchTransformer.wrap_struct(_batch(**columns)) + + +def _one_group(*batches): + """One group of the given batches, shaped as the group serializer yields it.""" + return iter([iter(batches)]) + + +def _one_cogroup(left, right): + """One co-group of a left and a right batch, as the co-group serializer yields it.""" + return iter([([left], [right])]) + + +def _grouped_arg_offsets(*dataframes): + """Encode ``arg_offsets`` from ``(key_cols, value_cols)`` per DataFrame. + + Mirrors BasePandasGroupExec.resolveArgOffsets: each DataFrame is laid out as + ``[length, num_keys, *key_cols, *value_cols]``. + """ + offsets: list = [] + for key_cols, value_cols in dataframes: + group = [len(key_cols), *key_cols, *value_cols] + offsets += [len(group), *group] + return offsets + + +def _scalar_handler(handler_cls, udf): + """Build a scalar handler whose one UDF reads column 0 and returns LongType. + + The scalar UDF tuple is ``(func, args_offsets, kwargs_offsets, return_type)``. + """ + return handler_cls(udfs=[(udf, [0], {}, LongType())], runner_conf=_RunnerConf(), eval_conf=None) + + +def _grouped_handler(handler_cls, udf, arg_offsets, num_udf_args): + """Build a grouped/cogrouped-map handler with the shared return type. + + The grouped-map UDF tuple is ``(func, arg_offsets, return_type, num_udf_args)``. + """ + return handler_cls( + udfs=[(udf, arg_offsets, _RETURN_TYPE, num_udf_args)], + runner_conf=_RunnerConf(), + eval_conf=None, + ) + + +# arg_offsets for one DataFrame with key column 0 and value column 1. +_GROUP_OFFSETS = _grouped_arg_offsets(([0], [1])) +# arg_offsets for two DataFrames (co-group), each with key column 0 and value column 1. +_COGROUP_OFFSETS = _grouped_arg_offsets(([0], [1]), ([0], [1])) +_RETURN_TYPE = StructType([StructField("v", LongType())]) + + @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowEvalTypeHandlerRegistrationTests(unittest.TestCase): def test_arrow_eval_types_are_registered(self): @@ -65,184 +134,125 @@ def test_arrow_eval_types_are_registered(self): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowScalarUDFHandlerTests(unittest.TestCase): - def test_end_to_end_output(self): - # One UDF reading column 0 (a pa.Array) and returning column + 1. + def test_invokes_udf_per_batch(self): def add_one(col): return pa.array([v.as_py() + 1 for v in col], type=pa.int64()) - udfs = [(add_one, [0], {}, LongType())] - handler = ArrowScalarUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - - batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3], type=pa.int64())], ["_0"]) - out = list(handler.run(0, iter([batch]))) + handler = _scalar_handler(ArrowScalarUDFHandler, add_one) + 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]]) - self.assertEqual(len(out), 1) - self.assertEqual(out[0].num_columns, 1) - self.assertEqual(out[0].column(0).to_pylist(), [2, 3, 4]) - - def test_output_schema_enforced(self): + def test_coerces_output_to_return_type(self): # The UDF returns int32, but the declared return type is LongType (int64). - # run must enforce the declared schema onto the output batch. def add_one(col): return pa.array([v.as_py() + 1 for v in col], type=pa.int32()) - udfs = [(add_one, [0], {}, LongType())] - handler = ArrowScalarUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - - batch = pa.RecordBatch.from_arrays([pa.array([10, 20], type=pa.int64())], ["_0"]) - out = list(handler.run(0, iter([batch]))) - # The int32 the UDF produced is coerced to the declared LongType (int64). + handler = _scalar_handler(ArrowScalarUDFHandler, 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]) @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowScalarIterUDFHandlerTests(unittest.TestCase): - def test_end_to_end_output(self): - # The UDF receives an iterator of the single argument column and yields - # an iterator of pa.Array; the handler assembles each into a RecordBatch. + def test_invokes_udf_over_batch_stream(self): def add_one(col_iter): for col in col_iter: yield pa.array([v.as_py() + 1 for v in col], type=pa.int64()) - udfs = [(add_one, [0], {}, LongType())] - handler = ArrowScalarIterUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - - batches = [ - pa.RecordBatch.from_arrays([pa.array([1, 2], type=pa.int64())], ["_0"]), - pa.RecordBatch.from_arrays([pa.array([3], type=pa.int64())], ["_0"]), - ] - out = list(handler.run(0, iter(batches))) + handler = _scalar_handler(ArrowScalarIterUDFHandler, 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_row_count_mismatch_is_rejected(self): - from pyspark.errors import PySparkRuntimeError - + def test_rejects_row_count_mismatch(self): # Emitting more rows than were consumed must fail (fail-fast row limit). def too_many(col_iter): for col in col_iter: yield pa.array(list(range(len(col) + 1)), type=pa.int64()) - udfs = [(too_many, [0], {}, LongType())] - handler = ArrowScalarIterUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - batch = pa.RecordBatch.from_arrays([pa.array([1, 2], type=pa.int64())], ["_0"]) + handler = _scalar_handler(ArrowScalarIterUDFHandler, too_many) with self.assertRaises(PySparkRuntimeError): - list(handler.run(0, iter([batch]))) + list(handler.run(0, iter([_batch(a=[1, 2])]))) @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowMapUDFHandlerTests(unittest.TestCase): - def test_end_to_end_output(self): - from pyspark.sql.conversion import ArrowBatchTransformer - - # mapInArrow exchanges a single struct column on the wire; the handler - # flattens it for the UDF and re-wraps the UDF's output. - def double_a(batch_iter): + def test_maps_batch_stream(self): + # mapInArrow flattens the wire struct for the UDF and re-wraps its output. + def double_v(batch_iter): for batch in batch_iter: - doubled = pa.array([v.as_py() * 2 for v in batch.column(0)], type=pa.int64()) - yield pa.RecordBatch.from_arrays([doubled], ["a"]) - - inner = pa.RecordBatch.from_arrays([pa.array([1, 2, 3], type=pa.int64())], ["a"]) - wrapped = ArrowBatchTransformer.wrap_struct(inner) - - udfs = [(double_a, None, None, None)] - handler = ArrowMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - out = list(handler.run(0, iter([wrapped]))) + yield _batch(v=[c.as_py() * 2 for c in batch.column("v")]) - self.assertEqual(len(out), 1) - # Output is a single struct column; its "a" field carries the doubled values. - self.assertEqual(out[0].num_columns, 1) - self.assertEqual(out[0].column(0).field("a").to_pylist(), [2, 4, 6]) + handler = ArrowMapUDFHandler( + udfs=[(double_v, None, None, None)], runner_conf=_RunnerConf(), eval_conf=None + ) + out = list(handler.run(0, iter([_struct_batch(v=[1, 2, 3])]))) + self.assertEqual(out[0].column(0).field("v").to_pylist(), [2, 4, 6]) @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowGroupedMapUDFHandlerTests(unittest.TestCase): - # arg_offsets encoding for one DataFrame with key column 0 and value column 1: - # [group_len=3, num_keys=1, key_offset=0, value_offset=1] - _ARG_OFFSETS = [3, 1, 0, 1] - - def _grouped_input(self): - from pyspark.sql.conversion import ArrowBatchTransformer - - inner = pa.RecordBatch.from_arrays( - [pa.array([7, 7], type=pa.int64()), pa.array([1, 2], type=pa.int64())], ["k", "v"] - ) - wrapped = ArrowBatchTransformer.wrap_struct(inner) - # One group, whose batches arrive as an iterator (matching the group serializer). - return iter([iter([wrapped])]) - - def test_values_only(self): - return_type = StructType([StructField("v", LongType())]) - + def test_applies_udf_per_group(self): def grouped_udf(value_table): return pa.table({"v": pa.array([c.as_py() * 10 for c in value_table.column("v")])}) - udfs = [(grouped_udf, self._ARG_OFFSETS, return_type, 1)] - handler = ArrowGroupedMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - out = list(handler.run(0, self._grouped_input())) + handler = _grouped_handler(ArrowGroupedMapUDFHandler, grouped_udf, _GROUP_OFFSETS, 1) + out = list(handler.run(0, _one_group(_struct_batch(k=[7, 7], v=[1, 2])))) self.assertEqual(out[0].column(0).field("v").to_pylist(), [10, 20]) - def test_key_and_values(self): - return_type = StructType([StructField("v", LongType())]) - + def test_passes_key_when_udf_takes_key(self): def grouped_udf(key, value_table): - # key is the grouping-key tuple; add it to every value. k = key[0].as_py() return pa.table({"v": pa.array([c.as_py() + k for c in value_table.column("v")])}) - udfs = [(grouped_udf, self._ARG_OFFSETS, return_type, 2)] - handler = ArrowGroupedMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - out = list(handler.run(0, self._grouped_input())) + handler = _grouped_handler(ArrowGroupedMapUDFHandler, grouped_udf, _GROUP_OFFSETS, 2) + out = list(handler.run(0, _one_group(_struct_batch(k=[7, 7], v=[1, 2])))) self.assertEqual(out[0].column(0).field("v").to_pylist(), [8, 9]) @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowGroupedMapIterUDFHandlerTests(unittest.TestCase): - def test_end_to_end_output(self): - from pyspark.sql.conversion import ArrowBatchTransformer - - return_type = StructType([StructField("v", LongType())]) - + def test_applies_udf_per_group(self): def grouped_udf(value_batches): for batch in value_batches: - yield pa.RecordBatch.from_arrays( - [pa.array([c.as_py() + 1 for c in batch.column("v")], type=pa.int64())], ["v"] - ) + yield _batch(v=[c.as_py() + 1 for c in batch.column("v")]) - inner = pa.RecordBatch.from_arrays( - [pa.array([7], type=pa.int64()), pa.array([41], type=pa.int64())], ["k", "v"] - ) - wrapped = ArrowBatchTransformer.wrap_struct(inner) - udfs = [(grouped_udf, [3, 1, 0, 1], return_type, 1)] - handler = ArrowGroupedMapIterUDFHandler( - udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None - ) - out = list(handler.run(0, iter([iter([wrapped])]))) + handler = _grouped_handler(ArrowGroupedMapIterUDFHandler, grouped_udf, _GROUP_OFFSETS, 1) + out = list(handler.run(0, _one_group(_struct_batch(k=[7], v=[41])))) self.assertEqual(out[0].column(0).field("v").to_pylist(), [42]) + def test_passes_key_when_udf_takes_key(self): + def grouped_udf(key, value_batches): + k = key[0].as_py() + for batch in value_batches: + yield _batch(v=[c.as_py() + k for c in batch.column("v")]) + + handler = _grouped_handler(ArrowGroupedMapIterUDFHandler, grouped_udf, _GROUP_OFFSETS, 2) + out = list(handler.run(0, _one_group(_struct_batch(k=[10], v=[5])))) + self.assertEqual(out[0].column(0).field("v").to_pylist(), [15]) + @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowCoGroupedMapUDFHandlerTests(unittest.TestCase): - def test_end_to_end_output(self): - return_type = StructType([StructField("out", LongType())]) - + # Co-group batches arrive un-wrapped (columns k, v), unlike the grouped-map wire format. + def test_applies_udf_per_cogroup(self): def cogrouped_udf(left_values, right_values): - total = left_values.column("lv")[0].as_py() + right_values.column("rv")[0].as_py() - return pa.table({"out": pa.array([total], type=pa.int64())}) + total = left_values.column("v")[0].as_py() + right_values.column("v")[0].as_py() + return pa.table({"v": pa.array([total], type=pa.int64())}) - # A co-group deserializes to a pair of lists of (non-struct-wrapped) batches. - left = pa.RecordBatch.from_arrays( - [pa.array([5], type=pa.int64()), pa.array([10], type=pa.int64())], ["k", "lv"] - ) - right = pa.RecordBatch.from_arrays( - [pa.array([5], type=pa.int64()), pa.array([20], type=pa.int64())], ["k", "rv"] - ) - # Two DataFrames, each key column 0 and value column 1. - arg_offsets = [3, 1, 0, 1, 3, 1, 0, 1] - udfs = [(cogrouped_udf, arg_offsets, return_type, 2)] - handler = ArrowCoGroupedMapUDFHandler(udfs=udfs, runner_conf=_RunnerConf(), eval_conf=None) - out = list(handler.run(0, iter([([left], [right])]))) - self.assertEqual(out[0].column(0).field("out").to_pylist(), [30]) + handler = _grouped_handler(ArrowCoGroupedMapUDFHandler, cogrouped_udf, _COGROUP_OFFSETS, 2) + out = list(handler.run(0, _one_cogroup(_batch(k=[5], v=[10]), _batch(k=[5], v=[20])))) + self.assertEqual(out[0].column(0).field("v").to_pylist(), [30]) + + def test_passes_key_when_udf_takes_key(self): + def cogrouped_udf(key, left_values, right_values): + k = key[0].as_py() + total = left_values.column("v")[0].as_py() + right_values.column("v")[0].as_py() + return pa.table({"v": pa.array([total + k], type=pa.int64())}) + + handler = _grouped_handler(ArrowCoGroupedMapUDFHandler, cogrouped_udf, _COGROUP_OFFSETS, 3) + out = list(handler.run(0, _one_cogroup(_batch(k=[5], v=[10]), _batch(k=[5], v=[20])))) + self.assertEqual(out[0].column(0).field("v").to_pylist(), [35]) @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) @@ -259,13 +269,10 @@ def arrow_bytes(batches): ArrowStreamSerializer().dump_stream(iter(batches), buf) return buf.getvalue() - left = pa.RecordBatch.from_arrays([pa.array([1, 2])], ["_0"]) - right = pa.RecordBatch.from_arrays([pa.array([9])], ["_0"]) - stream = io.BytesIO() write_int(2, stream) # two DataFrames in the co-group - stream.write(arrow_bytes([left])) - stream.write(arrow_bytes([right])) + stream.write(arrow_bytes([_batch(v=[1, 2])])) + stream.write(arrow_bytes([_batch(v=[9])])) write_int(0, stream) # end of stream stream.seek(0) From 57612de4ab29730567821d1f37873c40d6c6a309 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:54:53 +0000 Subject: [PATCH 14/20] test: move eval type handler tests under pyspark.sql.tests They exercise SQL Arrow-UDF machinery, so they belong in the pyspark-sql module rather than the top-level pyspark-core tests: the framework tests move to pyspark.sql.tests and the Arrow handler tests to pyspark.sql.tests.arrow. Co-authored-by: Isaac --- dev/sparktestsupport/modules.py | 4 ++-- .../tests/arrow}/test_arrow_eval_type_handlers.py | 0 .../pyspark/{ => sql}/tests/test_base_eval_type_handlers.py | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename python/pyspark/{tests => sql/tests/arrow}/test_arrow_eval_type_handlers.py (100%) rename python/pyspark/{ => sql}/tests/test_base_eval_type_handlers.py (100%) diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 74a9aee355a18..1ac634cd9cfdf 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -612,8 +612,6 @@ def __hash__(self): "pyspark.tests.test_taskcontext", "pyspark.tests.test_util", "pyspark.tests.test_worker", - "pyspark.tests.test_base_eval_type_handlers", - "pyspark.tests.test_arrow_eval_type_handlers", "pyspark.tests.test_stage_sched", "pyspark.tests.test_zero_copy_byte_stream", # unittests for upstream projects @@ -692,10 +690,12 @@ def __hash__(self): "pyspark.sql.tests.test_group", "pyspark.sql.tests.test_sql", "pyspark.sql.tests.test_job_cancellation", + "pyspark.sql.tests.test_base_eval_type_handlers", "pyspark.sql.tests.arrow.test_arrow", "pyspark.sql.tests.arrow.test_arrow_map", "pyspark.sql.tests.arrow.test_arrow_cogrouped_map", "pyspark.sql.tests.arrow.test_arrow_cogrouped_map_misc", + "pyspark.sql.tests.arrow.test_arrow_eval_type_handlers", "pyspark.sql.tests.arrow.test_arrow_grouped_map", "pyspark.sql.tests.arrow.test_arrow_python_aggregator", "pyspark.sql.tests.arrow.test_arrow_python_udf", diff --git a/python/pyspark/tests/test_arrow_eval_type_handlers.py b/python/pyspark/sql/tests/arrow/test_arrow_eval_type_handlers.py similarity index 100% rename from python/pyspark/tests/test_arrow_eval_type_handlers.py rename to python/pyspark/sql/tests/arrow/test_arrow_eval_type_handlers.py diff --git a/python/pyspark/tests/test_base_eval_type_handlers.py b/python/pyspark/sql/tests/test_base_eval_type_handlers.py similarity index 100% rename from python/pyspark/tests/test_base_eval_type_handlers.py rename to python/pyspark/sql/tests/test_base_eval_type_handlers.py From b8ac03b97fce5bc966c8ba0a580fb966074ae03e Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:00:38 +0000 Subject: [PATCH 15/20] test: co-locate eval type handler tests in pyspark.eval_handlers.tests Put the tests in a tests subpackage next to the code they exercise, mirroring the per-component tests convention (pyspark.sql.tests, pyspark.ml.tests, ...), instead of the top-level pyspark tests. Co-authored-by: Isaac --- dev/sparktestsupport/modules.py | 4 ++-- python/pyspark/eval_handlers/tests/__init__.py | 16 ++++++++++++++++ .../tests}/test_arrow_eval_type_handlers.py | 0 .../tests/test_base_eval_type_handlers.py | 0 4 files changed, 18 insertions(+), 2 deletions(-) create mode 100644 python/pyspark/eval_handlers/tests/__init__.py rename python/pyspark/{sql/tests/arrow => eval_handlers/tests}/test_arrow_eval_type_handlers.py (100%) rename python/pyspark/{sql => eval_handlers}/tests/test_base_eval_type_handlers.py (100%) diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 1ac634cd9cfdf..f1c78b63b92f7 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -690,12 +690,12 @@ def __hash__(self): "pyspark.sql.tests.test_group", "pyspark.sql.tests.test_sql", "pyspark.sql.tests.test_job_cancellation", - "pyspark.sql.tests.test_base_eval_type_handlers", + "pyspark.eval_handlers.tests.test_base_eval_type_handlers", + "pyspark.eval_handlers.tests.test_arrow_eval_type_handlers", "pyspark.sql.tests.arrow.test_arrow", "pyspark.sql.tests.arrow.test_arrow_map", "pyspark.sql.tests.arrow.test_arrow_cogrouped_map", "pyspark.sql.tests.arrow.test_arrow_cogrouped_map_misc", - "pyspark.sql.tests.arrow.test_arrow_eval_type_handlers", "pyspark.sql.tests.arrow.test_arrow_grouped_map", "pyspark.sql.tests.arrow.test_arrow_python_aggregator", "pyspark.sql.tests.arrow.test_arrow_python_udf", diff --git a/python/pyspark/eval_handlers/tests/__init__.py b/python/pyspark/eval_handlers/tests/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/eval_handlers/tests/__init__.py @@ -0,0 +1,16 @@ +# +# 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. +# diff --git a/python/pyspark/sql/tests/arrow/test_arrow_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py similarity index 100% rename from python/pyspark/sql/tests/arrow/test_arrow_eval_type_handlers.py rename to python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py diff --git a/python/pyspark/sql/tests/test_base_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py similarity index 100% rename from python/pyspark/sql/tests/test_base_eval_type_handlers.py rename to python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py From 21ccf363ec9e3325df2d06803bad586ab8765c63 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:02:39 +0000 Subject: [PATCH 16/20] test: group eval_handlers test goals together in modules.py Keep the pyspark.sql.tests.* entries contiguous by listing the pyspark.eval_handlers.tests goals as their own block at the end rather than interleaved. Co-authored-by: Isaac --- dev/sparktestsupport/modules.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index f1c78b63b92f7..990b061b784bc 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -690,8 +690,6 @@ def __hash__(self): "pyspark.sql.tests.test_group", "pyspark.sql.tests.test_sql", "pyspark.sql.tests.test_job_cancellation", - "pyspark.eval_handlers.tests.test_base_eval_type_handlers", - "pyspark.eval_handlers.tests.test_arrow_eval_type_handlers", "pyspark.sql.tests.arrow.test_arrow", "pyspark.sql.tests.arrow.test_arrow_map", "pyspark.sql.tests.arrow.test_arrow_cogrouped_map", @@ -751,6 +749,8 @@ def __hash__(self): "pyspark.sql.tests.coercion.test_python_udf_return_type", "pyspark.sql.tests.df_golden.test_df_golden", "pyspark.sql.tests.df_golden.test_df_golden_framework", + "pyspark.eval_handlers.tests.test_base_eval_type_handlers", + "pyspark.eval_handlers.tests.test_arrow_eval_type_handlers", ], ) From 21e65d22dffe49cd43392426cf21d388209340b7 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:05:00 +0000 Subject: [PATCH 17/20] test: order the eval_handlers test goals alphabetically Co-authored-by: Isaac --- dev/sparktestsupport/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 990b061b784bc..c7ea346946ffc 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -749,8 +749,8 @@ def __hash__(self): "pyspark.sql.tests.coercion.test_python_udf_return_type", "pyspark.sql.tests.df_golden.test_df_golden", "pyspark.sql.tests.df_golden.test_df_golden_framework", - "pyspark.eval_handlers.tests.test_base_eval_type_handlers", "pyspark.eval_handlers.tests.test_arrow_eval_type_handlers", + "pyspark.eval_handlers.tests.test_base_eval_type_handlers", ], ) From 68245506d566ebb2cc8eaa5c56fab55222dacf0a Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:09:08 +0000 Subject: [PATCH 18/20] test: list eval_handlers test goals before the pyspark.sql ones Co-authored-by: Isaac --- dev/sparktestsupport/modules.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index c7ea346946ffc..215839be14750 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -668,6 +668,8 @@ def __hash__(self): "pyspark.sql.observation", "pyspark.sql.tvf", # unittests + "pyspark.eval_handlers.tests.test_arrow_eval_type_handlers", + "pyspark.eval_handlers.tests.test_base_eval_type_handlers", "pyspark.sql.tests.test_artifact", "pyspark.sql.tests.test_catalog", "pyspark.sql.tests.test_column", @@ -749,8 +751,6 @@ def __hash__(self): "pyspark.sql.tests.coercion.test_python_udf_return_type", "pyspark.sql.tests.df_golden.test_df_golden", "pyspark.sql.tests.df_golden.test_df_golden_framework", - "pyspark.eval_handlers.tests.test_arrow_eval_type_handlers", - "pyspark.eval_handlers.tests.test_base_eval_type_handlers", ], ) From bce1fe745d30de8b60bdc78427446b6589883248 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 19:08:31 +0000 Subject: [PATCH 19/20] test: exempt pyspark.eval_handlers.tests from mypy like other test packages Co-authored-by: Isaac --- python/mypy.ini | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/mypy.ini b/python/mypy.ini index 5baa77c370c51..ca4ad20db04ca 100644 --- a/python/mypy.ini +++ b/python/mypy.ini @@ -70,6 +70,9 @@ disable_error_code = attr-defined,arg-type,call-arg,union-attr ; Ignore errors in tests +[mypy-pyspark.eval_handlers.tests.*] +ignore_errors = True + [mypy-pyspark.ml.tests.*] ignore_errors = True From d200d853ad6ed9407ac368cb4555db9f1e5c42fa Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Mon, 21 Sep 2026 16:20:17 +0000 Subject: [PATCH 20/20] refactor: register Arrow handlers without pyarrow, check version at run Decouple handler registration from pyarrow availability: _arrow imports pyarrow lazily (inside run and TYPE_CHECKING) so its handlers always register, and each handler calls require_minimum_pyarrow_version in __init__. A missing or too-old pyarrow now raises a clear error when the Arrow eval type runs, instead of leaving it unregistered and failing later with "Unknown eval type". The package __init__ imports _arrow unconditionally again. Co-authored-by: Isaac --- python/pyspark/eval_handlers/__init__.py | 14 ++------ python/pyspark/eval_handlers/_arrow.py | 45 +++++++++++++++++------- 2 files changed, 35 insertions(+), 24 deletions(-) diff --git a/python/pyspark/eval_handlers/__init__.py b/python/pyspark/eval_handlers/__init__.py index ff5439a652cf2..2f90d4b68b6a3 100644 --- a/python/pyspark/eval_handlers/__init__.py +++ b/python/pyspark/eval_handlers/__init__.py @@ -24,16 +24,8 @@ ``read_udfs`` looks up via ``get_eval_type_handler``. Importing this package imports the concrete handler submodules so they register. -``_arrow`` requires pyarrow and imports it at module top, so it is only imported -when pyarrow is available; the Arrow eval types it serves cannot run without it. +``_arrow`` imports pyarrow lazily, so it registers its handlers without pyarrow +installed and defers the pyarrow check to when a handler runs. """ -try: - from pyspark.sql.pandas.utils import require_minimum_pyarrow_version - - require_minimum_pyarrow_version() -except Exception: - # pyarrow is missing or too old; the Arrow eval types _arrow serves cannot run anyway. - pass -else: - from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import +from pyspark.eval_handlers import _arrow # noqa: F401 # registers handlers on import diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 54e50d391bd98..62e36b9ff792c 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -18,8 +18,11 @@ """Handlers for the Arrow-native UDF eval types (the UDF exchanges ``pa.Array`` / ``pa.RecordBatch`` values directly, without a pandas conversion). -This module imports ``pyarrow`` at the top level, so the package ``__init__`` only -imports it when pyarrow is available; callers must do the same. +pyarrow is imported lazily (inside ``run`` and the type-checking block) so the module +stays importable and its handlers register without pyarrow installed. Each handler +calls ``require_minimum_pyarrow_version`` in ``__init__``, so a missing or too-old +pyarrow surfaces a clear error when the handler runs rather than leaving the eval type +unregistered. """ from __future__ import annotations @@ -28,8 +31,6 @@ from collections.abc import Iterator from typing import TYPE_CHECKING, Any -import pyarrow as pa - from pyspark.eval_handlers._base import ( BatchEvalTypeHandler, CoGroupedEvalTypeHandler, @@ -45,18 +46,18 @@ ) from pyspark.sql.conversion import ArrowBatchTransformer from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type +from pyspark.sql.pandas.utils import require_minimum_pyarrow_version from pyspark.sql.types import StructField, StructType from pyspark.util import PythonEvalType if TYPE_CHECKING: - # Annotation-only, so they are not imported at runtime. ``from __future__ import - # annotations`` keeps every annotation below an unevaluated name rather than a - # string literal. + import pyarrow as pa + from pyspark.eval_handlers._typing import CoGroupedBatch, GroupedBatch from pyspark.worker_util import EvalConf, RunnerConf -class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler[pa.RecordBatch]): +class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]): """SQL_COGROUPED_MAP_ARROW_UDF (applyInArrow on a cogroup): the single UDF receives the two sides' value tables and returns one pa.Table, coerced to the declared schema.""" @@ -66,6 +67,7 @@ class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler[pa.RecordBatch]): def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: + require_minimum_pyarrow_version() super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One COGROUPED_MAP_ARROW UDF expected here." self._cogrouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] @@ -78,6 +80,8 @@ def __init__( def run(self, split_index: int, data: Iterator[CoGroupedBatch]) -> Iterator[pa.RecordBatch]: """Apply cogroupBy Arrow UDF.""" + import pyarrow as pa + select_columns = ArrowBatchTransformer.select_columns def table_from_batches(batches: list[pa.RecordBatch], cols: list[int]) -> pa.Table: @@ -109,7 +113,7 @@ def table_from_batches(batches: list[pa.RecordBatch], cols: list[int]) -> pa.Tab yield ArrowBatchTransformer.wrap_struct(batch) -class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): +class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): """SQL_GROUPED_MAP_ARROW_ITER_UDF: the single UDF receives each group as an iterator of RecordBatches and returns an iterator of RecordBatches, coerced to the declared schema.""" @@ -119,6 +123,7 @@ class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: + require_minimum_pyarrow_version() super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One GROUPED_MAP_ARROW_ITER UDF expected here." self._grouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] @@ -134,6 +139,8 @@ def __init__( def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: """Apply groupBy Arrow UDF (iterator variant).""" + import pyarrow as pa + key_offsets = self._key_offsets value_offsets = self._value_offsets for group in data: @@ -176,7 +183,7 @@ def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.Rec pass -class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): +class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): """SQL_GROUPED_MAP_ARROW_UDF (applyInArrow): the single UDF receives each group as one pa.Table and returns one pa.Table, coerced to the declared schema.""" @@ -186,6 +193,7 @@ class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler[pa.RecordBatch]): def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: + require_minimum_pyarrow_version() super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One GROUPED_MAP_ARROW UDF expected here." self._grouped_udf, arg_offsets, return_type, self._num_udf_args = udfs[0] @@ -199,6 +207,8 @@ def __init__( def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]: """Apply groupBy Arrow UDF (non-iterator variant).""" + import pyarrow as pa + key_offsets = self._key_offsets value_offsets = self._value_offsets for group in data: @@ -240,7 +250,7 @@ def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.Rec yield ArrowBatchTransformer.wrap_struct(batch) -class ArrowMapUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): +class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): """SQL_MAP_ARROW_ITER_UDF (mapInArrow): the single UDF receives the input RecordBatch stream and yields a RecordBatch stream, exchanged as flattened columns on the wire and wrapped back into a single struct column.""" @@ -250,11 +260,14 @@ class ArrowMapUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: + require_minimum_pyarrow_version() super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One MAP_ARROW_ITER UDF expected here." self._udf_func = udfs[0][0] def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + import pyarrow as pa + # Pre-processing input_batches = map(ArrowBatchTransformer.flatten_struct, data) @@ -280,7 +293,7 @@ def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.R yield from map(ArrowBatchTransformer.wrap_struct, verified_iter) -class ArrowScalarIterUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): +class ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): """SQL_SCALAR_ARROW_ITER_UDF: the UDF receives an iterator of the argument columns and yields an iterator of pa.Array; enforce the declared type on each result and verify the total row count matches the input.""" @@ -290,6 +303,7 @@ class ArrowScalarIterUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: + require_minimum_pyarrow_version() super().__init__(udfs, runner_conf, eval_conf) assert len(udfs) == 1, "One SCALAR_ARROW_ITER UDF expected here." self._udf_func, self._args_offsets, _, return_type = udfs[0] @@ -298,6 +312,8 @@ def __init__( ) def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + import pyarrow as pa + args_offsets = self._args_offsets num_input_rows = 0 @@ -334,7 +350,7 @@ def process_results() -> Iterator[pa.RecordBatch]: verify_iterator_exhausted(args_iter) -class ArrowScalarUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): +class ArrowScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): """SQL_SCALAR_ARROW_UDF: invoke each UDF once per input RecordBatch, coerce the result to the declared schema, and check the row count.""" @@ -343,6 +359,7 @@ class ArrowScalarUDFHandler(BatchEvalTypeHandler[pa.RecordBatch]): def __init__( self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf ) -> None: + require_minimum_pyarrow_version() super().__init__(udfs, runner_conf, eval_conf) self._col_names = ["_%d" % i for i in range(len(udfs))] self._combined_arrow_schema = to_arrow_schema( @@ -352,6 +369,8 @@ def __init__( ) def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: + import pyarrow as pa + for batch in data: output_batch = pa.RecordBatch.from_arrays( [