You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2026/04/02 03:42:30 UTC
(spark) branch master updated: Revert "[SPARK-56189][PYTHON] Refactor SQL_WINDOW_AGG_ARROW_UDF"
This is an automated email from the ASF dual-hosted git repository.
ruifengz 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 0d8e031b82fa Revert "[SPARK-56189][PYTHON] Refactor SQL_WINDOW_AGG_ARROW_UDF"
0d8e031b82fa is described below
commit 0d8e031b82fa41a2a9306630c300dc1a29564fc7
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Thu Apr 2 11:42:15 2026 +0800
Revert "[SPARK-56189][PYTHON] Refactor SQL_WINDOW_AGG_ARROW_UDF"
This reverts commit 3433c386089ccc9c9e2b24a3faf0dcd8f0984b30.
to restore CI https://github.com/apache/spark/actions/runs/23879531895/job/69630001459
Closes #55152 from zhengruifeng/revert-56189.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
python/pyspark/worker.py | 179 +++++++++++++++++++++++++++++------------------
1 file changed, 111 insertions(+), 68 deletions(-)
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 3d70a9d77561..bc4eae2e2148 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -68,6 +68,7 @@ from pyspark.sql.pandas.serializers import (
TransformWithStateInPySparkRowSerializer,
TransformWithStateInPySparkRowInitStateSerializer,
ArrowStreamAggPandasUDFSerializer,
+ ArrowStreamAggArrowUDFSerializer,
ArrowBatchUDFSerializer,
ArrowStreamUDTFSerializer,
ArrowStreamArrowUDTFSerializer,
@@ -1088,6 +1089,26 @@ def wrap_window_agg_pandas_udf(
)
+def wrap_window_agg_arrow_udf(f, args_offsets, kwargs_offsets, return_type, runner_conf, udf_index):
+ window_bound_types_str = runner_conf.get("window_bound_types")
+ window_bound_type = [t.strip().lower() for t in window_bound_types_str.split(",")][udf_index]
+ if window_bound_type == "bounded":
+ return wrap_bounded_window_agg_arrow_udf(
+ f, args_offsets, kwargs_offsets, return_type, runner_conf
+ )
+ elif window_bound_type == "unbounded":
+ return wrap_unbounded_window_agg_arrow_udf(
+ f, args_offsets, kwargs_offsets, return_type, runner_conf
+ )
+ else:
+ raise PySparkRuntimeError(
+ errorClass="INVALID_WINDOW_BOUND_TYPE",
+ messageParameters={
+ "window_bound_type": window_bound_type,
+ },
+ )
+
+
def wrap_unbounded_window_agg_pandas_udf(f, args_offsets, kwargs_offsets, return_type, runner_conf):
func, args_kwargs_offsets = wrap_kwargs_support(f, args_offsets, kwargs_offsets)
@@ -1107,6 +1128,27 @@ def wrap_unbounded_window_agg_pandas_udf(f, args_offsets, kwargs_offsets, return
)
+def wrap_unbounded_window_agg_arrow_udf(f, args_offsets, kwargs_offsets, return_type, runner_conf):
+ func, args_kwargs_offsets = wrap_kwargs_support(f, args_offsets, kwargs_offsets)
+
+ # This is similar to wrap_unbounded_window_agg_pandas_udf, the only difference
+ # is that this function is for arrow udf.
+ arrow_return_type = to_arrow_type(
+ return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types
+ )
+
+ def wrapped(*series):
+ import pyarrow as pa
+
+ result = func(*series)
+ return pa.repeat(result, len(series[0]))
+
+ return (
+ args_kwargs_offsets,
+ lambda *a: (wrapped(*a), arrow_return_type),
+ )
+
+
def wrap_bounded_window_agg_pandas_udf(f, args_offsets, kwargs_offsets, return_type, runner_conf):
# args_offsets should have at least 2 for begin_index, end_index.
assert len(args_offsets) >= 2, len(args_offsets)
@@ -1147,6 +1189,35 @@ def wrap_bounded_window_agg_pandas_udf(f, args_offsets, kwargs_offsets, return_t
)
+def wrap_bounded_window_agg_arrow_udf(f, args_offsets, kwargs_offsets, return_type, runner_conf):
+ # args_offsets should have at least 2 for begin_index, end_index.
+ assert len(args_offsets) >= 2, len(args_offsets)
+ func, args_kwargs_offsets = wrap_kwargs_support(f, args_offsets[2:], kwargs_offsets)
+
+ arrow_return_type = to_arrow_type(
+ return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types
+ )
+
+ def wrapped(begin_index, end_index, *series):
+ import pyarrow as pa
+
+ assert isinstance(begin_index, pa.Int32Array), type(begin_index)
+ assert isinstance(end_index, pa.Int32Array), type(end_index)
+
+ result = []
+ for i in range(len(begin_index)):
+ offset = begin_index[i].as_py()
+ length = end_index[i].as_py() - offset
+ series_slices = [s.slice(offset=offset, length=length) for s in series]
+ result.append(func(*series_slices))
+ return pa.array(result)
+
+ return (
+ args_offsets[:2] + args_kwargs_offsets,
+ lambda *a: (wrapped(*a), arrow_return_type),
+ )
+
+
def wrap_kwargs_support(f, args_offsets, kwargs_offsets):
if len(kwargs_offsets):
keys = list(kwargs_offsets.keys())
@@ -1361,7 +1432,9 @@ def read_single_udf(pickleSer, infile, eval_type, runner_conf, udf_index):
func, args_offsets, kwargs_offsets, return_type, runner_conf, udf_index
)
elif eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF:
- return func, args_offsets, kwargs_offsets, return_type
+ return wrap_window_agg_arrow_udf(
+ func, args_offsets, kwargs_offsets, return_type, runner_conf, udf_index
+ )
elif eval_type == PythonEvalType.SQL_BATCHED_UDF:
return wrap_udf(func, args_offsets, kwargs_offsets, return_type)
else:
@@ -2575,7 +2648,7 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
):
ser = ArrowStreamGroupSerializer(write_start_stream=True)
elif eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF:
- ser = ArrowStreamSerializer(write_start_stream=True, num_dfs=1)
+ ser = ArrowStreamAggArrowUDFSerializer(safecheck=True, arrow_cast=True)
elif eval_type in (
PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
@@ -2896,72 +2969,6 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
# profiling is not supported for UDF
return func, None, ser, 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 func(split_index: int, batches: Iterator[Any]) -> Iterator[pa.RecordBatch]:
- for group_batches in batches:
- batch_list = list(group_batches)
- 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)
-
- # profiling is not supported for UDF
- return func, None, ser, ser
-
is_scalar_iter = eval_type == PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
is_map_pandas_iter = eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
@@ -3386,6 +3393,42 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
)
return f(series_iter)
+ elif eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF:
+ import pyarrow as pa
+
+ # For SQL_WINDOW_AGG_ARROW_UDF,
+ # convert iterator of batch columns to a concatenated RecordBatch
+ def mapper(a):
+ # a is Iterator[Tuple[pa.Array, ...]] - convert to RecordBatch
+ batches = []
+ for batch_columns in a:
+ # batch_columns is Tuple[pa.Array, ...] - convert to RecordBatch
+ batch = pa.RecordBatch.from_arrays(
+ batch_columns, names=["_%d" % i for i in range(len(batch_columns))]
+ )
+ batches.append(batch)
+
+ # Concatenate all batches into one
+ if hasattr(pa, "concat_batches"):
+ concatenated_batch = pa.concat_batches(batches)
+ else:
+ # pyarrow.concat_batches not supported before 19.0.0
+ # remove this once we drop support for old versions
+ concatenated_batch = pa.RecordBatch.from_struct_array(
+ pa.concat_arrays([b.to_struct_array() for b in batches])
+ )
+
+ # Extract series using offsets (concatenated_batch.columns[o] gives pa.Array)
+ result = tuple(
+ f(*[concatenated_batch.columns[o] for o in arg_offsets]) for arg_offsets, f in udfs
+ )
+ # In the special case of a single UDF this will return a single result rather
+ # than a tuple of results; this is the format that the JVM side expects.
+ if len(result) == 1:
+ return result[0]
+ else:
+ return result
+
elif eval_type in (
PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org