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/10/22 23:46:18 UTC
(spark) branch branch-4.x updated: [SPARK-60357][PYTHON] Return a named record from read_single_udf for every eval type
This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new bd223d21f367 [SPARK-60357][PYTHON] Return a named record from read_single_udf for every eval type
bd223d21f367 is described below
commit bd223d21f367c33554a847b7e491852f36760aba
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Thu Oct 22 23:45:51 2026 +0000
[SPARK-60357][PYTHON] Return a named record from read_single_udf for every eval type
### What changes were proposed in this pull request?
`read_single_udf` now returns one `UDF` record (`func`, `raw_func`, `args_offsets`, `kwargs_offsets`, `return_type`) for every eval type, and `read_single_udtf` a `UDTF` record (`udtf`, `args_offsets`, `kwargs_offsets`, `return_type`). Both are frozen dataclasses, so they can only be read by field name, not unpacked positionally. Each eval type handler reads the fields it needs by name and derives its own values (`num_udf_args`, the unwrapped `Aggregator`, the SQL_BATCHED_UDF kwargs/` [...]
### Why are the changes needed?
Part of SPARK-59415. Each tuple shape was a private contract between `read_single_udf` and the one handler that consumes it; the chain only existed to pad a shared positional shape with `None` placeholders.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Updated and extended `pyspark.eval_handlers.tests` (record fields from `read_single_udf`, wrapped `func` vs. unwrapped `raw_func`, SQL_BATCHED_UDF kwargs binding and `toInternal` conversion, grouped map `num_udf_args`, incremental `Aggregator`). Ran the per-eval-type `pyspark.sql.tests` UDF/UDTF/Arrow/pandas/stateful modules locally.
Worker microbenchmark vs. master (median +- stdev of 14 runs):
| Case | Before | After | Diff |
|---|---|---|---|
| SQL_BATCHED_UDF, LongType | 72.0+-1.4ms | 64.2+-0.6ms | -10.9% |
| SQL_BATCHED_UDF, DateType (toInternal) | 90.9+-0.8ms | 89.0+-3.4ms | -2.1% |
| ScalarArrowUDFTimeBench (sm_batch_few_col, identity_udf) | 16.5+-0.7ms | 15.1+-1.0ms | -8.3% |
| ScalarPandasUDFTimeBench (sm_batch_few_col, identity_udf) | 153.6+-2.4ms | 153.6+-3.3ms | +0.0% |
| ArrowBatchedUDFTimeBench (sm_batch_few_col, identity_udf) | 29.0+-0.2ms | 29.4+-0.3ms | +1.3% |
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Opus 5.5
This pull request and its description were written by Isaac.
Closes #59604 from Yicong-Huang/SPARK-60357.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit 958019207fa06fd5f7f6c429bf99cdcb1882cbb6)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/eval_handlers/_arrow.py | 167 ++++++++++-----------
python/pyspark/eval_handlers/_base.py | 6 +-
python/pyspark/eval_handlers/_elementwise.py | 31 ++--
python/pyspark/eval_handlers/_pandas.py | 107 +++++++------
python/pyspark/eval_handlers/_stateful.py | 41 ++---
python/pyspark/eval_handlers/_typing.py | 43 +++++-
python/pyspark/eval_handlers/_udtf.py | 39 ++---
python/pyspark/eval_handlers/_vanilla.py | 41 +++--
.../tests/test_arrow_eval_type_handlers.py | 91 ++++++-----
.../tests/test_base_eval_type_handlers.py | 36 ++++-
.../tests/test_elementwise_eval_type_handlers.py | 60 +++++---
.../tests/test_pandas_eval_type_handlers.py | 61 ++++----
.../tests/test_stateful_eval_type_handlers.py | 21 ++-
.../tests/test_udtf_eval_type_handlers.py | 34 ++++-
.../tests/test_vanilla_eval_type_handlers.py | 35 ++++-
python/pyspark/worker.py | 93 ++----------
16 files changed, 503 insertions(+), 403 deletions(-)
diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py
index 217135b9577d..ed7fd87f6e3d 100644
--- a/python/pyspark/eval_handlers/_arrow.py
+++ b/python/pyspark/eval_handlers/_arrow.py
@@ -27,6 +27,7 @@ unregistered.
from __future__ import annotations
+import inspect
import itertools
import time
from collections.abc import Iterator
@@ -67,7 +68,7 @@ from pyspark.worker_metrics import WorkerMetrics
if TYPE_CHECKING:
import pyarrow as pa
- from pyspark.eval_handlers._typing import CoGroupedBatch, GroupedBatch
+ from pyspark.eval_handlers._typing import UDF, CoGroupedBatch, GroupedBatch
from pyspark.worker_util import EvalConf, RunnerConf
@@ -78,18 +79,19 @@ class ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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)
+ udf = udfs[0]
+ self._cogrouped_udf = udf.func
+ # signature was lost when wrapping it
+ self._num_udf_args = len(inspect.getfullargspec(udf.raw_func).args)
+ parsed_offsets = extract_key_value_indexes(udf.args_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
+ udf.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]:
@@ -135,21 +137,22 @@ class ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF
_iter_based = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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)
+ udf = udfs[0]
+ self._grouped_udf = udf.func
+ # signature was lost when wrapping it
+ self._num_udf_args = len(inspect.getfullargspec(udf.raw_func).args)
+ parsed_offsets = extract_key_value_indexes(udf.args_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
+ udf.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]:
@@ -205,19 +208,20 @@ class ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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)
+ udf = udfs[0]
+ self._grouped_udf = udf.func
+ # signature was lost when wrapping it
+ self._num_udf_args = len(inspect.getfullargspec(udf.raw_func).args)
+ parsed_offsets = extract_key_value_indexes(udf.args_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
+ udf.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]:
@@ -272,14 +276,12 @@ class ArrowGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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)]),
+ StructType([StructField(n, udf.return_type) for n, udf in zip(self._col_names, udfs)]),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -290,11 +292,11 @@ class ArrowGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
for group in data:
concatenated = ArrowBatchTransformer.concat_batches(group)
results = [
- udf_func(
- *[concatenated.column(o) for o in args_offsets],
- **{k: concatenated.column(v) for k, v in kwargs_offsets.items()},
+ udf.func(
+ *[concatenated.column(o) for o in udf.args_offsets],
+ **{k: concatenated.column(v) for k, v in udf.kwargs_offsets.items()},
)
- for udf_func, args_offsets, kwargs_offsets, _ in self._udfs
+ for udf in self._udfs
]
result_arrays = [pa.array([r]) for r in results]
batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names)
@@ -307,15 +309,14 @@ class ArrowGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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]
+ udf = udfs[0]
+ self._udf_func, self._args_offsets = udf.func, udf.args_offsets
self._return_schema = to_arrow_schema(
- StructType([StructField("_0", return_type)]),
+ StructType([StructField("_0", udf.return_type)]),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -347,15 +348,13 @@ class ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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)]),
+ StructType([StructField(n, udf.return_type) for n, udf in zip(self._col_names, udfs)]),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -368,7 +367,9 @@ class ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
num_rows = concatenated.num_rows
result_arrays = []
- for udf_index, (udf_func, args_offsets, kwargs_offsets, _) in enumerate(self._udfs):
+ for udf_index, udf in enumerate(self._udfs):
+ udf_func = udf.func
+ args_offsets, kwargs_offsets = udf.args_offsets, udf.kwargs_offsets
bound_type = self._window_bound_types[udf_index]
if bound_type == "unbounded":
result = udf_func(
@@ -417,9 +418,7 @@ class ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.Recor
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
# The leading `num_grouping_keys` input columns are the grouping keys (see the operator);
@@ -429,6 +428,9 @@ class ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.Recor
len(grouping_key_schema.fields) if grouping_key_schema is not None else 0
)
+ # The pickled function is the Aggregator object itself, whose zero/reduce/merge/finish
+ # methods are called directly, so take it unwrapped rather than as a plain callable.
+ self._aggs = [udf.raw_func for udf in udfs]
self._buffer_col_names = ["_%d" % i for i in range(len(udfs))]
self._buffer_arrow_types = [
to_arrow_type(
@@ -436,20 +438,17 @@ class ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.Recor
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
- for agg, _, _, _ in udfs
+ for agg in self._aggs
]
# Buffer field names are invariant across groups; compute them once per aggregator.
- self._field_names_by_udf = [
- [f.name for f in agg.bufferSchema.fields] for agg, _, _, _ in udfs
- ]
+ self._field_names_by_udf = [[f.name for f in agg.bufferSchema.fields] for agg in self._aggs]
# The aggregator's `reduce` receives a single positional tuple, so any named arguments at
# the call site are appended after the positional ones, in call order (kwargs_offsets
# preserves that order). This mirrors how a Python call `f(*args, **kwargs)` would order
# them into one value tuple.
self._input_offsets_by_udf = [
- list(args_offsets) + list(kwargs_offsets.values())
- for _, args_offsets, kwargs_offsets, _ in udfs
+ list(udf.args_offsets) + list(udf.kwargs_offsets.values()) for udf in udfs
]
# Input columns actually consumed per batch: the leading grouping keys plus every
# aggregator input, deduplicated. A UDF input may reuse a grouping-key column (the operator
@@ -511,11 +510,11 @@ class ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.Recor
hashable_key = hashable_grouping_key(key_values)
entry = groups.get(hashable_key)
if entry is None:
- buffers = [agg.zero() for agg, _, _, _ in self._udfs]
+ buffers = [agg.zero() for agg in self._aggs]
groups[hashable_key] = (key_values, buffers)
else:
buffers = entry[1]
- for i, (agg, _, _, _) in enumerate(self._udfs):
+ for i, agg in enumerate(self._aggs):
cols_i = udf_cols[i]
buffers[i] = agg.reduce(buffers[i], tuple(c[r] for c in cols_i))
if cap is not None and len(groups) >= cap:
@@ -547,39 +546,37 @@ class ArrowGroupedAggIncrementalFinalUDFHandler(GroupedEvalTypeHandler["pa.Recor
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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)]),
+ StructType([StructField(n, udf.return_type) for n, udf in zip(self._col_names, udfs)]),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
+ # The Aggregator object itself, unwrapped (see ArrowGroupedAggIncrementalPartialUDFHandler).
+ self._aggs = [udf.raw_func for udf in udfs]
# Buffer field names are invariant across groups and batches; compute once per aggregator.
- self._field_names_by_udf = [
- [f.name for f in agg.bufferSchema.fields] for agg, _, _, _ in udfs
- ]
+ self._field_names_by_udf = [[f.name for f in agg.bufferSchema.fields] for agg in self._aggs]
def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]:
import pyarrow as pa
for group in data:
- merged: list = [None] * len(self._udfs)
+ merged: list = [None] * len(self._aggs)
for batch in group:
- for i, (agg, args_offsets, _, _) in enumerate(self._udfs):
+ for i, (udf, agg) in enumerate(zip(self._udfs, self._aggs)):
field_names = self._field_names_by_udf[i]
m = merged[i]
- for row in batch.column(args_offsets[0]).to_pylist():
+ for row in batch.column(udf.args_offsets[0]).to_pylist():
if row is None:
continue
partial = tuple(row[name] for name in field_names)
m = partial if m is None else agg.merge(m, partial)
merged[i] = m
results = []
- for i, (agg, _, _, _) in enumerate(self._udfs):
+ for i, agg in enumerate(self._aggs):
m = merged[i] if merged[i] is not None else agg.zero()
results.append(agg.finish(m))
# Type each output array explicitly (mirroring the PARTIAL stage) so a non-trivial
@@ -601,18 +598,18 @@ class ArrowWindowAggIncrementalUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch
eval_type = PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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)]),
+ StructType([StructField(n, udf.return_type) for n, udf in zip(self._col_names, udfs)]),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
+ # The Aggregator object itself, unwrapped (see ArrowGroupedAggIncrementalPartialUDFHandler).
+ self._aggs = [udf.raw_func for udf in udfs]
def run(self, split_index: int, data: Iterator[GroupedBatch]) -> Iterator[pa.RecordBatch]:
import pyarrow as pa
@@ -629,7 +626,8 @@ class ArrowWindowAggIncrementalUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch
num_rows = concatenated.num_rows
result_arrays = []
- for udf_index, (agg, args_offsets, kwargs_offsets, _) in enumerate(self._udfs):
+ for udf_index, (udf, agg) in enumerate(zip(self._udfs, self._aggs)):
+ args_offsets, kwargs_offsets = udf.args_offsets, udf.kwargs_offsets
bound_type = self._window_bound_types[udf_index]
result_type = self._return_schema.field(udf_index).type
if bound_type == "unbounded":
@@ -686,13 +684,11 @@ class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF
_iter_based = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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]
+ self._udf_func = udfs[0].func
def run(self, split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
import pyarrow as pa
@@ -730,15 +726,14 @@ class ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF
_iter_based = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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]
+ udf = udfs[0]
+ self._udf_func, self._args_offsets = udf.func, udf.args_offsets
self._arrow_return_type = to_arrow_type(
- return_type, timezone="UTC", prefers_large_types=runner_conf.use_large_var_types
+ udf.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]:
@@ -786,14 +781,12 @@ 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
- ) -> None:
+ def __init__(self, udfs: list[UDF], 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(
- StructType([StructField(n, rt) for n, (_, _, _, rt) in zip(self._col_names, udfs)]),
+ StructType([StructField(n, udf.return_type) for n, udf in zip(self._col_names, udfs)]),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -804,11 +797,11 @@ class ArrowScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
for batch in data:
output_batch = pa.RecordBatch.from_arrays(
[
- udf_func(
- *[batch.column(o) for o in args_offsets],
- **{k: batch.column(v) for k, v in kwargs_offsets.items()},
+ udf.func(
+ *[batch.column(o) for o in udf.args_offsets],
+ **{k: batch.column(v) for k, v in udf.kwargs_offsets.items()},
)
- for udf_func, args_offsets, kwargs_offsets, _ in self._udfs
+ for udf in self._udfs
],
self._col_names,
)
@@ -828,17 +821,15 @@ class ArrowBatchedUDFArrowHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
variant_conf = "use_legacy_pandas_udf_conversion"
variant_value = False
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
# --- UDF preparation ---
self._udf_infos = []
- for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in udfs:
+ for udf in udfs:
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
- udf_func, udf_args_offsets, udf_kwargs_offsets
+ udf.func, udf.args_offsets, udf.kwargs_offsets
)
zero_arg = len(args_kwargs_offsets) == 0
self._udf_infos.append(
@@ -847,12 +838,12 @@ class ArrowBatchedUDFArrowHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
args_kwargs_offsets or (0,),
zero_arg,
to_arrow_type(
- udf_return_type,
+ udf.return_type,
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
),
LocalDataToArrowConversion._create_converter(
- udf_return_type,
+ udf.return_type,
none_on_identity=True,
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
),
diff --git a/python/pyspark/eval_handlers/_base.py b/python/pyspark/eval_handlers/_base.py
index ce9d2a468a2e..a4793c5ad07a 100644
--- a/python/pyspark/eval_handlers/_base.py
+++ b/python/pyspark/eval_handlers/_base.py
@@ -26,6 +26,8 @@ from collections.abc import Iterator
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Optional
from pyspark.eval_handlers._typing import (
+ UDF,
+ UDTF,
CoGroupedBatch,
GroupedBatch,
InputBatch,
@@ -149,9 +151,9 @@ class EvalTypeHandler(Generic[InputBatch, OutputBatch], metaclass=_EvalTypeHandl
_iter_based: ClassVar[bool] = False
def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
+ self, udfs: "list[UDF] | list[UDTF]", runner_conf: "RunnerConf", eval_conf: "EvalConf"
) -> None:
- self._udfs = udfs
+ self._udfs: list[Any] = udfs
self._runner_conf = runner_conf
self._eval_conf = eval_conf
diff --git a/python/pyspark/eval_handlers/_elementwise.py b/python/pyspark/eval_handlers/_elementwise.py
index cfcae6c81d30..2a836cabaf6e 100644
--- a/python/pyspark/eval_handlers/_elementwise.py
+++ b/python/pyspark/eval_handlers/_elementwise.py
@@ -62,6 +62,7 @@ from pyspark.util import PythonEvalType, fail_on_stopiteration
if TYPE_CHECKING:
import pyarrow as pa
+ from pyspark.eval_handlers._typing import UDF
from pyspark.worker_util import EvalConf, RunnerConf
@@ -179,9 +180,7 @@ class ArrowElementwiseUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
# This path exchanges data with the JVM over Arrow, so PyArrow is required. Fail with a
# clear message rather than a bare ImportError from `import pyarrow`.
require_minimum_pyarrow_version()
@@ -199,9 +198,9 @@ class ArrowElementwiseUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
nesting = eval_conf.elementwise_nesting
self._udf_infos = []
for udf_index, udf in enumerate(udfs):
- udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type = udf
+ udf_return_type = udf.return_type
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
- udf_func, udf_args_offsets, udf_kwargs_offsets
+ udf.func, udf.args_offsets, udf.kwargs_offsets
)
depth = nesting[udf_index] if nesting is not None else 1
# Each argument arrives as ``array^depth<T>``; convert its leaves with the element type
@@ -308,9 +307,7 @@ class _VectorizedElementwiseUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
whole flat columns. The conversions below are the Arrow flavor (``pa.Array`` in and out);
``_PandasElementwiseUDFHandler`` overrides them for pandas."""
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
@@ -390,17 +387,15 @@ class _ScalarElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
``array<R>`` using the input's offsets - one row in, one row out, one Python round trip per
batch. Subclasses supply the flavor's conversions and result check."""
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
super().__init__(udfs, runner_conf, eval_conf)
nesting = eval_conf.elementwise_nesting
self._udf_infos = []
for udf_index, udf in enumerate(udfs):
- udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type = udf
+ udf_return_type = udf.return_type
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
- udf_func, udf_args_offsets, udf_kwargs_offsets
+ udf.func, udf.args_offsets, udf.kwargs_offsets
)
udf_output_schema = _elementwise_output_schema(udf_return_type, runner_conf)
depth = nesting[udf_index] if nesting is not None else 1
@@ -518,14 +513,14 @@ class _ScalarIterElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
_iter_based = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One SCALAR_*_ITER_ELEMENTWISE UDF expected here."
- self._udf_func, self._args_offsets, kwargs_offsets, self._return_type = udfs[0]
- assert not kwargs_offsets, "Iterator UDFs do not take keyword arguments."
+ udf = udfs[0]
+ self._udf_func, self._args_offsets = udf.func, udf.args_offsets
+ self._return_type = udf.return_type
+ assert not udf.kwargs_offsets, "Iterator UDFs do not take keyword arguments."
# Each argument arrives as ``array^depth<T>``; the vectorized function must see the leaf
# element type ``T`` (arguments may differ, e.g. an outer column repeated into an aligned
diff --git a/python/pyspark/eval_handlers/_pandas.py b/python/pyspark/eval_handlers/_pandas.py
index ba7070657d6b..744c785520e4 100644
--- a/python/pyspark/eval_handlers/_pandas.py
+++ b/python/pyspark/eval_handlers/_pandas.py
@@ -28,6 +28,7 @@ dependency surfaces a clear error when the handler runs.
from __future__ import annotations
+import inspect
from collections.abc import Iterator
from typing import TYPE_CHECKING, Any
@@ -62,7 +63,7 @@ from pyspark.util import PythonEvalType, fail_on_stopiteration
if TYPE_CHECKING:
import pyarrow as pa
- from pyspark.eval_handlers._typing import CoGroupedBatch, GroupedBatch
+ from pyspark.eval_handlers._typing import UDF, CoGroupedBatch, GroupedBatch
from pyspark.worker_util import EvalConf, RunnerConf
@@ -75,14 +76,12 @@ class PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
self._return_schema = StructType(
- [StructField("_%d" % i, rt) for i, (_, _, _, rt) in enumerate(udfs)]
+ [StructField("_%d" % i, udf.return_type) for i, udf in enumerate(udfs)]
)
def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") -> "Iterator[pa.RecordBatch]":
@@ -103,10 +102,11 @@ class PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
# Process: evaluate each UDF column-wise on pandas Series.
results = []
- for udf_func, args_offsets, kwargs_offsets, return_type in self._udfs:
- result = udf_func(
- *[pandas_columns[o] for o in args_offsets],
- **{k: pandas_columns[v] for k, v in kwargs_offsets.items()},
+ for udf in self._udfs:
+ return_type = udf.return_type
+ result = udf.func(
+ *[pandas_columns[o] for o in udf.args_offsets],
+ **{k: pandas_columns[v] for k, v in udf.kwargs_offsets.items()},
)
if not hasattr(result, "__len__"):
pd_type = (
@@ -151,16 +151,16 @@ class PandasScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
_iter_based = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
import pandas as pd
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One SCALAR_PANDAS_ITER UDF expected here."
- self._udf_func, self._args_offsets, _, self._return_type = udfs[0]
+ udf = udfs[0]
+ self._udf_func, self._args_offsets = udf.func, udf.args_offsets
+ self._return_type = udf.return_type
self._return_schema = StructType([StructField("_0", self._return_type)])
self._expected_iter_type = (
Iterator[pd.DataFrame]
@@ -223,16 +223,14 @@ class PandasMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
_iter_based = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
import pandas as pd
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One MAP_PANDAS_ITER UDF expected here."
- self._udf_func, _, _, self._return_type = udfs[0]
+ self._udf_func, self._return_type = udfs[0].func, udfs[0].return_type
self._return_schema = StructType([StructField("_0", self._return_type)])
self._expected_iter_type = (
Iterator[pd.DataFrame]
@@ -291,15 +289,16 @@ class PandasCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One COGROUPED_MAP_PANDAS UDF expected here."
- self._cogrouped_udf, arg_offsets, self._return_type, self._num_udf_args = udfs[0]
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ udf = udfs[0]
+ self._cogrouped_udf, self._return_type = udf.func, udf.return_type
+ # signature was lost when wrapping it
+ self._num_udf_args = len(inspect.getfullargspec(udf.raw_func).args)
+ parsed_offsets = extract_key_value_indexes(udf.args_offsets)
self._left_key_offsets, self._left_value_offsets = parsed_offsets[0]
self._right_key_offsets, self._right_value_offsets = parsed_offsets[1]
self._return_schema = StructType([StructField("_0", self._return_type)])
@@ -373,15 +372,16 @@ class PandasGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
_iter_based = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One GROUPED_MAP_PANDAS_ITER UDF expected here."
- self._grouped_udf, arg_offsets, self._return_type, self._num_udf_args = udfs[0]
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ udf = udfs[0]
+ self._grouped_udf, self._return_type = udf.func, udf.return_type
+ # signature was lost when wrapping it
+ self._num_udf_args = len(inspect.getfullargspec(udf.raw_func).args)
+ parsed_offsets = extract_key_value_indexes(udf.args_offsets)
assert len(parsed_offsets) == 1, (
"Expected one pair of offsets for GROUPED_MAP_PANDAS_ITER UDF."
)
@@ -453,15 +453,16 @@ class PandasGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One GROUPED_MAP_PANDAS UDF expected here."
- self._grouped_udf, arg_offsets, self._return_type, self._num_udf_args = udfs[0]
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ udf = udfs[0]
+ self._grouped_udf, self._return_type = udf.func, udf.return_type
+ # signature was lost when wrapping it
+ self._num_udf_args = len(inspect.getfullargspec(udf.raw_func).args)
+ parsed_offsets = extract_key_value_indexes(udf.args_offsets)
assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_PANDAS UDF."
self._key_offsets, self._value_offsets = parsed_offsets[0]
self._return_schema = StructType([StructField("_0", self._return_type)])
@@ -535,14 +536,12 @@ class PandasGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
self._return_schema = StructType(
- [StructField("_%d" % i, rt) for i, (_, _, _, rt) in enumerate(udfs)]
+ [StructField("_%d" % i, udf.return_type) for i, udf in enumerate(udfs)]
)
def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]":
@@ -560,11 +559,11 @@ class PandasGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
)
results = [
- udf_func(
- *[all_series[o] for o in args_offsets],
- **{k: all_series[v] for k, v in kwargs_offsets.items()},
+ udf.func(
+ *[all_series[o] for o in udf.args_offsets],
+ **{k: all_series[v] for k, v in udf.kwargs_offsets.items()},
)
- for udf_func, args_offsets, kwargs_offsets, _ in self._udfs
+ for udf in self._udfs
]
yield PandasToArrowConversion.from_pandas(
[pd.Series([r]) for r in results],
@@ -586,15 +585,14 @@ class PandasGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One GROUPED_AGG_PANDAS_ITER UDF expected here."
- self._udf_func, self._args_offsets, _, return_type = udfs[0]
- self._return_schema = StructType([StructField("_0", return_type)])
+ udf = udfs[0]
+ self._udf_func, self._args_offsets = udf.func, udf.args_offsets
+ self._return_schema = StructType([StructField("_0", udf.return_type)])
def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]":
import pandas as pd
@@ -639,15 +637,13 @@ class PandasWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
self._window_bound_types = runner_conf.window_bound_types
self._return_schema = StructType(
- [StructField("_%d" % i, rt) for i, (_, _, _, rt) in enumerate(udfs)]
+ [StructField("_%d" % i, udf.return_type) for i, udf in enumerate(udfs)]
)
def run(self, split_index: int, data: "Iterator[GroupedBatch]") -> "Iterator[pa.RecordBatch]":
@@ -667,7 +663,9 @@ class PandasWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
num_rows = table.num_rows
result_series = []
- for udf_index, (udf_func, args_offsets, kwargs_offsets, _) in enumerate(self._udfs):
+ for udf_index, udf in enumerate(self._udfs):
+ udf_func = udf.func
+ args_offsets, kwargs_offsets = udf.args_offsets, udf.kwargs_offsets
bound_type = self._window_bound_types[udf_index]
if bound_type == "unbounded":
result = udf_func(
@@ -733,18 +731,17 @@ class ArrowBatchedUDFLegacyPandasHandler(BatchEvalTypeHandler["pa.RecordBatch"])
variant_conf = "use_legacy_pandas_udf_conversion"
variant_value = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: "RunnerConf", eval_conf: "EvalConf"
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
# --- UDF preparation ---
self._udf_infos = []
- for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in udfs:
+ for udf in udfs:
+ udf_return_type = udf.return_type
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
- udf_func, udf_args_offsets, udf_kwargs_offsets
+ udf.func, udf.args_offsets, udf.kwargs_offsets
)
zero_arg = len(args_kwargs_offsets) == 0
# Legacy coerces String/Binary for Arrow compatibility
diff --git a/python/pyspark/eval_handlers/_stateful.py b/python/pyspark/eval_handlers/_stateful.py
index 564b25f339fa..e8a1f49660ce 100644
--- a/python/pyspark/eval_handlers/_stateful.py
+++ b/python/pyspark/eval_handlers/_stateful.py
@@ -67,6 +67,7 @@ if TYPE_CHECKING:
import pandas as pd
import pyarrow as pa
+ from pyspark.eval_handlers._typing import UDF
from pyspark.worker_util import EvalConf, RunnerConf
@@ -97,9 +98,7 @@ class PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
]
)
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
@@ -108,8 +107,8 @@ class PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
# See FlatMapGroupsInPandasWithStateExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes.
- self._func, arg_offsets, self._return_type = udfs[0]
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ self._func, self._return_type = udfs[0].func, udfs[0].return_type
+ parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._value_offsets = parsed_offsets[0][1]
@@ -421,19 +420,17 @@ class PandasTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"])
eval_type = PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One TRANSFORM_WITH_STATE_PANDAS UDF expected here."
- self._func, arg_offsets, self._return_type = udfs[0]
+ self._func, self._return_type = udfs[0].func, udfs[0].return_type
# See TransformWithStateInPandasExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
assert len(parsed_offsets) == 1, (
"Expected one pair of offsets for TRANSFORM_WITH_STATE_PANDAS UDF."
)
@@ -576,15 +573,13 @@ class PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
eval_type = PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One TRANSFORM_WITH_STATE_PANDAS_INIT_STATE UDF expected here."
- self._func, arg_offsets, self._return_type = udfs[0]
+ self._func, self._return_type = udfs[0].func, udfs[0].return_type
# See TransformWithStateInPandasExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes.
@@ -593,7 +588,7 @@ class PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
# [groupingKeyOffsets, dedupDataOffsets],
# [initStateGroupingOffsets, dedupInitDataOffsets]
# ]
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._init_key_offsets = parsed_offsets[1][0]
self._output_schema = StructType([StructField("_0", self._return_type)])
@@ -798,18 +793,16 @@ class RowTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One TRANSFORM_WITH_STATE_PYTHON_ROW UDF expected here."
- self._func, arg_offsets, self._return_type = udfs[0]
+ self._func, self._return_type = udfs[0].func, udfs[0].return_type
# See TransformWithStateInPySparkExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes.
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._stateful_processor_api_client = StatefulProcessorApiClient(
@@ -907,14 +900,12 @@ class RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
eval_type = PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE UDF expected here."
- self._func, arg_offsets, self._return_type = udfs[0]
+ self._func, self._return_type = udfs[0].func, udfs[0].return_type
# See TransformWithStateInPandasExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes.
@@ -923,7 +914,7 @@ class RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
# [groupingKeyOffsets, dedupDataOffsets],
# [initStateGroupingOffsets, dedupInitDataOffsets]
# ]
- parsed_offsets = extract_key_value_indexes(arg_offsets)
+ parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._init_key_offsets = parsed_offsets[1][0]
diff --git a/python/pyspark/eval_handlers/_typing.py b/python/pyspark/eval_handlers/_typing.py
index 4198c64fa182..8d533bd1278f 100644
--- a/python/pyspark/eval_handlers/_typing.py
+++ b/python/pyspark/eval_handlers/_typing.py
@@ -21,12 +21,15 @@ The Arrow element types are forward refs so pyarrow stays a type-checking-only
import at runtime, as elsewhere in the worker.
"""
-from collections.abc import Iterator
-from typing import TYPE_CHECKING, TypeVar
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
if TYPE_CHECKING:
import pyarrow as pa
+ from pyspark.sql.types import StructType
+
# Grouped-category input element, matching what the serializer yields: a group
# serializer yields a lazy iterator; a co-group serializer materializes each side.
GroupedBatch = Iterator["pa.RecordBatch"]
@@ -35,3 +38,39 @@ CoGroupedBatch = tuple[list["pa.RecordBatch"], list["pa.RecordBatch"]]
# Handler input and output stream element types.
InputBatch = TypeVar("InputBatch")
OutputBatch = TypeVar("OutputBatch")
+
+
+@dataclass(frozen=True)
+class UDF:
+ """One UDF as read by ``read_single_udf``; each handler takes the fields it needs."""
+
+ # The chained function wrapped for profiling and, for most eval types, in
+ # fail_on_stopiteration.
+ func: Callable[..., Any]
+ # The chained function before any wrapping, for handlers that need the pickled object
+ # itself (an incremental Aggregator) or its signature.
+ raw_func: Any
+ args_offsets: list[int]
+ kwargs_offsets: dict[str, int]
+ # The DataType of the last function in the chain; a StructType for the eval types that
+ # return rows.
+ return_type: Any
+
+
+class UDTFInstance(Protocol):
+ """An instance of the user's UDTF class, or the PartitionedUDTF wrapping one: an object whose
+ methods are called, not a function. ``terminate`` and ``cleanup`` are optional."""
+
+ def eval(self, *args: Any, **kwargs: Any) -> Any: ...
+
+
+@dataclass(frozen=True)
+class UDTF:
+ """The UDTF as read by ``read_single_udtf``."""
+
+ # Wrapped in the handler's partition_wrapper when its TABLE argument has PARTITION BY
+ # expressions.
+ udtf: UDTFInstance
+ args_offsets: list[int]
+ kwargs_offsets: dict[str, int]
+ return_type: "StructType"
diff --git a/python/pyspark/eval_handlers/_udtf.py b/python/pyspark/eval_handlers/_udtf.py
index e470b2c69db4..c17b7e82653f 100644
--- a/python/pyspark/eval_handlers/_udtf.py
+++ b/python/pyspark/eval_handlers/_udtf.py
@@ -18,8 +18,8 @@
"""Handlers for the user-defined table function (UDTF) eval types, and the wrappers that
implement TABLE argument PARTITION BY semantics on top of a UDTF instance.
-Each handler receives a single ``(udtf, args_offsets, kwargs_offsets, return_type)`` tuple
-from ``read_single_udtf``, where ``udtf`` is the instantiated UDTF.
+Each handler receives the single ``UDTF`` record ``read_single_udtf`` returns, which holds
+the instantiated UDTF.
pyarrow is imported lazily (inside the handlers and the type-checking block) so the module
stays importable and its handlers register without pyarrow installed.
@@ -34,7 +34,7 @@ from typing import TYPE_CHECKING, Any, Callable, ClassVar, Optional, Tuple
from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
from pyspark.eval_handlers._base import BatchEvalTypeHandler, EvalTypeHandler
-from pyspark.eval_handlers._typing import InputBatch, OutputBatch
+from pyspark.eval_handlers._typing import UDTF, InputBatch, OutputBatch
from pyspark.eval_handlers.utils import wrap_kwargs_support
from pyspark.eval_handlers.verification import verify_arrow_result, verify_pandas_result
from pyspark.serializers import BatchedSerializer, CPickleSerializer, Serializer
@@ -603,14 +603,13 @@ class UDTFEvalTypeHandler(EvalTypeHandler[InputBatch, OutputBatch]):
# Wraps the UDTF when its TABLE argument has PARTITION BY expressions; see read_single_udtf.
partition_wrapper: ClassVar[type[PartitionedUDTF]] = UDTFWithPartitions
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One UDTF expected here."
- udtf, args_offsets, kwargs_offsets, self._return_type = udfs[0]
+ udtf = udfs[0].udtf
+ self._return_type = udfs[0].return_type
self._eval_method, self._args_kwargs_offsets = wrap_kwargs_support(
- getattr(udtf, "eval"), args_offsets, kwargs_offsets
+ udtf.eval, udfs[0].args_offsets, udfs[0].kwargs_offsets
)
self._terminate_method = getattr(udtf, "terminate", None)
self._cleanup_method = getattr(udtf, "cleanup", None)
@@ -648,16 +647,14 @@ class TableUDTFHandler(EvalTypeHandler[Any, tuple]):
# Wraps the UDTF when its TABLE argument has PARTITION BY expressions; see read_single_udtf.
partition_wrapper: ClassVar[type[PartitionedUDTF]] = UDTFWithPartitions
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
super().__init__(udfs, runner_conf, eval_conf)
assert len(udfs) == 1, "One UDTF expected here."
- udtf, args_offsets, kwargs_offsets, return_type = udfs[0]
+ udtf, return_type = udfs[0].udtf, udfs[0].return_type
self._check_output_row_against_schema = build_null_checker(return_type)
eval_func_kwargs_support, self._args_kwargs_offsets = wrap_kwargs_support(
- getattr(udtf, "eval"), args_offsets, kwargs_offsets
+ udtf.eval, udfs[0].args_offsets, udfs[0].kwargs_offsets
)
self._eval = self._wrap_udtf(eval_func_kwargs_support, return_type)
@@ -776,9 +773,7 @@ class ArrowUDTFHandler(
eval_type = PythonEvalType.SQL_ARROW_UDTF
partition_wrapper = ArrowUDTFWithPartition
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
import pyarrow as pa
super().__init__(udfs, runner_conf, eval_conf)
@@ -888,9 +883,7 @@ class ArrowTableUDTFHandler(
``spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled``. Subclasses implement
``_to_columns`` and ``_convert_result``."""
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
super().__init__(udfs, runner_conf, eval_conf)
self._return_type_size = len(self._return_type)
self._check_output_row_against_schema = build_null_checker(self._return_type)
@@ -942,9 +935,7 @@ class ArrowTableUDTFLegacyPandasHandler(ArrowTableUDTFHandler):
variant_conf = "use_legacy_pandas_udtf_conversion"
variant_value = True
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
import pandas # noqa: F401
super().__init__(udfs, runner_conf, eval_conf)
@@ -1054,9 +1045,7 @@ class ArrowTableUDTFArrowHandler(ArrowTableUDTFHandler):
variant_conf = "use_legacy_pandas_udtf_conversion"
variant_value = False
- def __init__(
- self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
- ) -> None:
+ def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf: EvalConf) -> None:
import pyarrow # noqa: F401
super().__init__(udfs, runner_conf, eval_conf)
diff --git a/python/pyspark/eval_handlers/_vanilla.py b/python/pyspark/eval_handlers/_vanilla.py
index 043ec2cb44f7..2ad0e4767fe5 100644
--- a/python/pyspark/eval_handlers/_vanilla.py
+++ b/python/pyspark/eval_handlers/_vanilla.py
@@ -21,33 +21,54 @@ Input rows arrive as pickled batches and are evaluated one row at a time.
"""
import os
-from collections.abc import Iterator
-from typing import Any
+from collections.abc import Callable, Iterator
+from typing import TYPE_CHECKING, Any
from pyspark.eval_handlers._base import EvalTypeHandler
+from pyspark.eval_handlers.utils import wrap_kwargs_support
from pyspark.serializers import BatchedSerializer, CPickleSerializer, Serializer
from pyspark.util import PythonEvalType
+if TYPE_CHECKING:
+ from pyspark.eval_handlers._typing import UDF
+ from pyspark.worker_util import EvalConf, RunnerConf
+
class BatchedUDFHandler(EvalTypeHandler[Any, Any]):
- """SQL_BATCHED_UDF. ``read_single_udf`` prepares each UDF as an ``(arg_offsets, eval_func)``
- pair. Apply every one to each input row: a single UDF yields its bare result, multiple UDFs
- yield a tuple of results, which is the shape the JVM side expects. The number of UDFs is
- fixed, so the single-result case is handled once here rather than by unwrapping a one-element
- tuple per row."""
+ """SQL_BATCHED_UDF. Prepare each UDF as an ``(arg_offsets, eval_func)`` pair and apply every
+ one to each input row: a single UDF yields its bare result, multiple UDFs yield a tuple of
+ results, which is the shape the JVM side expects. The number of UDFs is fixed, so the
+ single-result case is handled once here rather than by unwrapping a one-element tuple per
+ row."""
eval_type = PythonEvalType.SQL_BATCHED_UDF
+ def __init__(self, udfs: "list[UDF]", runner_conf: "RunnerConf", eval_conf: "EvalConf") -> None:
+ super().__init__(udfs, runner_conf, eval_conf)
+ self._udf_infos = [self._prepare_udf(udf) for udf in udfs]
+
+ @staticmethod
+ def _prepare_udf(udf: "UDF") -> tuple[list[int], Callable[..., Any]]:
+ """Apply kwargs binding, and convert each result to the internal representation only
+ when the return type requires it."""
+ func, args_kwargs_offsets = wrap_kwargs_support(
+ udf.func, udf.args_offsets, udf.kwargs_offsets
+ )
+ if udf.return_type.needConversion():
+ toInternal = udf.return_type.toInternal
+ return args_kwargs_offsets, lambda *a: toInternal(func(*a))
+ return args_kwargs_offsets, func
+
@property
def serializer(self) -> Serializer:
batch_size = int(os.environ.get("PYTHON_UDF_BATCH_SIZE", "100"))
return BatchedSerializer(CPickleSerializer(), batch_size)
def run(self, split_index: int, data: Iterator[Any]) -> Iterator[Any]:
- if len(self._udfs) == 1:
- arg_offsets, f = self._udfs[0]
+ if len(self._udf_infos) == 1:
+ arg_offsets, f = self._udf_infos[0]
return (f(*[row[offset] for offset in arg_offsets]) for row in data)
return (
- tuple(f(*[row[offset] for offset in arg_offsets]) for arg_offsets, f in self._udfs)
+ tuple(f(*[row[offset] for offset in arg_offsets]) for arg_offsets, f in self._udf_infos)
for row in data
)
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 b62283a20c0f..63ce1c546314 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
@@ -29,11 +29,12 @@ from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._typing import UDF
from pyspark.eval_handlers.utils import hashable_grouping_key
from pyspark.sql.pandas.serializers import ArrowStreamCoGroupSerializer, ArrowStreamSerializer
from pyspark.sql.types import DoubleType, LongType, StringType, StructField, StructType
from pyspark.testing.utils import PySparkBaseTestCase, have_pyarrow, pyarrow_requirement_message
-from pyspark.util import PythonEvalType
+from pyspark.util import PythonEvalType, fail_on_stopiteration
from pyspark.worker_metrics import WorkerMetrics
with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
@@ -124,23 +125,27 @@ class ArrowEvalTypeHandlerTestCase(PySparkBaseTestCase):
return offsets
@staticmethod
- 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 _udf(func, args_offsets, kwargs_offsets=None, return_type=None, raw_func=None):
+ """A ``UDF`` record; ``raw_func`` defaults to ``func``."""
+ return UDF(
+ func=func,
+ raw_func=func if raw_func is None else raw_func,
+ args_offsets=args_offsets,
+ kwargs_offsets=kwargs_offsets or {},
+ return_type=return_type or LongType(),
)
@classmethod
- def _grouped_handler(cls, 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)``.
- """
+ def _grouped_handler(cls, handler_cls, udf, arg_offsets):
+ """Build a grouped/cogrouped-map handler with the shared return type. As from
+ ``read_single_udf``, ``func`` is wrapped (losing the signature) while ``raw_func``
+ keeps it."""
return handler_cls(
- udfs=[(udf, arg_offsets, cls._RETURN_TYPE, num_udf_args)],
+ udfs=[
+ cls._udf(
+ lambda *a: udf(*a), arg_offsets, return_type=cls._RETURN_TYPE, raw_func=udf
+ )
+ ],
runner_conf=RunnerConf({}),
eval_conf=None,
)
@@ -194,7 +199,7 @@ class ArrowBatchedUDFArrowHandlerTests(ArrowEvalTypeHandlerTestCase):
def _handler(self, *udfs, runner_conf=None, input_type=None):
"""Build the handler from ``(func, args_offsets, kwargs_offsets, return_type)`` tuples."""
return ArrowBatchedUDFArrowHandler(
- udfs=list(udfs),
+ udfs=[self._udf(*u) for u in udfs],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=EvalConf({"input_type": (input_type or self._INPUT_TYPE).json()}),
)
@@ -381,7 +386,9 @@ class ArrowScalarUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
def add_one(col):
return pa.array([v.as_py() + 1 for v in col], type=pa.int64())
- handler = self._scalar_handler(ArrowScalarUDFHandler, add_one)
+ handler = ArrowScalarUDFHandler(
+ udfs=[self._udf(add_one, [0])], runner_conf=RunnerConf({}), eval_conf=None
+ )
out = list(handler.run(0, iter([self._batch(a=[1, 2, 3])])))
self.assertEqual([b.column(0).to_pylist() for b in out], [[2, 3, 4]])
@@ -390,7 +397,9 @@ class ArrowScalarUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
def add_one(col):
return pa.array([v.as_py() + 1 for v in col], type=pa.int32())
- handler = self._scalar_handler(ArrowScalarUDFHandler, add_one)
+ handler = ArrowScalarUDFHandler(
+ udfs=[self._udf(add_one, [0])], runner_conf=RunnerConf({}), eval_conf=None
+ )
out = list(handler.run(0, iter([self._batch(a=[10, 20])])))
self.assertEqual(out[0].schema.field(0).type, pa.int64())
self.assertEqual(out[0].column(0).to_pylist(), [11, 21])
@@ -403,7 +412,9 @@ class ArrowScalarIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
for col in col_iter:
yield pa.array([v.as_py() + 1 for v in col], type=pa.int64())
- handler = self._scalar_handler(ArrowScalarIterUDFHandler, add_one)
+ handler = ArrowScalarIterUDFHandler(
+ udfs=[self._udf(add_one, [0])], runner_conf=RunnerConf({}), eval_conf=None
+ )
out = list(handler.run(0, iter([self._batch(a=[1, 2]), self._batch(a=[3])])))
self.assertEqual([b.column(0).to_pylist() for b in out], [[2, 3], [4]])
@@ -413,7 +424,9 @@ class ArrowScalarIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
for col in col_iter:
yield pa.array(list(range(len(col) + 1)), type=pa.int64())
- handler = self._scalar_handler(ArrowScalarIterUDFHandler, too_many)
+ handler = ArrowScalarIterUDFHandler(
+ udfs=[self._udf(too_many, [0])], runner_conf=RunnerConf({}), eval_conf=None
+ )
with self.assertRaises(PySparkRuntimeError):
list(handler.run(0, iter([self._batch(a=[1, 2])])))
@@ -427,7 +440,7 @@ class ArrowMapUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
yield self._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
+ udfs=[self._udf(double_v, [])], runner_conf=RunnerConf({}), eval_conf=None
)
out = list(handler.run(0, iter([self._struct_batch(v=[1, 2, 3])])))
self.assertEqual(out[0].column(0).field("v").to_pylist(), [2, 4, 6])
@@ -442,9 +455,7 @@ class ArrowGroupedMapUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
def grouped_udf(value_table):
return pa.table({"v": pa.array([c.as_py() * 10 for c in value_table.column("v")])})
- handler = self._grouped_handler(
- ArrowGroupedMapUDFHandler, grouped_udf, self._GROUP_OFFSETS, 1
- )
+ handler = self._grouped_handler(ArrowGroupedMapUDFHandler, grouped_udf, self._GROUP_OFFSETS)
out = list(handler.run(0, self._one_group(self._struct_batch(k=[7, 7], v=[1, 2]))))
self.assertEqual(out[0].column(0).field("v").to_pylist(), [10, 20])
@@ -453,9 +464,7 @@ class ArrowGroupedMapUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
k = key[0].as_py()
return pa.table({"v": pa.array([c.as_py() + k for c in value_table.column("v")])})
- handler = self._grouped_handler(
- ArrowGroupedMapUDFHandler, grouped_udf, self._GROUP_OFFSETS, 2
- )
+ handler = self._grouped_handler(ArrowGroupedMapUDFHandler, grouped_udf, self._GROUP_OFFSETS)
out = list(handler.run(0, self._one_group(self._struct_batch(k=[7, 7], v=[1, 2]))))
self.assertEqual(out[0].column(0).field("v").to_pylist(), [8, 9])
@@ -471,7 +480,7 @@ class ArrowGroupedMapIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
yield self._batch(v=[c.as_py() + 1 for c in batch.column("v")])
handler = self._grouped_handler(
- ArrowGroupedMapIterUDFHandler, grouped_udf, self._GROUP_OFFSETS, 1
+ ArrowGroupedMapIterUDFHandler, grouped_udf, self._GROUP_OFFSETS
)
out = list(handler.run(0, self._one_group(self._struct_batch(k=[7], v=[41]))))
self.assertEqual(out[0].column(0).field("v").to_pylist(), [42])
@@ -483,7 +492,7 @@ class ArrowGroupedMapIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
yield self._batch(v=[c.as_py() + k for c in batch.column("v")])
handler = self._grouped_handler(
- ArrowGroupedMapIterUDFHandler, grouped_udf, self._GROUP_OFFSETS, 2
+ ArrowGroupedMapIterUDFHandler, grouped_udf, self._GROUP_OFFSETS
)
out = list(handler.run(0, self._one_group(self._struct_batch(k=[10], v=[5]))))
self.assertEqual(out[0].column(0).field("v").to_pylist(), [15])
@@ -506,7 +515,7 @@ class ArrowCoGroupedMapUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return pa.table({"v": pa.array([total], type=pa.int64())})
handler = self._grouped_handler(
- ArrowCoGroupedMapUDFHandler, cogrouped_udf, self._COGROUP_OFFSETS, 2
+ ArrowCoGroupedMapUDFHandler, cogrouped_udf, self._COGROUP_OFFSETS
)
out = list(
handler.run(
@@ -522,7 +531,7 @@ class ArrowCoGroupedMapUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return pa.table({"v": pa.array([total + k], type=pa.int64())})
handler = self._grouped_handler(
- ArrowCoGroupedMapUDFHandler, cogrouped_udf, self._COGROUP_OFFSETS, 3
+ ArrowCoGroupedMapUDFHandler, cogrouped_udf, self._COGROUP_OFFSETS
)
out = list(
handler.run(
@@ -541,7 +550,7 @@ class ArrowGroupedAggUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return sum(c.as_py() for c in col)
handler = ArrowGroupedAggUDFHandler(
- udfs=[(sum_udf, [0], {}, LongType())], runner_conf=RunnerConf({}), eval_conf=None
+ udfs=[self._udf(sum_udf, [0])], runner_conf=RunnerConf({}), eval_conf=None
)
out = list(handler.run(0, self._one_group(self._batch(v=[1, 2, 3]), self._batch(v=[4]))))
self.assertEqual(out[0].column("_0").to_pylist(), [10])
@@ -555,7 +564,7 @@ class ArrowGroupedAggIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
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
+ udfs=[self._udf(sum_iter_udf, [0])], runner_conf=RunnerConf({}), eval_conf=None
)
out = list(handler.run(0, self._one_group(self._batch(v=[10]), self._batch(v=[20]))))
self.assertEqual(out[0].column("_0").to_pylist(), [30])
@@ -569,7 +578,7 @@ class ArrowWindowAggUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return sum(c.as_py() for c in col)
handler = ArrowWindowAggUDFHandler(
- udfs=[(sum_udf, [0], {}, LongType())],
+ udfs=[self._udf(sum_udf, [0])],
runner_conf=RunnerConf({"window_bound_types": "unbounded"}),
eval_conf=None,
)
@@ -582,7 +591,7 @@ class ArrowWindowAggUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return sum(c.as_py() for c in col)
handler = ArrowWindowAggUDFHandler(
- udfs=[(sum_udf, [0, 1, 2], {}, LongType())],
+ udfs=[self._udf(sum_udf, [0, 1, 2])],
runner_conf=RunnerConf({"window_bound_types": "bounded"}),
eval_conf=None,
)
@@ -600,8 +609,10 @@ class ArrowGroupedAggIncrementalPartialUDFHandlerTests(ArrowEvalTypeHandlerTestC
_KEY_SCHEMA = StructType([StructField("k", LongType())])
def _handler(self, cap=10000):
+ # As from read_single_udf: func is wrapped, raw_func is the Aggregator itself.
+ agg = _SumAggregator()
return ArrowGroupedAggIncrementalPartialUDFHandler(
- udfs=[(_SumAggregator(), [1], {}, LongType())],
+ udfs=[self._udf(fail_on_stopiteration(agg), [1], raw_func=agg)],
runner_conf=RunnerConf({"spark.sql.execution.arrow.maxRecordsPerBatch": str(cap)}),
eval_conf=EvalConf({"grouping_key_schema": self._KEY_SCHEMA.json()}),
)
@@ -658,8 +669,9 @@ class ArrowGroupedAggIncrementalPartialUDFHandlerTests(ArrowEvalTypeHandlerTestC
[pa.array([float("nan"), float("nan")]), pa.array([1, 2], type=pa.int64())],
["k", "v"],
)
+ agg = _SumAggregator()
handler = ArrowGroupedAggIncrementalPartialUDFHandler(
- udfs=[(_SumAggregator(), [1], {}, LongType())],
+ udfs=[self._udf(fail_on_stopiteration(agg), [1], raw_func=agg)],
runner_conf=RunnerConf({}),
eval_conf=EvalConf(
{"grouping_key_schema": StructType([StructField("k", DoubleType())]).json()}
@@ -680,8 +692,10 @@ class ArrowGroupedAggIncrementalFinalUDFHandlerTests(ArrowEvalTypeHandlerTestCas
return pa.RecordBatch.from_arrays([pa.array(buffers, type=self._BUFFER_TYPE)], ["buf"])
def _handler(self):
+ # As from read_single_udf: func is wrapped, raw_func is the Aggregator itself.
+ agg = _SumAggregator()
return ArrowGroupedAggIncrementalFinalUDFHandler(
- udfs=[(_SumAggregator(), [0], {}, LongType())],
+ udfs=[self._udf(fail_on_stopiteration(agg), [0], raw_func=agg)],
runner_conf=RunnerConf({}),
eval_conf=None,
)
@@ -708,7 +722,8 @@ class ArrowGroupedAggIncrementalFinalUDFHandlerTests(ArrowEvalTypeHandlerTestCas
class ArrowWindowAggIncrementalUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
def _handler(self, agg, bound_type, args_offsets):
return ArrowWindowAggIncrementalUDFHandler(
- udfs=[(agg, args_offsets, {}, LongType())],
+ # As from read_single_udf: func is wrapped, raw_func is the Aggregator itself.
+ udfs=[self._udf(fail_on_stopiteration(agg), args_offsets, raw_func=agg)],
runner_conf=RunnerConf({"window_bound_types": bound_type}),
eval_conf=None,
)
diff --git a/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py
index 26f414a5c1e4..fdd0d76ebaf0 100644
--- a/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_base_eval_type_handlers.py
@@ -22,8 +22,10 @@ Handler-flavor tests live alongside their module, e.g. the Arrow handlers in
"""
import os
+from types import SimpleNamespace
from unittest.mock import patch
+from pyspark.errors import PySparkRuntimeError
from pyspark.eval_handlers._base import (
BatchEvalTypeHandler,
CoGroupedEvalTypeHandler,
@@ -32,15 +34,21 @@ from pyspark.eval_handlers._base import (
_eval_type_handlers,
get_eval_type_handler,
)
+from pyspark.eval_handlers._typing import UDF
+from pyspark.serializers import CloudPickleSerializer
from pyspark.sql.pandas.serializers import (
ArrowStreamCoGroupSerializer,
ArrowStreamGroupSerializer,
ArrowStreamSerializer,
)
+from pyspark.sql.types import LongType
from pyspark.testing.utils import PySparkBaseTestCase
+from pyspark.util import PythonEvalType
with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
- from pyspark.worker_util import RunnerConf
+ from pyspark.worker import read_single_udf
+ from pyspark.worker_message import UDFInfo
+ from pyspark.worker_util import RunnerConf, pickleSer
class EvalTypeHandlerTests(PySparkBaseTestCase):
@@ -264,6 +272,32 @@ class EvalTypeHandlerTests(PySparkBaseTestCase):
self.assertNotIn(unused_eval_type, _eval_type_handlers)
+class ReadSingleUDFTests(PySparkBaseTestCase):
+ def test_returns_udf_record(self):
+ def stop(x):
+ raise StopIteration(x)
+
+ udf_info = UDFInfo(
+ udfs=[CloudPickleSerializer().dumps((stop, LongType()))],
+ args=[1],
+ kwargs={"y": 2},
+ result_id=0,
+ )
+ # read_single_udf only reads these two attributes of the handler class.
+ handler_cls = SimpleNamespace(eval_type=PythonEvalType.SQL_BATCHED_UDF, is_iter_based=False)
+ udf = read_single_udf(pickleSer, udf_info, handler_cls, RunnerConf({}))
+
+ self.assertIsInstance(udf, UDF)
+ self.assertEqual(udf.args_offsets, [1])
+ self.assertEqual(udf.kwargs_offsets, {"y": 2})
+ self.assertEqual(udf.return_type, LongType())
+ # func is wrapped in fail_on_stopiteration; raw_func is the unwrapped function.
+ with self.assertRaises(PySparkRuntimeError):
+ udf.func(1)
+ with self.assertRaises(StopIteration):
+ udf.raw_func(1)
+
+
if __name__ == "__main__":
from pyspark.testing import main
diff --git a/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py
index 3fffbafcb4d4..a03a4726a878 100644
--- a/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_elementwise_eval_type_handlers.py
@@ -29,6 +29,7 @@ from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._typing import UDF
from pyspark.sql.pandas.serializers import ArrowStreamSerializer
from pyspark.sql.types import ArrayType, LongType, StringType, StructField, StructType
from pyspark.testing.utils import (
@@ -99,7 +100,16 @@ class ElementwiseEvalTypeHandlerTestCase(PySparkBaseTestCase):
if nesting is not None:
eval_conf["elementwise_nesting"] = ",".join(map(str, nesting))
return handler_cls(
- udfs=list(udfs),
+ udfs=[
+ UDF(
+ func=func,
+ raw_func=func,
+ args_offsets=args_offsets,
+ kwargs_offsets=kwargs_offsets,
+ return_type=return_type,
+ )
+ for func, args_offsets, kwargs_offsets, return_type in udfs
+ ],
runner_conf=RunnerConf(runner_conf or {"spark.sql.session.timeZone": "UTC"}),
eval_conf=EvalConf(eval_conf),
)
@@ -217,11 +227,8 @@ class ScalarElementwiseUDFHandlerTestsMixin:
handler_cls: type
- def _scalar_handler(self, *udfs, **kwargs):
- return self._handler(self.handler_cls, *udfs, **kwargs)
-
def test_serializer(self):
- handler = self._scalar_handler((lambda x: x, [0], {}, LongType()))
+ handler = self._handler(self.handler_cls, (lambda x: x, [0], {}, LongType()))
self.assertIsInstance(handler.serializer, ArrowStreamSerializer)
def test_invokes_udf_per_batch(self):
@@ -231,21 +238,22 @@ class ScalarElementwiseUDFHandlerTestsMixin:
calls.append(len(x))
return self._add(x, 1)
- handler = self._scalar_handler((udf, [0], {}, LongType()))
+ handler = self._handler(self.handler_cls, (udf, [0], {}, LongType()))
out = list(handler.run(0, iter([self._batch()])))
self.assertEqual(calls, [3])
self.assertEqual(out[0].schema.names, ["_0"])
self.assertEqual(self._column(out), self._A_PLUS_ONE)
def test_nesting_greater_than_one(self):
- handler = self._scalar_handler(
- (lambda x: self._add(x, 1), [1], {}, LongType()), nesting=[2]
+ handler = self._handler(
+ self.handler_cls, (lambda x: self._add(x, 1), [1], {}, LongType()), nesting=[2]
)
out = list(handler.run(0, iter([self._batch()])))
self.assertEqual(self._column(out), self._B_PLUS_ONE)
def test_fused_udfs_renest_by_own_shape(self):
- handler = self._scalar_handler(
+ handler = self._handler(
+ self.handler_cls,
(lambda x: self._add(x, 1), [0], {}, LongType()),
(lambda x: self._add(x, 1), [1], {}, LongType()),
(lambda x, y: self._mul(x, y), [0, 2], {}, LongType()),
@@ -258,31 +266,35 @@ class ScalarElementwiseUDFHandlerTestsMixin:
self.assertEqual(self._column(out, 2), self._A_TIMES_C * 2)
def test_passes_arg_by_keyword_offset(self):
- handler = self._scalar_handler((lambda x, y: self._mul(x, y), [0], {"y": 2}, LongType()))
+ handler = self._handler(
+ self.handler_cls, (lambda x, y: self._mul(x, y), [0], {"y": 2}, LongType())
+ )
out = list(handler.run(0, iter([self._batch()])))
self.assertEqual(self._column(out), self._A_TIMES_C)
def test_preserves_large_list(self):
- handler = self._scalar_handler(
- (lambda x: self._add(x, 1), [1], {}, LongType()), nesting=[2]
+ handler = self._handler(
+ self.handler_cls, (lambda x: self._add(x, 1), [1], {}, LongType()), nesting=[2]
)
out = list(handler.run(0, iter([self._batch(large_list=True)])))
self.assertEqual(out[0].schema.field(0).type, pa.large_list(pa.large_list(pa.int64())))
self.assertEqual(self._column(out), self._B_PLUS_ONE)
def test_all_empty_and_null_arrays(self):
- handler = self._scalar_handler((lambda x: self._add(x, 1), [0], {}, LongType()))
+ handler = self._handler(self.handler_cls, (lambda x: self._add(x, 1), [0], {}, LongType()))
batch = self._batch(a=[[], None], b=[None, None], c=[[], None])
out = list(handler.run(0, iter([batch])))
self.assertEqual(self._column(out), [[], None])
def test_rejects_result_length_mismatch(self):
- handler = self._scalar_handler((lambda x, y: self._mul(x, y)[:1], [0, 2], {}, LongType()))
+ handler = self._handler(
+ self.handler_cls, (lambda x, y: self._mul(x, y)[:1], [0, 2], {}, LongType())
+ )
with self.assertRaisesRegex(PySparkRuntimeError, "RESULT_ROWS_MISMATCH"):
list(handler.run(0, iter([self._batch()])))
def test_rejects_non_array_like_result(self):
- handler = self._scalar_handler((lambda x: 1, [0], {}, LongType()))
+ handler = self._handler(self.handler_cls, (lambda x: 1, [0], {}, LongType()))
with self.assertRaisesRegex(PySparkTypeError, "UDF_RETURN_TYPE"):
list(handler.run(0, iter([self._batch()])))
@@ -310,13 +322,15 @@ class ArrowScalarElementwiseUDFHandlerTests(
seen.append(x)
return x
- handler = self._scalar_handler((udf, [0], {}, LongType()))
+ handler = self._handler(self.handler_cls, (udf, [0], {}, LongType()))
list(handler.run(0, iter([self._batch()])))
self.assertIsInstance(seen[0], pa.Array)
self.assertEqual(seen[0].to_pylist(), [1, 2, 3])
def test_coerces_result_to_return_type(self):
- handler = self._scalar_handler((lambda x: pc.cast(x, pa.int32()), [0], {}, LongType()))
+ handler = self._handler(
+ self.handler_cls, (lambda x: pc.cast(x, pa.int32()), [0], {}, LongType())
+ )
out = list(handler.run(0, iter([self._batch()])))
self.assertEqual(out[0].schema.field(0).type, pa.list_(pa.int64()))
self.assertEqual(self._column(out), self._A)
@@ -345,7 +359,7 @@ class PandasScalarElementwiseUDFHandlerTests(
seen.append(x)
return x
- handler = self._scalar_handler((udf, [0], {}, LongType()))
+ handler = self._handler(self.handler_cls, (udf, [0], {}, LongType()))
list(handler.run(0, iter([self._batch()])))
self.assertIsInstance(seen[0], pd.Series)
self.assertIsNone(seen[0].name)
@@ -358,20 +372,22 @@ class PandasScalarElementwiseUDFHandlerTests(
[pa.array([[{"v": 1}, {"v": 2}], None], type=pa.list_(pa.struct([("v", pa.int64())])))],
["s"],
)
- handler = self._scalar_handler(
- (lambda df: df["v"] + 1, [0], {}, LongType()), input_type=input_type
+ handler = self._handler(
+ self.handler_cls, (lambda df: df["v"] + 1, [0], {}, LongType()), input_type=input_type
)
self.assertEqual(self._column(list(handler.run(0, iter([batch])))), [[2, 3], None])
def test_struct_return_type_requires_dataframe(self):
return_type = StructType([StructField("v", LongType())])
- handler = self._scalar_handler((lambda x: x, [0], {}, return_type))
+ handler = self._handler(self.handler_cls, (lambda x: x, [0], {}, return_type))
with self.assertRaises(PySparkValueError):
list(handler.run(0, iter([self._batch()])))
def test_struct_return_type(self):
return_type = StructType([StructField("v", LongType())])
- handler = self._scalar_handler((lambda x: pd.DataFrame({"v": x + 1}), [0], {}, return_type))
+ handler = self._handler(
+ self.handler_cls, (lambda x: pd.DataFrame({"v": x + 1}), [0], {}, return_type)
+ )
out = list(handler.run(0, iter([self._batch()])))
self.assertEqual(self._column(out), [[{"v": 2}, {"v": 3}], [], None, [{"v": 4}]])
diff --git a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
index 55e416239c5d..9b22dcae9d82 100644
--- a/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_pandas_eval_type_handlers.py
@@ -31,6 +31,7 @@ from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._typing import UDF
from pyspark.sql.types import BinaryType, LongType, StringType, StructField, StructType
from pyspark.testing.utils import (
PySparkBaseTestCase,
@@ -90,11 +91,19 @@ class PandasEvalTypeHandlerTestCase(PySparkBaseTestCase):
return offsets
@classmethod
- def _grouped_handler(cls, handler_cls, udf, arg_offsets, num_udf_args, runner_conf=None):
- """Build a grouped/cogrouped-map handler from its
- ``(func, arg_offsets, return_type, num_udf_args)`` UDF tuple."""
+ def _grouped_handler(cls, handler_cls, udf, arg_offsets, runner_conf=None):
+ """Build a grouped/cogrouped-map handler. As from ``read_single_udf``, ``func`` is
+ wrapped (losing the signature) while ``raw_func`` keeps it."""
return handler_cls(
- udfs=[(udf, arg_offsets, cls._GROUPED_RETURN_TYPE, num_udf_args)],
+ udfs=[
+ UDF(
+ func=lambda *a: udf(*a),
+ raw_func=udf,
+ args_offsets=arg_offsets,
+ kwargs_offsets={},
+ return_type=cls._GROUPED_RETURN_TYPE,
+ )
+ ],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=None,
)
@@ -106,12 +115,18 @@ class PandasEvalTypeHandlerTestCase(PySparkBaseTestCase):
@staticmethod
def _udf(func, return_type=None, args=(0,), kwargs=None):
- """A pandas UDF tuple ``(func, args_offsets, kwargs_offsets, return_type)``."""
- return (func, list(args), dict(kwargs or {}), return_type or LongType())
+ """A ``UDF`` record for a pandas UDF."""
+ return UDF(
+ func=func,
+ raw_func=func,
+ args_offsets=list(args),
+ kwargs_offsets=dict(kwargs or {}),
+ return_type=return_type or LongType(),
+ )
@staticmethod
def _agg_handler(handler_cls, *udfs, runner_conf=None):
- """Build a grouped-agg/window-agg handler from one or more ``_udf`` tuples."""
+ """Build a grouped-agg/window-agg handler from one or more ``_udf`` records."""
return handler_cls(
udfs=list(udfs), runner_conf=runner_conf or RunnerConf({}), eval_conf=None
)
@@ -158,7 +173,7 @@ class ArrowBatchedUDFLegacyPandasHandlerTests(PandasEvalTypeHandlerTestCase):
_INPUT_TYPE = StructType([StructField("a", LongType()), StructField("b", LongType())])
def _handler(self, *udfs, runner_conf=None, input_type=None):
- """Build the handler from one or more ``_udf`` tuples."""
+ """Build the handler from one or more ``_udf`` records."""
return ArrowBatchedUDFLegacyPandasHandler(
udfs=list(udfs),
runner_conf=runner_conf or RunnerConf({}),
@@ -218,7 +233,7 @@ class ArrowBatchedUDFLegacyPandasHandlerTests(PandasEvalTypeHandlerTestCase):
class PandasScalarUDFHandlerTests(PandasEvalTypeHandlerTestCase):
@staticmethod
def _handler(*udfs, runner_conf=None):
- """Build a PandasScalarUDFHandler from one or more ``_udf`` tuples."""
+ """Build a PandasScalarUDFHandler from one or more ``_udf`` records."""
return PandasScalarUDFHandler(
udfs=list(udfs), runner_conf=runner_conf or RunnerConf({}), eval_conf=None
)
@@ -305,7 +320,7 @@ class PandasScalarIterUDFHandlerTests(PandasEvalTypeHandlerTestCase):
def _handler(udf, return_type=None, args=(0,), runner_conf=None):
"""Build a PandasScalarIterUDFHandler whose one UDF reads the given arg offsets."""
return PandasScalarIterUDFHandler(
- udfs=[(udf, list(args), {}, return_type or LongType())],
+ udfs=[PandasEvalTypeHandlerTestCase._udf(udf, return_type, args)],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=None,
)
@@ -385,9 +400,9 @@ class PandasMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
@staticmethod
def _handler(udf, return_type, runner_conf=None):
- """Build a PandasMapUDFHandler from its ``(func, None, None, return_type)`` UDF tuple."""
+ """Build a PandasMapUDFHandler; map UDFs take no arg offsets."""
return PandasMapUDFHandler(
- udfs=[(udf, None, None, return_type)],
+ udfs=[PandasEvalTypeHandlerTestCase._udf(udf, return_type, args=())],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=None,
)
@@ -454,7 +469,6 @@ class PandasGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
PandasGroupedMapUDFHandler,
lambda df: pd.DataFrame({"v": df["b"] * 2}),
self._KEY0_VALUE1,
- 1,
)
groups = [
iter([self._batch(a=[1, 1], b=[1, 2]), self._batch(a=[1], b=[3])]),
@@ -471,7 +485,6 @@ class PandasGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
PandasGroupedMapUDFHandler,
lambda key, df: pd.DataFrame({"v": [key[0]] * len(df)}),
self._KEY0_VALUE1,
- 2,
)
out = list(handler.run(0, iter([iter([self._batch(a=[7, 7], b=[1, 2])])])))
self.assertEqual(self._struct_values(out), [7, 7])
@@ -481,14 +494,13 @@ class PandasGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
PandasGroupedMapUDFHandler,
lambda df: pd.DataFrame({"v": df["b"].astype("int32")}),
self._KEY0_VALUE1,
- 1,
)
out = list(handler.run(0, iter([iter([self._batch(a=[1], b=[3])])])))
self.assertEqual(out[0].column(0).type.field("v").type, pa.int64())
def test_rejects_non_dataframe_result(self):
handler = self._grouped_handler(
- PandasGroupedMapUDFHandler, lambda df: df["b"], self._KEY0_VALUE1, 1
+ PandasGroupedMapUDFHandler, lambda df: df["b"], self._KEY0_VALUE1
)
with self.assertRaises(PySparkTypeError):
list(handler.run(0, iter([iter([self._batch(a=[1], b=[3])])])))
@@ -500,7 +512,6 @@ class PandasGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
PandasGroupedMapUDFHandler,
lambda df: pd.DataFrame({"v": df["b"]}),
self._KEY0_VALUE1,
- 1,
runner_conf=conf,
)
values = list(range(100))
@@ -520,9 +531,7 @@ class PandasGroupedMapIterUDFHandlerTests(PandasEvalTypeHandlerTestCase):
for df in df_iter:
yield pd.DataFrame({"v": df["b"] * 2})
- handler = self._grouped_handler(
- PandasGroupedMapIterUDFHandler, double, self._KEY0_VALUE1, 1
- )
+ handler = self._grouped_handler(PandasGroupedMapIterUDFHandler, double, self._KEY0_VALUE1)
groups = [
iter([self._batch(a=[1, 1], b=[1, 2]), self._batch(a=[1], b=[3])]),
iter([self._batch(a=[2], b=[5])]),
@@ -535,9 +544,7 @@ class PandasGroupedMapIterUDFHandlerTests(PandasEvalTypeHandlerTestCase):
for df in df_iter:
yield pd.DataFrame({"v": [key[0]] * len(df)})
- handler = self._grouped_handler(
- PandasGroupedMapIterUDFHandler, with_key, self._KEY0_VALUE1, 2
- )
+ handler = self._grouped_handler(PandasGroupedMapIterUDFHandler, with_key, self._KEY0_VALUE1)
out = list(handler.run(0, iter([iter([self._batch(a=[7, 7], b=[1, 2])])])))
self.assertEqual(self._struct_values(out), [7, 7])
@@ -548,7 +555,7 @@ class PandasGroupedMapIterUDFHandlerTests(PandasEvalTypeHandlerTestCase):
yield pd.DataFrame({"v": next(df_iter)["b"]})
handler = self._grouped_handler(
- PandasGroupedMapIterUDFHandler, first_only, self._KEY0_VALUE1, 1
+ PandasGroupedMapIterUDFHandler, first_only, self._KEY0_VALUE1
)
group1 = iter([self._batch(a=[1], b=[1]), self._batch(a=[1], b=[2])])
group2 = iter([self._batch(a=[2], b=[3])])
@@ -562,7 +569,7 @@ class PandasGroupedMapIterUDFHandlerTests(PandasEvalTypeHandlerTestCase):
yield df["b"]
handler = self._grouped_handler(
- PandasGroupedMapIterUDFHandler, yields_series, self._KEY0_VALUE1, 1
+ PandasGroupedMapIterUDFHandler, yields_series, self._KEY0_VALUE1
)
with self.assertRaises(PySparkTypeError):
list(handler.run(0, iter([iter([self._batch(a=[1], b=[3])])])))
@@ -578,7 +585,6 @@ class PandasCoGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
PandasCoGroupedMapUDFHandler,
lambda left, right: pd.DataFrame({"v": [left["b"].sum() + right["b"].sum()]}),
self._offsets,
- 2,
)
cogroup = (
[self._batch(a=[1], b=[1]), self._batch(a=[1], b=[2])],
@@ -594,7 +600,6 @@ class PandasCoGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
PandasCoGroupedMapUDFHandler,
lambda key, left, right: pd.DataFrame({"v": [key[0]]}),
self._offsets,
- 3,
)
cogroups = [
([self._batch(a=[4], b=[1])], [self._batch(a=[], b=[])]),
@@ -605,7 +610,7 @@ class PandasCoGroupedMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
def test_rejects_non_dataframe_result(self):
handler = self._grouped_handler(
- PandasCoGroupedMapUDFHandler, lambda left, right: 1, self._offsets, 2
+ PandasCoGroupedMapUDFHandler, lambda left, right: 1, self._offsets
)
cogroup = ([self._batch(a=[1], b=[1])], [self._batch(a=[1], b=[2])])
with self.assertRaises(PySparkTypeError):
diff --git a/python/pyspark/eval_handlers/tests/test_stateful_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_stateful_eval_type_handlers.py
index b3d4be27e4ae..6f8646926f01 100644
--- a/python/pyspark/eval_handlers/tests/test_stateful_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_stateful_eval_type_handlers.py
@@ -30,6 +30,7 @@ from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._typing import UDF
from pyspark.serializers import CPickleSerializer
from pyspark.sql.types import LongType, Row, StructField, StructType
from pyspark.testing.utils import (
@@ -116,7 +117,15 @@ class TransformWithStateHandlerTestCase(PySparkBaseTestCase):
@classmethod
def _handler(cls, handler_cls, udf, arg_offsets, runner_conf=None):
return handler_cls(
- udfs=[(udf, arg_offsets, cls._RETURN_TYPE)],
+ udfs=[
+ UDF(
+ func=udf,
+ raw_func=udf,
+ args_offsets=arg_offsets,
+ kwargs_offsets={},
+ return_type=cls._RETURN_TYPE,
+ )
+ ],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=cls._eval_conf(),
)
@@ -226,7 +235,15 @@ class PandasGroupedMapWithStateUDFHandlerTests(PySparkBaseTestCase):
@classmethod
def _handler(cls, udf, runner_conf=None):
return PandasGroupedMapWithStateUDFHandler(
- udfs=[(udf, cls._ARG_OFFSETS, cls._RETURN_TYPE)],
+ udfs=[
+ UDF(
+ func=udf,
+ raw_func=udf,
+ args_offsets=cls._ARG_OFFSETS,
+ kwargs_offsets={},
+ return_type=cls._RETURN_TYPE,
+ )
+ ],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=EvalConf({"state_value_schema": cls._STATE_SCHEMA.json()}),
)
diff --git a/python/pyspark/eval_handlers/tests/test_udtf_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_udtf_eval_type_handlers.py
index 684d596e42e8..038a5a2cba04 100644
--- a/python/pyspark/eval_handlers/tests/test_udtf_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_udtf_eval_type_handlers.py
@@ -17,9 +17,9 @@
"""Tests for the UDTF eval type handlers (``_udtf``).
-The handler gets the ``(udtf, args_offsets, kwargs_offsets, return_type)`` tuple
-``read_single_udtf`` produces; each test builds that tuple around a UDTF instance,
-calls ``run(0, <input batches>)`` and asserts on the output batches.
+The handler gets the ``UDTF`` record ``read_single_udtf`` returns; each test builds that
+record around a UDTF instance, calls ``run(0, <input batches>)`` and asserts on the output
+batches.
"""
import os
@@ -28,6 +28,7 @@ from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._typing import UDTF
from pyspark.eval_handlers._udtf import TableUDTFHandler, UDTFWithPartitions
from pyspark.serializers import BatchedSerializer, CPickleSerializer
from pyspark.sql.functions import SkipRestOfInputTableException
@@ -137,7 +138,14 @@ class TableUDTFHandlerTests(UDTFEvalTypeHandlerTestCase):
def _handler(self, udtf, args=(0,), kwargs=None, return_type=None):
return TableUDTFHandler(
- udfs=[(udtf, list(args), kwargs or {}, return_type or self._ROW_RETURN_TYPE)],
+ udfs=[
+ UDTF(
+ udtf=udtf,
+ args_offsets=list(args),
+ kwargs_offsets=kwargs or {},
+ return_type=return_type or self._ROW_RETURN_TYPE,
+ )
+ ],
runner_conf=RunnerConf({}),
eval_conf=EvalConf({}),
)
@@ -241,7 +249,14 @@ class ArrowUDTFHandlerTests(UDTFEvalTypeHandlerTestCase):
def _handler(self, udtf, args=(0,), kwargs=None, table_arg_offsets="0"):
conf = {} if table_arg_offsets is None else {"table_arg_offsets": table_arg_offsets}
return ArrowUDTFHandler(
- udfs=[(udtf, list(args), kwargs or {}, self._RETURN_TYPE)],
+ udfs=[
+ UDTF(
+ udtf=udtf,
+ args_offsets=list(args),
+ kwargs_offsets=kwargs or {},
+ return_type=self._RETURN_TYPE,
+ )
+ ],
runner_conf=RunnerConf({}),
eval_conf=EvalConf(conf),
)
@@ -401,7 +416,14 @@ class ArrowTableUDTFHandlerTestsMixin:
def _handler(self, udtf, args=(0,), kwargs=None, return_type=None):
return self.handler_cls(
- udfs=[(udtf, list(args), kwargs or {}, return_type or self._ROW_RETURN_TYPE)],
+ udfs=[
+ UDTF(
+ udtf=udtf,
+ args_offsets=list(args),
+ kwargs_offsets=kwargs or {},
+ return_type=return_type or self._ROW_RETURN_TYPE,
+ )
+ ],
runner_conf=RunnerConf({}),
eval_conf=EvalConf({"input_type": self._ROW_INPUT_TYPE.json()}),
)
diff --git a/python/pyspark/eval_handlers/tests/test_vanilla_eval_type_handlers.py b/python/pyspark/eval_handlers/tests/test_vanilla_eval_type_handlers.py
index 0b3e878ec876..72a4daa6fa37 100644
--- a/python/pyspark/eval_handlers/tests/test_vanilla_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_vanilla_eval_type_handlers.py
@@ -17,17 +17,20 @@
"""Tests for the plain Python (pickle) UDF eval type handler (``_vanilla``).
-The handler is built from ``(arg_offsets, eval_func)`` pairs, the shape ``read_single_udf``
-prepares for SQL_BATCHED_UDF, and fed pickled-row input via ``run(0, <rows>)``.
+The handler is built from ``UDF`` records, as ``read_single_udf`` returns them, and fed
+pickled-row input via ``run(0, <rows>)``.
"""
+import datetime
import io
import os
from unittest.mock import patch
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._typing import UDF
from pyspark.eval_handlers._vanilla import BatchedUDFHandler
from pyspark.serializers import BatchedSerializer, CPickleSerializer
+from pyspark.sql.types import DateType, LongType
from pyspark.testing.utils import PySparkBaseTestCase
from pyspark.util import PythonEvalType
@@ -37,9 +40,21 @@ with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
class BatchedUDFHandlerTests(PySparkBaseTestCase):
@staticmethod
- def _handler(*udfs):
+ def _handler(*udfs, kwargs_offsets=None, return_type=None):
+ """Build the handler from ``(arg_offsets, func)`` pairs."""
return BatchedUDFHandler(
- udfs=list(udfs), runner_conf=RunnerConf({}), eval_conf=EvalConf({})
+ udfs=[
+ UDF(
+ func=f,
+ raw_func=f,
+ args_offsets=arg_offsets,
+ kwargs_offsets=kwargs_offsets or {},
+ return_type=return_type or LongType(),
+ )
+ for arg_offsets, f in udfs
+ ],
+ runner_conf=RunnerConf({}),
+ eval_conf=EvalConf({}),
)
def test_registered(self):
@@ -71,6 +86,18 @@ class BatchedUDFHandlerTests(PySparkBaseTestCase):
rows = [(1, 10), (2, 20)]
self.assertEqual(list(handler.run(0, iter(rows))), [(2, 10), (3, 40)])
+ def test_kwargs_are_bound_by_name(self):
+ handler = self._handler(([1], lambda a, b: a - b), kwargs_offsets={"b": 0})
+ rows = [(1, 10), (2, 20)]
+ self.assertEqual(list(handler.run(0, iter(rows))), [9, 18])
+
+ def test_result_converted_when_return_type_needs_conversion(self):
+ handler = self._handler(
+ ([0], lambda d: d + datetime.timedelta(days=1)), return_type=DateType()
+ )
+ rows = [(datetime.date(1970, 1, 1),), (datetime.date(1970, 1, 10),)]
+ self.assertEqual(list(handler.run(0, iter(rows))), [1, 10])
+
def test_empty_input(self):
handler = self._handler(([0], lambda a: a))
self.assertEqual(list(handler.run(0, iter([]))), [])
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index f9e305707b23..931c4a390a57 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -35,8 +35,8 @@ from pyspark.accumulators import (
)
from pyspark.errors import PySparkRuntimeError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._typing import UDF, UDTF
from pyspark.eval_handlers._udtf import TableUDTFHandler, UDTFEvalTypeHandler
-from pyspark.eval_handlers.utils import wrap_kwargs_support
from pyspark.logger.worker_io import capture_outputs
from pyspark.messages import (
SparkMessageReceiver,
@@ -201,86 +201,20 @@ def read_single_udf(pickleSer, udf_info, handler_cls, runner_conf):
# when they are processed in a for loop, raise them as RuntimeError's instead
func = fail_on_stopiteration(profiling_func)
- args_offsets, kwargs_offsets = udf_info.args, udf_info.kwargs
-
- # The last returnType will be the return type of UDF. Eval types are grouped below by the
- # shape of the value they return.
-
- # Incremental Python aggregators: the pickled "function" is the Aggregator object itself, whose
- # zero/reduce/merge/finish methods the worker calls directly. Return it unwrapped (not through
- # fail_on_stopiteration, which would treat it as a plain callable).
- if eval_type in (
- PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
- PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
- PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
- ):
- return chained_func, args_offsets, kwargs_offsets, return_type
-
- # Scalar, aggregation and window UDFs: (func, args_offsets, kwargs_offsets, return_type).
- if eval_type in (
- PythonEvalType.SQL_ARROW_BATCHED_UDF,
- PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
- PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
- PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
- 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,
- PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF,
- PythonEvalType.SQL_SCALAR_ARROW_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
- PythonEvalType.SQL_SCALAR_PANDAS_UDF,
- PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
- PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
- ):
- return func, args_offsets, kwargs_offsets, return_type
- # Grouped-map and cogrouped-map UDFs: (func, args_offsets, return_type, num_udf_args).
- elif eval_type in (
- PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF,
- PythonEvalType.SQL_COGROUPED_MAP_PANDAS_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,
- ):
- # signature was lost when wrapping it
- num_udf_args = len(inspect.getfullargspec(chained_func).args)
- return func, args_offsets, return_type, num_udf_args
- # Grouped-map-with-state and transform-with-state UDFs: (func, args_offsets, return_type).
- elif eval_type in (
- PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
- PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
- PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF,
- PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF,
- PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF,
- ):
- return func, args_offsets, return_type
- # Map iterator UDFs take no offsets.
- elif eval_type == PythonEvalType.SQL_MAP_PANDAS_ITER_UDF:
- return func, None, None, return_type
- elif eval_type == PythonEvalType.SQL_MAP_ARROW_ITER_UDF:
- return func, None, None, None
- # Batched (plain Python) UDFs: (args_kwargs_offsets, eval func); apply kwargs binding and
- # convert each result to the internal representation only when the return type requires it.
- elif eval_type == PythonEvalType.SQL_BATCHED_UDF:
- func, args_kwargs_offsets = wrap_kwargs_support(func, args_offsets, kwargs_offsets)
- if return_type.needConversion():
- toInternal = return_type.toInternal
- return args_kwargs_offsets, lambda *a: toInternal(func(*a))
- else:
- return args_kwargs_offsets, lambda *a: func(*a)
- else:
- raise ValueError("Unknown eval type: {}".format(eval_type))
+ # The last returnType will be the return type of UDF.
+ return UDF(
+ func=func,
+ raw_func=chained_func,
+ args_offsets=udf_info.args,
+ kwargs_offsets=udf_info.kwargs,
+ return_type=return_type,
+ )
# Read a serialized user-defined table function (UDTF) and prepare it for evaluation: wrap
# the constructor to receive the AnalyzeResult, instantiate the UDTF (wrapped in the handler's
# partition_wrapper when its TABLE argument has PARTITION BY expressions), and check the 'eval'
-# signature against the call arguments. Returns the same
-# (func, args_offsets, kwargs_offsets, return_type) shape as read_single_udf, with the UDTF
-# instance in the func slot.
+# signature against the call arguments.
def read_single_udtf(pickleSer, udtf_info, handler_cls):
if udtf_info.pickled_analyze_result is not None:
pickled_analyze_result = pickleSer.loads(udtf_info.pickled_analyze_result)
@@ -356,7 +290,12 @@ def read_single_udtf(pickleSer, udtf_info, handler_cls):
messageParameters={"name": udtf_info.name, "reason": str(e)},
) from None
- return udtf, udtf_info.args, udtf_info.kwargs, return_type
+ return UDTF(
+ udtf=udtf,
+ args_offsets=udtf_info.args,
+ kwargs_offsets=udtf_info.kwargs,
+ return_type=return_type,
+ )
def invoke_udf(message_receiver: SparkMessageReceiver, outfile: BinaryIO):
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org