From 9b99be5aad4be4c493510988ddacba59fb175c52 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 18 Sep 2026 05:31:18 +0000 Subject: [PATCH 1/4] refactor: migrate Arrow agg/window UDF eval types to EvalTypeHandler pipeline --- python/pyspark/eval_handlers/_arrow.py | 159 +++++++++++++++++++++++++ python/pyspark/worker.py | 142 ---------------------- 2 files changed, 159 insertions(+), 142 deletions(-) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 62e36b9ff792c..9f32144ba2f41 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -31,6 +31,7 @@ from collections.abc import Iterator from typing import TYPE_CHECKING, Any +from pyspark.errors import PySparkRuntimeError from pyspark.eval_handlers._base import ( BatchEvalTypeHandler, CoGroupedEvalTypeHandler, @@ -250,6 +251,164 @@ def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.Rec yield ArrowBatchTransformer.wrap_struct(batch) +def _concat_group_batches(batch_list: list["pa.RecordBatch"]) -> "pa.RecordBatch": + """Concatenate a group's RecordBatches into a single one, with a fallback for + pyarrow before 19.0.0 (which lacks ``pa.concat_batches``). Remove the fallback + once support for those versions is dropped.""" + import pyarrow as pa + + if hasattr(pa, "concat_batches"): + return pa.concat_batches(batch_list) + return pa.RecordBatch.from_struct_array( + pa.concat_arrays([b.to_struct_array() for b in batch_list]) + ) + + +class ArrowGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): + """SQL_GROUPED_AGG_ARROW_UDF: each UDF reduces its input columns over the whole + group to a single scalar; emit one row per group with one column per UDF, + coerced to the declared schema.""" + + eval_type = PythonEvalType.SQL_GROUPED_AGG_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) + self._col_names = ["_%d" % i for i in range(len(udfs))] + self._return_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[GroupedBatch]) -> Iterator[pa.RecordBatch]: + import pyarrow as pa + + for group in data: + batch_list = list(group) + if not batch_list: + continue + concatenated = _concat_group_batches(batch_list) + results = [ + udf_func( + *[concatenated.column(o) for o in args_offsets], + **{k: concatenated.column(v) for k, v in kwargs_offsets.items()}, + ) + for udf_func, args_offsets, kwargs_offsets, _ in self._udfs + ] + result_arrays = [pa.array([r]) for r in results] + batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names) + yield ArrowBatchTransformer.enforce_schema(batch, self._return_schema) + + +class ArrowGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): + """SQL_GROUPED_AGG_ARROW_ITER_UDF: the single UDF receives each group as an + iterator of its input columns and returns one scalar; emit one row per group.""" + + eval_type = PythonEvalType.SQL_GROUPED_AGG_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_AGG_ARROW_ITER UDF expected here." + self._udf_func, self._args_offsets, _, return_type = udfs[0] + self._return_schema = to_arrow_schema( + StructType([StructField("_0", 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]: + import pyarrow as pa + + args_offsets = self._args_offsets + + def extract_args(batch: pa.RecordBatch) -> Any: + args = tuple(batch.column(o) for o in args_offsets) + return args[0] if len(args) == 1 else args + + for group in data: + batch_iter = map(extract_args, group) + result = self._udf_func(batch_iter) + # Drain remaining batches to maintain stream position + for _ in batch_iter: + pass + batch = pa.RecordBatch.from_arrays([pa.array([result])], ["_0"]) + yield ArrowBatchTransformer.enforce_schema(batch, self._return_schema) + + +class ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]): + """SQL_WINDOW_AGG_ARROW_UDF: each UDF produces one value per input row over its + window frame -- an unbounded frame is computed once and repeated for every row, + a bounded frame slices the input columns per row -- emitted as one column per + UDF, coerced to the declared schema.""" + + eval_type = PythonEvalType.SQL_WINDOW_AGG_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) + window_bound_types_str = runner_conf.get("window_bound_types") + self._window_bound_types = [t.strip().lower() for t in window_bound_types_str.split(",")] + self._col_names = ["_%d" % i for i in range(len(udfs))] + self._return_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[GroupedBatch]) -> Iterator[pa.RecordBatch]: + import pyarrow as pa + + for group in data: + batch_list = list(group) + if not batch_list: + continue + concatenated = _concat_group_batches(batch_list) + num_rows = concatenated.num_rows + + result_arrays = [] + for udf_index, (udf_func, args_offsets, kwargs_offsets, _) in enumerate(self._udfs): + bound_type = self._window_bound_types[udf_index] + if bound_type == "unbounded": + result = udf_func( + *[concatenated.column(o) for o in args_offsets], + **{k: concatenated.column(v) for k, v in kwargs_offsets.items()}, + ) + result_arrays.append(pa.repeat(result, num_rows)) + elif bound_type == "bounded": + begin_col = concatenated.column(args_offsets[0]) + end_col = concatenated.column(args_offsets[1]) + results = [] + for i in range(num_rows): + offset = begin_col[i].as_py() + length = end_col[i].as_py() - offset + slices = [ + concatenated.column(o).slice(offset=offset, length=length) + for o in args_offsets[2:] + ] + kw_slices = { + k: concatenated.column(v).slice(offset=offset, length=length) + for k, v in kwargs_offsets.items() + } + results.append(udf_func(*slices, **kw_slices)) + result_arrays.append(pa.array(results)) + else: + raise PySparkRuntimeError( + errorClass="INVALID_WINDOW_BOUND_TYPE", + messageParameters={"window_bound_type": bound_type}, + ) + + batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names) + yield ArrowBatchTransformer.enforce_schema(batch, self._return_schema) + + 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 diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 40439a497785e..c13714d271585 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -1817,11 +1817,8 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, - PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, - PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF, PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, - PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE, PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF, PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF, @@ -1833,8 +1830,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): ): # NOTE: if timezone is set here, that implies respectSessionTimeZone is True if eval_type in ( - PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF, - PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF, # The map-side PARTIAL stage streams ordinary (multi-group) batches and hash-combines @@ -1843,7 +1838,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF, PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF, - PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF, PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF, ): @@ -1863,75 +1857,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf): num_udfs = len(udfs) - if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF: - import pyarrow as pa - - # Pre-compute target schema for output coercion - col_names = ["_%d" % i for i in range(len(udfs))] - return_schema = to_arrow_schema( - StructType([StructField(name, rt) for name, (_, _, _, rt) in zip(col_names, udfs)]), - timezone="UTC", - prefers_large_types=runner_conf.use_large_var_types, - ) - - def grouped_func( - split_index: int, data: Iterator["GroupedBatch"] - ) -> Iterator[pa.RecordBatch]: - for group in data: - batch_list = list(group) - if not batch_list: - continue - if hasattr(pa, "concat_batches"): - concatenated = pa.concat_batches(batch_list) - else: - # pyarrow.concat_batches not supported before 19.0.0 - # remove this once we drop support for old versions - concatenated = pa.RecordBatch.from_struct_array( - pa.concat_arrays([b.to_struct_array() for b in batch_list]) - ) - results = [ - udf_func( - *[concatenated.column(o) for o in args_offsets], - **{k: concatenated.column(v) for k, v in kwargs_offsets.items()}, - ) - for udf_func, args_offsets, kwargs_offsets, _ in udfs - ] - result_arrays = [pa.array([r]) for r in results] - batch = pa.RecordBatch.from_arrays(result_arrays, col_names) - yield ArrowBatchTransformer.enforce_schema(batch, return_schema) - - return grouped_func, ser - - if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF: - import pyarrow as pa - - assert num_udfs == 1, "One GROUPED_AGG_ARROW_ITER UDF expected here." - udf_func, args_offsets, kwargs_offsets, return_type = udfs[0] - - return_schema = to_arrow_schema( - StructType([StructField("_0", return_type)]), - timezone="UTC", - prefers_large_types=runner_conf.use_large_var_types, - ) - - def extract_args(batch): - args = tuple(batch.column(o) for o in args_offsets) - return args[0] if len(args) == 1 else args - - def grouped_func( - split_index: int, data: Iterator["GroupedBatch"] - ) -> Iterator[pa.RecordBatch]: - for group in data: - batch_iter = map(extract_args, group) - result = udf_func(batch_iter) - # Drain remaining batches to maintain stream position - for _ in batch_iter: - pass - batch = pa.RecordBatch.from_arrays([pa.array([result])], ["_0"]) - yield ArrowBatchTransformer.enforce_schema(batch, return_schema) - - return grouped_func, ser - if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF: import pyarrow as pa @@ -2184,73 +2109,6 @@ def grouped_func( return grouped_func, ser - if eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF: - import pyarrow as pa - - window_bound_types_str = runner_conf.get("window_bound_types") - window_bound_types = [t.strip().lower() for t in window_bound_types_str.split(",")] - - col_names = ["_%d" % i for i in range(len(udfs))] - return_schema = to_arrow_schema( - StructType([StructField(name, rt) for name, (_, _, _, rt) in zip(col_names, udfs)]), - timezone="UTC", - prefers_large_types=runner_conf.use_large_var_types, - ) - - def grouped_func( - split_index: int, data: Iterator["GroupedBatch"] - ) -> Iterator[pa.RecordBatch]: - for group in data: - batch_list = list(group) - if not batch_list: - continue - if hasattr(pa, "concat_batches"): - concatenated = pa.concat_batches(batch_list) - else: - # pyarrow.concat_batches not supported before 19.0.0 - # remove this once we drop support for old versions - concatenated = pa.RecordBatch.from_struct_array( - pa.concat_arrays([b.to_struct_array() for b in batch_list]) - ) - num_rows = concatenated.num_rows - - result_arrays = [] - for udf_index, (udf_func, args_offsets, kwargs_offsets, _) in enumerate(udfs): - bound_type = window_bound_types[udf_index] - if bound_type == "unbounded": - result = udf_func( - *[concatenated.column(o) for o in args_offsets], - **{k: concatenated.column(v) for k, v in kwargs_offsets.items()}, - ) - result_arrays.append(pa.repeat(result, num_rows)) - elif bound_type == "bounded": - begin_col = concatenated.column(args_offsets[0]) - end_col = concatenated.column(args_offsets[1]) - results = [] - for i in range(num_rows): - offset = begin_col[i].as_py() - length = end_col[i].as_py() - offset - slices = [ - concatenated.column(o).slice(offset=offset, length=length) - for o in args_offsets[2:] - ] - kw_slices = { - k: concatenated.column(v).slice(offset=offset, length=length) - for k, v in kwargs_offsets.items() - } - results.append(udf_func(*slices, **kw_slices)) - result_arrays.append(pa.array(results)) - else: - raise PySparkRuntimeError( - errorClass="INVALID_WINDOW_BOUND_TYPE", - messageParameters={"window_bound_type": bound_type}, - ) - - batch = pa.RecordBatch.from_arrays(result_arrays, col_names) - yield ArrowBatchTransformer.enforce_schema(batch, return_schema) - - return grouped_func, ser - if eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF: import pyarrow as pa From 5a2928775abc474588fa49c820f7d2f05cb7c6d2 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Tue, 22 Sep 2026 01:24:40 +0000 Subject: [PATCH 2/4] test: add handler unit tests for Arrow agg/window eval types --- .../tests/test_arrow_eval_type_handlers.py | 72 ++++++++++++++++++- 1 file changed, 71 insertions(+), 1 deletion(-) 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 index 79e876ed0df6e..33e9d808f4485 100644 --- a/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py +++ b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py @@ -36,11 +36,14 @@ from pyspark.eval_handlers._arrow import ( ArrowCoGroupedMapUDFHandler, + ArrowGroupedAggIterUDFHandler, + ArrowGroupedAggUDFHandler, ArrowGroupedMapIterUDFHandler, ArrowGroupedMapUDFHandler, ArrowMapUDFHandler, ArrowScalarIterUDFHandler, ArrowScalarUDFHandler, + ArrowWindowAggUDFHandler, ) from pyspark.sql.conversion import ArrowBatchTransformer @@ -52,6 +55,11 @@ class _RunnerConf: use_large_var_types = False assign_cols_by_name = True map_in_batch_legacy_accept_any_iterable = False + # Read by the window-agg handler via ``get``; overridden per instance for the bounded case. + window_bound_types = "unbounded" + + def get(self, key, default="", *, lower_str=True): + return getattr(self, key, default) def _batch(**columns): @@ -120,7 +128,7 @@ def _grouped_handler(handler_cls, udf, arg_offsets, num_udf_args): @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. + # Every migrated Arrow 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), @@ -128,6 +136,9 @@ def test_arrow_eval_types_are_registered(self): (PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF, ArrowGroupedMapUDFHandler), (PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF, ArrowGroupedMapIterUDFHandler), (PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, ArrowCoGroupedMapUDFHandler), + (PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF, ArrowGroupedAggUDFHandler), + (PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF, ArrowGroupedAggIterUDFHandler), + (PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF, ArrowWindowAggUDFHandler), ): self.assertIs(get_eval_type_handler(eval_type), handler_cls) @@ -255,6 +266,65 @@ def cogrouped_udf(key, left_values, right_values): self.assertEqual(out[0].column(0).field("v").to_pylist(), [35]) +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowGroupedAggUDFHandlerTests(unittest.TestCase): + # Grouped-agg batches arrive un-wrapped (flat columns); the UDF reads columns by offset and + # returns one scalar, emitted as a single-row batch per group. + def test_reduces_group_to_one_row(self): + def sum_udf(col): + return sum(c.as_py() for c in col) + + handler = ArrowGroupedAggUDFHandler( + udfs=[(sum_udf, [0], {}, LongType())], runner_conf=_RunnerConf(), eval_conf=None + ) + out = list(handler.run(0, _one_group(_batch(v=[1, 2, 3]), _batch(v=[4])))) + self.assertEqual(out[0].column("_0").to_pylist(), [10]) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowGroupedAggIterUDFHandlerTests(unittest.TestCase): + # The UDF receives the group's input columns as an iterator and returns one scalar. + def test_reduces_group_to_one_row(self): + def sum_iter_udf(col_iter): + return sum(c.as_py() for col in col_iter for c in col) + + handler = ArrowGroupedAggIterUDFHandler( + udfs=[(sum_iter_udf, [0], {}, LongType())], runner_conf=_RunnerConf(), eval_conf=None + ) + out = list(handler.run(0, _one_group(_batch(v=[10]), _batch(v=[20])))) + self.assertEqual(out[0].column("_0").to_pylist(), [30]) + + +@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) +class ArrowWindowAggUDFHandlerTests(unittest.TestCase): + # One output value per input row over the UDF's window frame. + def test_unbounded_frame_repeats_one_value(self): + def sum_udf(col): + return sum(c.as_py() for c in col) + + conf = _RunnerConf() + conf.window_bound_types = "unbounded" + handler = ArrowWindowAggUDFHandler( + udfs=[(sum_udf, [0], {}, LongType())], runner_conf=conf, eval_conf=None + ) + out = list(handler.run(0, _one_group(_batch(v=[1, 2, 3])))) + self.assertEqual(out[0].column("_0").to_pylist(), [6, 6, 6]) + + def test_bounded_frame_slices_per_row(self): + # args_offsets = [begin_col, end_col, *value_cols]; each row's frame is ``[begin, end)``. + def sum_udf(col): + return sum(c.as_py() for c in col) + + conf = _RunnerConf() + conf.window_bound_types = "bounded" + handler = ArrowWindowAggUDFHandler( + udfs=[(sum_udf, [0, 1, 2], {}, LongType())], runner_conf=conf, eval_conf=None + ) + # Row 0 frame [0, 1) -> [10]; row 1 frame [0, 2) -> [10, 20]. + out = list(handler.run(0, _one_group(_batch(begin=[0, 0], end=[1, 2], v=[10, 20])))) + self.assertEqual(out[0].column("_0").to_pylist(), [10, 30]) + + @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class CoGroupedBatchTests(unittest.TestCase): def test_deserialized_co_group_is_a_pair_of_lists(self): From 8d6670d7d84eb3982aa7b3822df7f1eb39024861 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Tue, 22 Sep 2026 04:12:05 +0000 Subject: [PATCH 3/4] test: drop unused lower_str param from test RunnerConf.get --- .../eval_handlers/tests/test_arrow_eval_type_handlers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 index 33e9d808f4485..ac66ce1f5f575 100644 --- a/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py +++ b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py @@ -58,7 +58,7 @@ class _RunnerConf: # Read by the window-agg handler via ``get``; overridden per instance for the bounded case. window_bound_types = "unbounded" - def get(self, key, default="", *, lower_str=True): + def get(self, key, default=""): return getattr(self, key, default) From 31c6952fa85b2dd74071e450526534e4c2f072c4 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Tue, 22 Sep 2026 04:18:34 +0000 Subject: [PATCH 4/4] refactor: add window_bound_types property to RunnerConf --- python/pyspark/eval_handlers/_arrow.py | 3 +-- .../tests/test_arrow_eval_type_handlers.py | 11 ++++------- python/pyspark/worker.py | 6 ++---- python/pyspark/worker_util.py | 6 ++++++ 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py index 9f32144ba2f41..21e10b68f9413 100644 --- a/python/pyspark/eval_handlers/_arrow.py +++ b/python/pyspark/eval_handlers/_arrow.py @@ -354,8 +354,7 @@ def __init__( ) -> None: require_minimum_pyarrow_version() super().__init__(udfs, runner_conf, eval_conf) - window_bound_types_str = runner_conf.get("window_bound_types") - self._window_bound_types = [t.strip().lower() for t in window_bound_types_str.split(",")] + self._window_bound_types = runner_conf.window_bound_types self._col_names = ["_%d" % i for i in range(len(udfs))] self._return_schema = to_arrow_schema( StructType([StructField(n, rt) for n, (_, _, _, rt) in zip(self._col_names, udfs)]), 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 index ac66ce1f5f575..fb84b1020f8d2 100644 --- a/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py +++ b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py @@ -55,11 +55,8 @@ class _RunnerConf: use_large_var_types = False assign_cols_by_name = True map_in_batch_legacy_accept_any_iterable = False - # Read by the window-agg handler via ``get``; overridden per instance for the bounded case. - window_bound_types = "unbounded" - - def get(self, key, default=""): - return getattr(self, key, default) + # One bound type per window UDF; the bounded test overrides this per instance. + window_bound_types = ["unbounded"] def _batch(**columns): @@ -303,7 +300,7 @@ def sum_udf(col): return sum(c.as_py() for c in col) conf = _RunnerConf() - conf.window_bound_types = "unbounded" + conf.window_bound_types = ["unbounded"] handler = ArrowWindowAggUDFHandler( udfs=[(sum_udf, [0], {}, LongType())], runner_conf=conf, eval_conf=None ) @@ -316,7 +313,7 @@ def sum_udf(col): return sum(c.as_py() for c in col) conf = _RunnerConf() - conf.window_bound_types = "bounded" + conf.window_bound_types = ["bounded"] handler = ArrowWindowAggUDFHandler( udfs=[(sum_udf, [0, 1, 2], {}, LongType())], runner_conf=conf, eval_conf=None ) diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index c13714d271585..bf8117a13174c 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -2118,8 +2118,7 @@ def grouped_func( # ``zero``) and produces the value with ``finish``, one output value per input row. A window # has no shuffle, so the intermediate buffer never leaves the worker (unlike the two-stage # groupBy path); ``merge`` is not used here. - window_bound_types_str = runner_conf.get("window_bound_types") - window_bound_types = [t.strip().lower() for t in window_bound_types_str.split(",")] + window_bound_types = runner_conf.window_bound_types col_names = ["_%d" % i for i in range(len(udfs))] return_schema = to_arrow_schema( @@ -2207,8 +2206,7 @@ def grouped_func( import pandas as pd import pyarrow as pa - window_bound_types_str = runner_conf.get("window_bound_types") - window_bound_types = [t.strip().lower() for t in window_bound_types_str.split(",")] + window_bound_types = runner_conf.window_bound_types col_names = ["_%d" % i for i in range(len(udfs))] output_schema = StructType( diff --git a/python/pyspark/worker_util.py b/python/pyspark/worker_util.py index 3a24bde38116b..f02f0cd8dc443 100644 --- a/python/pyspark/worker_util.py +++ b/python/pyspark/worker_util.py @@ -362,6 +362,12 @@ def prefer_int_ext_dtype(self) -> bool: def timezone(self) -> Optional[str]: return self.get("spark.sql.session.timeZone", None, lower_str=False) + @property + def window_bound_types(self) -> list[str]: + # Per-UDF window frame bound type ("unbounded" or "bounded"), one entry per window UDF, + # sent comma-separated by the window operator. + return [t.strip().lower() for t in self.get("window_bound_types").split(",")] + @property def arrow_max_records_per_batch(self) -> int: return int(self.get("spark.sql.execution.arrow.maxRecordsPerBatch", 10000))