You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/09/23 04:20:43 UTC
(spark) branch master updated: [SPARK-59704][PYTHON] Migrate Arrow aggregate and window UDF eval types to the EvalTypeHandler pipeline
This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new de8be6d566db [SPARK-59704][PYTHON] Migrate Arrow aggregate and window UDF eval types to the EvalTypeHandler pipeline
de8be6d566db is described below
commit de8be6d566dbf7d795f8a58512024f15ed4c000e
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Wed Sep 23 04:20:28 2026 +0000
[SPARK-59704][PYTHON] Migrate Arrow aggregate and window UDF eval types to the EvalTypeHandler pipeline
### What changes were proposed in this pull request?
This PR migrates the three non-incremental Arrow-native aggregate/window UDF eval types off the hand-written `if/elif` chain in `worker.py`'s `read_udfs` and onto the `EvalTypeHandler` pipeline introduced in SPARK-59425 and extended in SPARK-59622. The migrated types are `SQL_GROUPED_AGG_ARROW_UDF` (252), `SQL_GROUPED_AGG_ARROW_ITER_UDF` (254), and `SQL_WINDOW_AGG_ARROW_UDF` (253).
Each becomes a `GroupedEvalTypeHandler` subclass in `python/pyspark/eval_handlers/_arrow.py` (`ArrowGroupedAggUDFHandler`, `ArrowGroupedAggIterUDFHandler`, `ArrowWindowAggUDFHandler`), declaring its `eval_type` so it self-registers and is dispatched via `get_eval_type_handler`. The handlers follow the existing lazy-pyarrow convention: string forward-ref bases, `require_minimum_pyarrow_version()` in `__init__`, and `import pyarrow as pa` inside `run`. The window handler keeps the unbou [...]
This is a pure code move with no behavior change: the UDF logic, output coercion, and stream draining are identical to the removed branches.
### Why are the changes needed?
Consolidating the Arrow eval types onto one registry-dispatched pipeline removes the growing per-eval-type `if/elif` chain in `read_udfs`, so the execution model for each type lives in one cohesive handler class rather than being split between a serializer-selection block and a function-construction block. Follows the umbrella SPARK-59415.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing per-eval-type tests, classic and Spark Connect: `pyspark.sql.tests.arrow.test_arrow_udf_grouped_agg`, `pyspark.sql.tests.arrow.test_arrow_udf_window`, and their `connect.arrow.test_parity_*` counterparts all pass. Worker microbenchmarks for the three eval types (including the window bounded and unbounded frames) show no regression versus the pre-refactor baseline.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac
This pull request and its description were written by Isaac.
Closes #58956 from Yicong-Huang/SPARK-59415-arrow-agg-window.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/eval_handlers/_arrow.py | 158 +++++++++++++++++++++
.../tests/test_arrow_eval_type_handlers.py | 69 ++++++++-
python/pyspark/worker.py | 148 +------------------
python/pyspark/worker_util.py | 6 +
4 files changed, 234 insertions(+), 147 deletions(-)
diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py
index 62e36b9ff792..21e10b68f941 100644
--- a/python/pyspark/eval_handlers/_arrow.py
+++ b/python/pyspark/eval_handlers/_arrow.py
@@ -31,6 +31,7 @@ import itertools
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,163 @@ class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
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)
+ 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)]),
+ 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/eval_handlers/tests/test_arrow_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py
index 79e876ed0df6..fb84b1020f8d 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 @@ if have_pyarrow:
from pyspark.eval_handlers._arrow import (
ArrowCoGroupedMapUDFHandler,
+ ArrowGroupedAggIterUDFHandler,
+ ArrowGroupedAggUDFHandler,
ArrowGroupedMapIterUDFHandler,
ArrowGroupedMapUDFHandler,
ArrowMapUDFHandler,
ArrowScalarIterUDFHandler,
ArrowScalarUDFHandler,
+ ArrowWindowAggUDFHandler,
)
from pyspark.sql.conversion import ArrowBatchTransformer
@@ -52,6 +55,8 @@ class _RunnerConf:
use_large_var_types = False
assign_cols_by_name = True
map_in_batch_legacy_accept_any_iterable = False
+ # One bound type per window UDF; the bounded test overrides this per instance.
+ window_bound_types = ["unbounded"]
def _batch(**columns):
@@ -120,7 +125,7 @@ _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.
+ # 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 +133,9 @@ class ArrowEvalTypeHandlerRegistrationTests(unittest.TestCase):
(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 +263,65 @@ class ArrowCoGroupedMapUDFHandlerTests(unittest.TestCase):
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):
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 40439a497785..bf8117a13174 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 read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
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
@@ -2260,8 +2118,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
# ``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(
@@ -2349,8 +2206,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
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 3a24bde38116..f02f0cd8dc44 100644
--- a/python/pyspark/worker_util.py
+++ b/python/pyspark/worker_util.py
@@ -362,6 +362,12 @@ class RunnerConf(Conf):
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))
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org