diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 61f3e4c155705..215839be14750 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -612,7 +612,6 @@ def __hash__(self): "pyspark.tests.test_taskcontext", "pyspark.tests.test_util", "pyspark.tests.test_worker", - "pyspark.tests.test_eval_type_handlers", "pyspark.tests.test_stage_sched", "pyspark.tests.test_zero_copy_byte_stream", # unittests for upstream projects @@ -669,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", 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 diff --git a/python/pyspark/eval_handlers/__init__.py b/python/pyspark/eval_handlers/__init__.py index 4c51849909861..2f90d4b68b6a3 100644 --- a/python/pyspark/eval_handlers/__init__.py +++ b/python/pyspark/eval_handlers/__init__.py @@ -22,7 +22,10 @@ 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`` imports pyarrow lazily, so it registers its handlers without pyarrow +installed and defers the pyarrow check to when a handler runs. """ 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 9f5df996ad22d..62e36b9ff792c 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -16,24 +16,340 @@ # """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). +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 + +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.utils import extract_key_value_indexes +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.pandas.utils import require_minimum_pyarrow_version from pyspark.sql.types import StructField, StructType from pyspark.util import PythonEvalType if TYPE_CHECKING: import pyarrow as pa + from pyspark.eval_handlers._typing import CoGroupedBatch, GroupedBatch from pyspark.worker_util import EvalConf, RunnerConf +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: + 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] + 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 = 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.""" + import pyarrow as pa + + 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]) + + 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) + + +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: + 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] + 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 = 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 (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 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: + 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] + 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 = 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).""" + 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 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: + 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) + + # 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: + 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] + 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) -> 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.""" @@ -41,8 +357,9 @@ class ArrowScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]): 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: + 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( @@ -51,7 +368,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]": + def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: import pyarrow as pa for batch in data: 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/eval_handlers/tests/test_arrow_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py new file mode 100644 index 0000000000000..79e876ed0df6e --- /dev/null +++ b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py @@ -0,0 +1,291 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for the 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 +from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message +from pyspark.util import PythonEvalType + +if have_pyarrow: + import pyarrow as pa + + from pyspark.eval_handlers._arrow import ( + ArrowCoGroupedMapUDFHandler, + ArrowGroupedMapIterUDFHandler, + ArrowGroupedMapUDFHandler, + ArrowMapUDFHandler, + ArrowScalarIterUDFHandler, + ArrowScalarUDFHandler, + ) + from pyspark.sql.conversion import ArrowBatchTransformer + + +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 + + +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): + # 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) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowScalarUDFHandlerTests(unittest.TestCase): + 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()) + + 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]]) + + def test_coerces_output_to_return_type(self): + # The UDF returns int32, but the declared return type is LongType (int64). + def add_one(col): + return pa.array([v.as_py() + 1 for v in col], type=pa.int32()) + + 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_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()) + + 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_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()) + + handler = _scalar_handler(ArrowScalarIterUDFHandler, too_many) + with self.assertRaises(PySparkRuntimeError): + list(handler.run(0, iter([_batch(a=[1, 2])]))) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowMapUDFHandlerTests(unittest.TestCase): + 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: + yield _batch(v=[c.as_py() * 2 for c in batch.column("v")]) + + 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): + 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")])}) + + 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_passes_key_when_udf_takes_key(self): + def grouped_udf(key, value_table): + k = key[0].as_py() + return pa.table({"v": pa.array([c.as_py() + k for c in value_table.column("v")])}) + + 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_applies_udf_per_group(self): + def grouped_udf(value_batches): + for batch in value_batches: + yield _batch(v=[c.as_py() + 1 for c in batch.column("v")]) + + 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): + # 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("v")[0].as_py() + right_values.column("v")[0].as_py() + return pa.table({"v": pa.array([total], type=pa.int64())}) + + 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) +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 + + from pyspark.serializers import write_int + + def arrow_bytes(batches): + buf = io.BytesIO() + ArrowStreamSerializer().dump_stream(iter(batches), buf) + return buf.getvalue() + + stream = io.BytesIO() + write_int(2, stream) # two DataFrames in the co-group + 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) + + 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 + + main() diff --git a/python/pyspark/tests/test_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py similarity index 50% rename from python/pyspark/tests/test_eval_type_handlers.py rename to python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py index 4f7cd96a2d899..c5fff82312ac5 100644 --- a/python/pyspark/tests/test_eval_type_handlers.py +++ b/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py @@ -15,9 +15,14 @@ # 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._arrow import ArrowScalarUDFHandler from pyspark.eval_handlers._base import ( BatchEvalTypeHandler, CoGroupedEvalTypeHandler, @@ -31,9 +36,6 @@ ArrowStreamGroupSerializer, ArrowStreamSerializer, ) -from pyspark.sql.types import LongType -from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message -from pyspark.util import PythonEvalType class _RunnerConf: @@ -41,15 +43,11 @@ 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_category_bases_are_abstract(self): # The interface and the three category bases must not be instantiable: # they leave ``run`` abstract. @@ -104,19 +102,28 @@ def run(self, split_index, data): 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, - ) + 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 @@ -132,78 +139,6 @@ class _Abstract(BatchEvalTypeHandler["pa.RecordBatch"]): self.assertNotIn(unused_eval_type, _eval_type_handlers) -@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()) - - 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]))) - - 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): - 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): - 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). - 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 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/eval_handlers/utils.py b/python/pyspark/eval_handlers/utils.py new file mode 100644 index 0000000000000..48a7570203374 --- /dev/null +++ b/python/pyspark/eval_handlers/utils.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/eval_handlers/verification.py b/python/pyspark/eval_handlers/verification.py index f63097fae7804..3b8399984dee7 100644 --- a/python/pyspark/eval_handlers/verification.py +++ b/python/pyspark/eval_handlers/verification.py @@ -20,10 +20,59 @@ 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, cast, get_args, get_origin from pyspark.errors import PySparkRuntimeError, PySparkTypeError +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.""" @@ -37,6 +86,48 @@ def verify_result_row_count(result_length: int, expected: int) -> None: ) +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 cast(T, map(check_element, result)) + + 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_scalar_result(result: Any, num_rows: int) -> Any: """ Verify a scalar UDF result is array-like and has the expected number of rows. diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 28d2ae348ec40..b3a99ee9ff8a9 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 @@ -59,8 +53,13 @@ ) 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.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 @@ -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 @@ -3645,7 +3243,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 = (