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 2bbb9e248c76 [SPARK-60532][PYTHON] Chain UDF functions in
EvalTypeHandler and enforce single-UDF handlers in the base
2bbb9e248c76 is described below
commit 2bbb9e248c76f0c018671ab919d2338579a453a8
Author: Yicong Huang <[email protected]>
AuthorDate: Tue Oct 27 11:00:31 2026 +0000
[SPARK-60532][PYTHON] Chain UDF functions in EvalTypeHandler and enforce
single-UDF handlers in the base
### What changes were proposed in this pull request?
Two handler API changes under the EvalTypeHandler umbrella (SPARK-59415):
1. **Handlers take one UDF unless they declare otherwise.**
`EvalTypeHandler` gets a `_multi_udf` class flag, `False` by default. Handlers
that take any number of UDFs set it (scalar pandas/Arrow, batched,
Arrow-batched, elementwise, grouped/window agg, incremental agg/window). For
every other handler (grouped/cogrouped map, map iter, scalar iter, grouped agg
iter, scalar iter elementwise, stateful, UDTF), `__init__` fails fast with a
`PySparkRuntimeError` naming the eval type (as the [...]
2. **Chaining nested UDF calls moves from the worker into the handler.**
`worker.py` only deserializes (`unpickle_and_chain` and `chain` removed,
`read_command` inlined in `invoke_udf`), and the `UDF` record carries the
unchained `funcs` instead of `raw_func`. `EvalTypeHandler.__init__` chains them
into `_raw_funcs`, then wraps for profiling and `fail_on_stopiteration` as
before. A single function is kept as is, so incremental `Aggregator` objects,
`getfullargspec` and the data-source [...]
### Why are the changes needed?
Keeps the worker to reading the wire format and gives handlers one place
for their UDF invariants instead of repeating them per eval type.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New unit tests in `test_base_eval_type_handlers.py` (single-UDF guard,
`_multi_udf` per handler, chaining order, single function not wrapped); the
other handler tests are updated for the `funcs` field. Existing e2e tests pass,
including nested-UDF chaining in `test_udf` / `test_pandas_udf*`,
`test_udf_profiler`, `test_python_datasource`, `test_udtf`, and
transform-with-state.
ASV (`bench_eval_type`, 11 TimeBench classes) shows no regression;
per-bench geomean of after/before is 0.96-1.01. Representative run,
`GroupedMapArrowUDFTimeBench`, `identity_udf`:
| scenario | before | after |
|---|---|---|
| few_groups_sm | 5.01±0.08ms | 5.45±0.3ms |
| few_groups_lg | 40.4±2ms | 39.1±0.3ms |
| many_groups_sm | 162±4ms | 157±4ms |
| many_groups_lg | 104±2ms | 103±0.9ms |
| wide_cols | 147±2ms | 145±0.7ms |
| multi_key | 28.9±0.7ms | 28.4±0.1ms |
`few_groups_sm` does not reproduce in an interleaved `timeit` run of
`time_worker` (4.68ms before vs 4.65ms after, median); the per-batch path is
unchanged.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
This pull request and its description were written by Isaac.
Closes #59776 from Yicong-Huang/single-udf-chain-in-handler.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit 834d036612dd5853ff0db69730ccee80b38f704d)
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/eval_handlers/_arrow.py | 176 +++++++++++----------
python/pyspark/eval_handlers/_base.py | 48 +++++-
python/pyspark/eval_handlers/_elementwise.py | 45 +++---
python/pyspark/eval_handlers/_pandas.py | 154 ++++++++++--------
python/pyspark/eval_handlers/_stateful.py | 86 +++++-----
python/pyspark/eval_handlers/_typing.py | 6 +-
python/pyspark/eval_handlers/_udtf.py | 71 +++++----
python/pyspark/eval_handlers/_vanilla.py | 11 +-
.../tests/test_arrow_eval_type_handlers.py | 41 +++--
.../tests/test_base_eval_type_handlers.py | 150 ++++++++++++++----
.../tests/test_elementwise_eval_type_handlers.py | 8 +-
.../tests/test_pandas_eval_type_handlers.py | 24 +--
.../tests/test_stateful_eval_type_handlers.py | 8 +-
.../tests/test_udtf_eval_type_handlers.py | 14 +-
.../tests/test_vanilla_eval_type_handlers.py | 8 +-
python/pyspark/worker.py | 42 ++---
16 files changed, 516 insertions(+), 376 deletions(-)
diff --git a/python/pyspark/eval_handlers/_arrow.py
b/python/pyspark/eval_handlers/_arrow.py
index dd75347a0bf6..fb635c6b1fca 100644
--- a/python/pyspark/eval_handlers/_arrow.py
+++ b/python/pyspark/eval_handlers/_arrow.py
@@ -79,18 +79,18 @@ class
ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
# 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._num_udf_args = len(inspect.getfullargspec(self._raw_func).args)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.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(
- udf.return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ self._udf_desc.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]:
@@ -109,11 +109,11 @@ class
ArrowCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
right_values = table_from_batches(right_batches,
self._right_val_cols)
if self._num_udf_args == 2:
- result = self._funcs[0](left_values, right_values)
+ result = self._func(left_values, right_values)
else:
key_table = left_keys if left_keys.num_rows > 0 else right_keys
key = tuple(c[0] for c in key_table.columns)
- result = self._funcs[0](key, left_values, right_values)
+ result = self._func(key, left_values, right_values)
verify_return_type(result, pa.Table)
# Verify types (and reorder by name when configured).
@@ -136,21 +136,21 @@ class
ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_ITER_UDF
_iter_based = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
# 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._num_udf_args = len(inspect.getfullargspec(self._raw_func).args)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.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(
- udf.return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ self._udf_desc.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]:
@@ -179,10 +179,10 @@ class
ArrowGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
# Call UDF with iterator of batches
if self._num_udf_args == 1:
- result = self._funcs[0](value_batches)
+ result = self._func(value_batches)
else:
key = tuple(c[0] for c in keys.columns)
- result = self._funcs[0](key, value_batches)
+ result = self._func(key, value_batches)
# Verify (and reorder by name when configured) each output batch
for batch in verify_return_type(result, Iterator[pa.RecordBatch]):
@@ -206,19 +206,19 @@ class
ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
# 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._num_udf_args = len(inspect.getfullargspec(self._raw_func).args)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.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(
- udf.return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ self._udf_desc.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]:
@@ -248,10 +248,10 @@ class
ArrowGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
# Call UDF
value_table = pa.Table.from_batches(value_batches)
if self._num_udf_args == 1:
- result = self._funcs[0](value_table)
+ result = self._func(value_table)
else:
key = tuple(c[0] for c in keys.columns)
- result = self._funcs[0](key, value_table)
+ result = self._func(key, value_table)
verify_return_type(result, pa.Table)
# Verify types (and reorder by name when configured).
@@ -272,13 +272,16 @@ class
ArrowGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
coerced to the declared schema."""
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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))]
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._col_names = ["_%d" % i for i in range(len(udf_descs))]
self._return_schema = to_arrow_schema(
- StructType([StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udfs)]),
+ StructType(
+ [StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udf_descs)]
+ ),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -293,7 +296,7 @@ class
ArrowGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
*[concatenated.column(o) for o in udf.args_offsets],
**{k: concatenated.column(v) for k, v in
udf.kwargs_offsets.items()},
)
- for udf, func in zip(self._udfs, self._funcs)
+ for udf, func in zip(self._udf_descs, self._funcs)
]
result_arrays = [pa.array([r]) for r in results]
batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names)
@@ -306,13 +309,11 @@ class
ArrowGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
self._return_schema = to_arrow_schema(
- StructType([StructField("_0", udf.return_type)]),
+ StructType([StructField("_0", self._udf_desc.return_type)]),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -320,7 +321,7 @@ class
ArrowGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
def run(self, split_index: int, data: Iterator[GroupedBatch]) ->
Iterator[pa.RecordBatch]:
import pyarrow as pa
- args_offsets = self._udfs[0].args_offsets
+ args_offsets = self._udf_desc.args_offsets
def extract_args(batch: pa.RecordBatch) -> Any:
args = tuple(batch.column(o) for o in args_offsets)
@@ -328,7 +329,7 @@ class
ArrowGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
for group in data:
batch_iter = map(extract_args, group)
- result = self._funcs[0](batch_iter)
+ result = self._func(batch_iter)
# Drain remaining batches to maintain stream position
for _ in batch_iter:
pass
@@ -343,14 +344,17 @@ class
ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
UDF, coerced to the declared schema."""
eval_type = PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, 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._col_names = ["_%d" % i for i in range(len(udf_descs))]
self._return_schema = to_arrow_schema(
- StructType([StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udfs)]),
+ StructType(
+ [StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udf_descs)]
+ ),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -363,7 +367,7 @@ class
ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
num_rows = concatenated.num_rows
result_arrays = []
- for udf_index, (udf, udf_func) in enumerate(zip(self._udfs,
self._funcs)):
+ for udf_index, (udf, udf_func) in enumerate(zip(self._udf_descs,
self._funcs)):
args_offsets, kwargs_offsets = udf.args_offsets,
udf.kwargs_offsets
bound_type = self._window_bound_types[udf_index]
if bound_type == "unbounded":
@@ -412,10 +416,11 @@ class
ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.Recor
compares unequal to itself) are merged downstream."""
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
# The leading `num_grouping_keys` input columns are the grouping keys
(see the operator);
# `grouping_key_schema` carries their names/types so the emitted key
columns round-trip.
grouping_key_schema = eval_conf.grouping_key_schema
@@ -425,8 +430,8 @@ class
ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.Recor
# 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._aggs = self._raw_funcs
+ self._buffer_col_names = ["_%d" % i for i in range(len(udf_descs))]
self._buffer_arrow_types = [
to_arrow_type(
agg.bufferSchema,
@@ -443,7 +448,7 @@ class
ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.Recor
# 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(udf.args_offsets) + list(udf.kwargs_offsets.values()) for udf
in udfs
+ list(udf.args_offsets) + list(udf.kwargs_offsets.values()) for udf
in udf_descs
]
# 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
@@ -540,18 +545,21 @@ class
ArrowGroupedAggIncrementalFinalUDFHandler(GroupedEvalTypeHandler["pa.Recor
by the operator) still produces ``finish(zero)``."""
eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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))]
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._col_names = ["_%d" % i for i in range(len(udf_descs))]
self._return_schema = to_arrow_schema(
- StructType([StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udfs)]),
+ StructType(
+ [StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udf_descs)]
+ ),
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]
+ self._aggs = self._raw_funcs
# 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 self._aggs]
@@ -561,7 +569,7 @@ class
ArrowGroupedAggIncrementalFinalUDFHandler(GroupedEvalTypeHandler["pa.Recor
for group in data:
merged: list = [None] * len(self._aggs)
for batch in group:
- for i, (udf, agg) in enumerate(zip(self._udfs, self._aggs)):
+ for i, (udf, agg) in enumerate(zip(self._udf_descs,
self._aggs)):
field_names = self._field_names_by_udf[i]
m = merged[i]
for row in batch.column(udf.args_offsets[0]).to_pylist():
@@ -592,19 +600,22 @@ class
ArrowWindowAggIncrementalUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch
intermediate buffer never leaves the worker; ``merge`` is not used here."""
eval_type = PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, 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._col_names = ["_%d" % i for i in range(len(udf_descs))]
self._return_schema = to_arrow_schema(
- StructType([StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udfs)]),
+ StructType(
+ [StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udf_descs)]
+ ),
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]
+ self._aggs = self._raw_funcs
def run(self, split_index: int, data: Iterator[GroupedBatch]) ->
Iterator[pa.RecordBatch]:
import pyarrow as pa
@@ -621,7 +632,7 @@ class
ArrowWindowAggIncrementalUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch
num_rows = concatenated.num_rows
result_arrays = []
- for udf_index, (udf, agg) in enumerate(zip(self._udfs,
self._aggs)):
+ for udf_index, (udf, agg) in enumerate(zip(self._udf_descs,
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
@@ -679,10 +690,9 @@ class
ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF
_iter_based = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
+ super().__init__(udf_descs, runner_conf, eval_conf)
def run(self, split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
import pyarrow as pa
@@ -691,7 +701,7 @@ class
ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
input_batches = map(ArrowBatchTransformer.flatten_struct, data)
# invoke the UDF
- output_batches = self._funcs[0](input_batches)
+ output_batches = self._func(input_batches)
# The declared signature is Iterator[...], so a strict iterator is
required by
# default. With the legacy flag, accept any object Python can iterate
over -- via
@@ -720,19 +730,19 @@ class
ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_SCALAR_ARROW_ITER_UDF
_iter_based = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
self._arrow_return_type = to_arrow_type(
- udf.return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ self._udf_desc.return_type,
+ timezone="UTC",
+ prefers_large_types=runner_conf.use_large_var_types,
)
def run(self, split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
import pyarrow as pa
- args_offsets = self._udfs[0].args_offsets
+ args_offsets = self._udf_desc.args_offsets
num_input_rows = 0
def extract_args(batch: pa.RecordBatch) -> Any:
@@ -745,7 +755,7 @@ class
ArrowScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
args_iter = map(extract_args, data)
# Call UDF and verify result type (iterator of pa.Array)
- verified_iter = verify_return_type(self._funcs[0](args_iter),
Iterator[pa.Array])
+ verified_iter = verify_return_type(self._func(args_iter),
Iterator[pa.Array])
# Process results: enforce schema and assemble into RecordBatch
target_schema = pa.schema([pa.field("_0", self._arrow_return_type)])
@@ -773,13 +783,16 @@ class
ArrowScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
the result to the declared schema, and check the row count."""
eval_type = PythonEvalType.SQL_SCALAR_ARROW_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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))]
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._col_names = ["_%d" % i for i in range(len(udf_descs))]
self._combined_arrow_schema = to_arrow_schema(
- StructType([StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udfs)]),
+ StructType(
+ [StructField(n, udf.return_type) for n, udf in
zip(self._col_names, udf_descs)]
+ ),
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -794,7 +807,7 @@ class
ArrowScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
*[batch.column(o) for o in udf.args_offsets],
**{k: batch.column(v) for k, v in
udf.kwargs_offsets.items()},
)
- for udf, func in zip(self._udfs, self._funcs)
+ for udf, func in zip(self._udf_descs, self._funcs)
],
self._col_names,
)
@@ -813,16 +826,17 @@ class
ArrowBatchedUDFArrowHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_ARROW_BATCHED_UDF
variant_conf = "use_legacy_pandas_udf_conversion"
variant_value = False
+ _multi_udf = True
# run wraps the row-by-row evaluation in fail_on_stopiteration itself.
_fail_on_stopiteration = False
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
# --- UDF preparation ---
self._udf_infos = []
- for udf, func in zip(udfs, self._funcs):
+ for udf, func in zip(udf_descs, self._funcs):
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
func, udf.args_offsets, udf.kwargs_offsets
)
@@ -845,7 +859,7 @@ class
ArrowBatchedUDFArrowHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
)
)
assert self._udf_infos, "At least one Arrow UDF expected here."
- self._col_names = [f"_{i}" for i in range(len(udfs))]
+ self._col_names = [f"_{i}" for i in range(len(udf_descs))]
# --- Input preparation ---
self._arrow_to_py_converters = [
diff --git a/python/pyspark/eval_handlers/_base.py
b/python/pyspark/eval_handlers/_base.py
index 327f5189c4cd..98e9c91b6830 100644
--- a/python/pyspark/eval_handlers/_base.py
+++ b/python/pyspark/eval_handlers/_base.py
@@ -21,10 +21,12 @@ This leaf module imports only ``_typing``, the serializers
and ``pyspark.util``,
package ``__init__`` and the concrete-handler submodules can both import it.
"""
+import functools
from abc import ABCMeta, abstractmethod
from collections.abc import Iterator
from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, Optional
+from pyspark.errors import PySparkRuntimeError
from pyspark.eval_handlers._typing import (
UDF,
UDTF,
@@ -39,7 +41,7 @@ from pyspark.sql.pandas.serializers import (
ArrowStreamGroupSerializer,
ArrowStreamSerializer,
)
-from pyspark.util import fail_on_stopiteration
+from pyspark.util import PythonEvalType, fail_on_stopiteration
if TYPE_CHECKING:
import pyarrow as pa # noqa: F401 # only in the batch category's
forward-ref subscript
@@ -153,24 +155,56 @@ class EvalTypeHandler(Generic[InputBatch, OutputBatch],
metaclass=_EvalTypeHandl
# Whether __init__ wraps the UDF in fail_on_stopiteration, so a
StopIteration raised by
# the UDF call fails with STOP_ITERATION_OCCURRED instead of PEP 479's
generic RuntimeError.
_fail_on_stopiteration: ClassVar[bool] = True
+ # Whether the handler takes any number of UDFs. Otherwise it takes exactly
one UDF or UDTF:
+ # __init__ exposes a UDF as _udf_desc and its functions as _raw_func and
_func, and
+ # UDTFHandlerBase exposes a UDTF as _udtf_desc.
+ _multi_udf: ClassVar[bool] = False
def __init__(
- self, udfs: "list[UDF] | list[UDTF]", runner_conf: "RunnerConf",
eval_conf: "EvalConf"
+ self, udf_descs: "list[UDF] | list[UDTF]", runner_conf: "RunnerConf",
eval_conf: "EvalConf"
) -> None:
- self._udfs: list[Any] = udfs
+ if not self._multi_udf and len(udf_descs) != 1:
+ raise PySparkRuntimeError(
+ "One {} expected, but got {}.".format(self._expected_udf(),
len(udf_descs))
+ )
+ self._udf_descs: list[Any] = udf_descs
self._runner_conf = runner_conf
self._eval_conf = eval_conf
- # Each UDF's function wrapped for profiling and, unless
_fail_on_stopiteration is False,
- # in fail_on_stopiteration, at the indexes of _udfs; empty for a UDTF.
+ # Each UDF's functions chained into one before any wrapping, at the
indexes of _udf_descs: a
+ # single function is the pickled object itself (an incremental
Aggregator) with its
+ # signature. Empty for a UDTF.
+ self._raw_funcs: list[Any] = []
+ # Each raw function wrapped for profiling and, unless
_fail_on_stopiteration is False,
+ # in fail_on_stopiteration, at the indexes of _udf_descs; empty for a
UDTF.
self._funcs: list[Callable[..., Any]] = []
- for udf in udfs:
+ for udf in udf_descs:
if isinstance(udf, UDF):
- profiling_func = self._with_profiler(udf.raw_func,
udf.result_id)
+ raw_func = functools.reduce(lambda f, g: lambda *a: g(f(*a)),
udf.funcs)
+ self._raw_funcs.append(raw_func)
+ profiling_func = self._with_profiler(raw_func, udf.result_id)
self._funcs.append(
fail_on_stopiteration(profiling_func)
if self._fail_on_stopiteration
else profiling_func
)
+ if not self._multi_udf and isinstance(udf_descs[0], UDF):
+ self._udf_desc: UDF = udf_descs[0]
+ self._raw_func: Any = self._raw_funcs[0]
+ self._func: Callable[..., Any] = self._funcs[0]
+
+ def _expected_udf(self) -> str:
+ """What a single-UDF handler takes, as named when it is given another
number of UDFs:
+ its eval type without the ``SQL_`` prefix and ``_UDF`` suffix, e.g.
``MAP_ARROW_ITER
+ UDF``."""
+ name = next(
+ (
+ k
+ for k, v in vars(PythonEvalType).items()
+ if k.startswith("SQL_") and v == self.eval_type
+ ),
+ type(self).__name__,
+ )
+ return "{} UDF".format(name.removeprefix("SQL_").removesuffix("_UDF"))
def _with_profiler(self, f: Callable[..., Any], result_id: int) ->
Callable[..., Any]:
"""Return ``f`` profiled by the profiler its config selects, or ``f``
itself when none is
diff --git a/python/pyspark/eval_handlers/_elementwise.py
b/python/pyspark/eval_handlers/_elementwise.py
index 7646b8254026..6c8e0e9a6e6f 100644
--- a/python/pyspark/eval_handlers/_elementwise.py
+++ b/python/pyspark/eval_handlers/_elementwise.py
@@ -179,12 +179,13 @@ class
ArrowElementwiseUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
per element tuple, convert the results to the declared Arrow type and
re-nest them."""
eval_type = PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
# Element-wise UDFs back higher-order lambdas like transform(arr, x ->
udf(x)).
# ExtractPythonUDFFromLambda rewrites them so the UDF receives *all*
array elements
@@ -197,7 +198,7 @@ class
ArrowElementwiseUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
# UDF preparation
nesting = eval_conf.elementwise_nesting
self._udf_infos = []
- for udf_index, udf in enumerate(udfs):
+ for udf_index, udf in enumerate(udf_descs):
udf_return_type = udf.return_type
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
self._funcs[udf_index], udf.args_offsets, udf.kwargs_offsets
@@ -235,7 +236,7 @@ class
ArrowElementwiseUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
),
)
)
- self._col_names = [f"_{i}" for i in range(len(udfs))]
+ self._col_names = [f"_{i}" for i in range(len(udf_descs))]
def run(self, split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
import pyarrow as pa
@@ -306,9 +307,9 @@ 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[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
def _to_udf_inputs(self, flat_batch: pa.RecordBatch, input_schema:
StructType) -> list[Any]:
"""Adapt one flattened input batch (one aligned leaf Array per
argument) to the UDF's
@@ -386,12 +387,14 @@ 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[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
- super().__init__(udfs, runner_conf, eval_conf)
+ _multi_udf = True
+
+ def __init__(self, udf_descs: list[UDF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
+ super().__init__(udf_descs, runner_conf, eval_conf)
nesting = eval_conf.elementwise_nesting
self._udf_infos = []
- for udf_index, udf in enumerate(udfs):
+ for udf_index, udf in enumerate(udf_descs):
udf_return_type = udf.return_type
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
self._funcs[udf_index], udf.args_offsets, udf.kwargs_offsets
@@ -413,7 +416,7 @@ class
_ScalarElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
udf_input_schema,
)
)
- self._col_names = [f"_{i}" for i in range(len(udfs))]
+ self._col_names = [f"_{i}" for i in range(len(udf_descs))]
@abstractmethod
def _verify_result(self, result: Any, return_type: DataType, num_elements:
int) -> None:
@@ -512,12 +515,10 @@ class
_ScalarIterElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
_iter_based = True
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
- super().__init__(udfs, runner_conf, eval_conf)
+ def __init__(self, udf_descs: list[UDF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
+ super().__init__(udf_descs, runner_conf, eval_conf)
- assert len(udfs) == 1, "One SCALAR_*_ITER_ELEMENTWISE UDF expected
here."
- udf = udfs[0]
- assert not udf.kwargs_offsets, "Iterator UDFs do not take keyword
arguments."
+ assert not self._udf_desc.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
@@ -525,9 +526,11 @@ class
_ScalarIterElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
nesting = eval_conf.elementwise_nesting
self._depth = nesting[0] if nesting is not None else 1
self._udf_input_schema = _elementwise_input_schema(
- eval_conf.input_type, udf.args_offsets, self._depth
+ eval_conf.input_type, self._udf_desc.args_offsets, self._depth
+ )
+ self._udf_output_schema = _elementwise_output_schema(
+ self._udf_desc.return_type, runner_conf
)
- self._udf_output_schema = _elementwise_output_schema(udf.return_type,
runner_conf)
@abstractmethod
def _verify_return_type(self, results: Any) -> Iterator[Any]:
@@ -549,7 +552,7 @@ class
_ScalarIterElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
# shapes re-nest this batch's rows. The user function sees the
flat leaves as a
# pandas Series / DataFrame (pandas) or a pa.Array (Arrow), each
with its leaf type.
flat_batch, list_lengths_by_level, large_list_by_level =
_elementwise_flatten_inputs(
- batch, self._udfs[0].args_offsets, self._depth
+ batch, self._udf_desc.args_offsets, self._depth
)
pending_shapes.append((list_lengths_by_level, large_list_by_level,
flat_batch.num_rows))
num_input_elements += flat_batch.num_rows
@@ -558,7 +561,7 @@ class
_ScalarIterElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
flat_args_iter = map(extract_flat, data)
- verified_iter =
self._verify_return_type(self._funcs[0](flat_args_iter))
+ verified_iter = self._verify_return_type(self._func(flat_args_iter))
# Buffer the streamed flat element results and emit an ``array<R>``
row as soon as the
# shape at the head of the FIFO is fully covered. A row whose length
is 0 (an empty
@@ -601,7 +604,7 @@ class
_ScalarIterElementwiseUDFHandler(_VectorizedElementwiseUDFHandler):
nonlocal pending_chunks, pending_len, empty_schema,
num_output_elements
for result in verified_iter:
chunk = self._to_flat_batch(
- result, self._udfs[0].return_type, self._udf_output_schema
+ result, self._udf_desc.return_type, self._udf_output_schema
)
num_output_elements += chunk.num_rows
# Fail fast if the UDF over-produces, before the buffer grows
unbounded (the
@@ -658,7 +661,7 @@ class PandasScalarIterElementwiseUDFHandler(
pandas_iter_type = (
Iterator[pd.DataFrame]
- if isinstance(self._udfs[0].return_type, StructType)
+ if isinstance(self._udf_desc.return_type, StructType)
else Iterator[pd.Series]
)
return verify_return_type(results, pandas_iter_type)
diff --git a/python/pyspark/eval_handlers/_pandas.py
b/python/pyspark/eval_handlers/_pandas.py
index a9215e95bb21..f5301e9e74b1 100644
--- a/python/pyspark/eval_handlers/_pandas.py
+++ b/python/pyspark/eval_handlers/_pandas.py
@@ -75,13 +75,16 @@ class
PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
``pyspark.sql.conversion`` directly with the runner_conf-derived
parameters."""
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf"
+ ) -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
self._return_schema = StructType(
- [StructField("_%d" % i, udf.return_type) for i, udf in
enumerate(udfs)]
+ [StructField("_%d" % i, udf.return_type) for i, udf in
enumerate(udf_descs)]
)
def run(self, split_index: int, data: "Iterator[pa.RecordBatch]") ->
"Iterator[pa.RecordBatch]":
@@ -102,7 +105,7 @@ class
PandasScalarUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
# Process: evaluate each UDF column-wise on pandas Series.
results = []
- for udf, func in zip(self._udfs, self._funcs):
+ for udf, func in zip(self._udf_descs, self._funcs):
return_type = udf.return_type
result = func(
*[pandas_columns[o] for o in udf.args_offsets],
@@ -151,18 +154,18 @@ class
PandasScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF
_iter_based = True
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: 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."
- udf = udfs[0]
- self._return_schema = StructType([StructField("_0", udf.return_type)])
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._return_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
self._expected_iter_type = (
Iterator[pd.DataFrame]
- if isinstance(udf.return_type, StructType)
+ if isinstance(self._udf_desc.return_type, StructType)
else Iterator[pd.Series]
)
@@ -180,19 +183,19 @@ class
PandasScalarIterUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
prefer_int_ext_dtype=self._runner_conf.prefer_int_ext_dtype,
df_for_struct=True,
)
- args = tuple(pandas_columns[o] for o in self._udfs[0].args_offsets)
+ args = tuple(pandas_columns[o] for o in
self._udf_desc.args_offsets)
num_input_rows += batch.num_rows
return args[0] if len(args) == 1 else args
# Extract args from input batches (streaming), then call the UDF and
verify the
# result type (iterator of pd.Series / pd.DataFrame).
args_iter = map(extract_args, data)
- verified_iter = verify_return_type(self._funcs[0](args_iter),
self._expected_iter_type)
+ verified_iter = verify_return_type(self._func(args_iter),
self._expected_iter_type)
def convert(result: Any) -> "pa.RecordBatch":
verify_pandas_result(
result,
- self._udfs[0].return_type,
+ self._udf_desc.return_type,
assign_cols_by_name=True,
truncate_return_schema=True,
)
@@ -224,17 +227,18 @@ class
PandasMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
_iter_based = True
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: 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._return_schema = StructType([StructField("_0",
udfs[0].return_type)])
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._return_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
self._expected_iter_type = (
Iterator[pd.DataFrame]
- if isinstance(udfs[0].return_type, StructType)
+ if isinstance(self._udf_desc.return_type, StructType)
else Iterator[pd.Series]
)
@@ -250,7 +254,7 @@ class
PandasMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
df_for_struct=True,
)[0]
- result = self._funcs[0](dataframe_iter())
+ result = self._func(dataframe_iter())
# The declared signature is Iterator[...], so a strict iterator is
required by
# default. With the legacy flag, accept any object Python can iterate
over -- via
# iter(...), which honors both __iter__ and the sequence protocol
(__getitem__) --
@@ -267,7 +271,10 @@ class
PandasMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
for df in verify_return_type(result, self._expected_iter_type):
verify_pandas_result(
- df, self._udfs[0].return_type, assign_cols_by_name=True,
truncate_return_schema=True
+ df,
+ self._udf_desc.return_type,
+ assign_cols_by_name=True,
+ truncate_return_schema=True,
)
yield PandasToArrowConversion.from_pandas(
[df],
@@ -289,18 +296,18 @@ class
PandasCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
# 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._num_udf_args = len(inspect.getfullargspec(self._raw_func).args)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.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", udf.return_type)])
+ self._return_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
def run(self, split_index: int, data: "Iterator[CoGroupedBatch]") ->
"Iterator[pa.RecordBatch]":
"""Apply cogroupBy Pandas UDF.
@@ -329,7 +336,7 @@ class
PandasCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
right_df = pd.concat([right_series[o] for o in
self._right_value_offsets], axis=1)
if self._num_udf_args == 2:
- result = self._funcs[0](left_df, right_df)
+ result = self._func(left_df, right_df)
else:
key_series = (
[left_series[o] for o in self._left_key_offsets]
@@ -337,14 +344,14 @@ class
PandasCoGroupedMapUDFHandler(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
else [right_series[o] for o in self._right_key_offsets]
)
key = tuple(s.iloc[0] for s in key_series)
- result = self._funcs[0](key, left_df, right_df)
+ result = self._func(key, left_df, right_df)
del left_batches, right_batches, left_table, right_table
del left_series, right_series, left_df, right_df
verify_pandas_result(
result,
- self._udfs[0].return_type,
+ self._udf_desc.return_type,
self._runner_conf.assign_cols_by_name,
truncate_return_schema=False,
)
@@ -371,20 +378,20 @@ class
PandasGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
_iter_based = True
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
# 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._num_udf_args = len(inspect.getfullargspec(self._raw_func).args)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.args_offsets)
assert len(parsed_offsets) == 1, (
"Expected one pair of offsets for GROUPED_MAP_PANDAS_ITER UDF."
)
self._key_offsets, self._value_offsets = parsed_offsets[0]
- self._return_schema = StructType([StructField("_0", udf.return_type)])
+ self._return_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
def run(self, split_index: int, data: "Iterator[GroupedBatch]") ->
"Iterator[pa.RecordBatch]":
"""Apply groupBy Pandas UDF (iterator variant)."""
@@ -412,15 +419,15 @@ class
PandasGroupedMapIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
yield pd.concat([series[o] for o in value_offsets], axis=1)
if self._num_udf_args == 1:
- result = self._funcs[0](dataframe_iter())
+ result = self._func(dataframe_iter())
else:
key = tuple(first_series[o].iloc[0] for o in self._key_offsets)
- result = self._funcs[0](key, dataframe_iter())
+ result = self._func(key, dataframe_iter())
for df in result:
verify_pandas_result(
df,
- self._udfs[0].return_type,
+ self._udf_desc.return_type,
self._runner_conf.assign_cols_by_name,
truncate_return_schema=False,
)
@@ -451,18 +458,18 @@ class
PandasGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: 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."
- udf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
# 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._num_udf_args = len(inspect.getfullargspec(self._raw_func).args)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.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", udf.return_type)])
+ self._return_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
def run(self, split_index: int, data: "Iterator[GroupedBatch]") ->
"Iterator[pa.RecordBatch]":
# Bound each output RecordBatch toward the byte cap (no limit when
unset).
@@ -498,16 +505,16 @@ class
PandasGroupedMapUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
value_df = pd.concat([all_series[o] for o in self._value_offsets],
axis=1)
if self._num_udf_args == 1:
- result = self._funcs[0](value_df)
+ result = self._func(value_df)
else:
key = tuple(all_series[o].iloc[0] for o in self._key_offsets)
- result = self._funcs[0](key, value_df)
+ result = self._func(key, value_df)
del all_batches, table, all_series, value_df
verify_pandas_result(
result,
- self._udfs[0].return_type,
+ self._udf_desc.return_type,
self._runner_conf.assign_cols_by_name,
truncate_return_schema=False,
)
@@ -532,13 +539,16 @@ class
PandasGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
ArrowGroupedAggUDFHandler."""
eval_type = PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf"
+ ) -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
self._return_schema = StructType(
- [StructField("_%d" % i, udf.return_type) for i, udf in
enumerate(udfs)]
+ [StructField("_%d" % i, udf.return_type) for i, udf in
enumerate(udf_descs)]
)
def run(self, split_index: int, data: "Iterator[GroupedBatch]") ->
"Iterator[pa.RecordBatch]":
@@ -560,7 +570,7 @@ class
PandasGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
*[all_series[o] for o in udf.args_offsets],
**{k: all_series[v] for k, v in
udf.kwargs_offsets.items()},
)
- for udf, func in zip(self._udfs, self._funcs)
+ for udf, func in zip(self._udf_descs, self._funcs)
]
yield PandasToArrowConversion.from_pandas(
[pd.Series([r]) for r in results],
@@ -582,18 +592,18 @@ class
PandasGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: 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."
- udf = udfs[0]
- self._return_schema = StructType([StructField("_0", udf.return_type)])
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._return_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
def run(self, split_index: int, data: "Iterator[GroupedBatch]") ->
"Iterator[pa.RecordBatch]":
import pandas as pd
- args_offsets = self._udfs[0].args_offsets
+ args_offsets = self._udf_desc.args_offsets
def extract_series(batch: "pa.RecordBatch") -> Any:
# Convert one RecordBatch to a pandas Series per column, then
select args:
@@ -609,7 +619,7 @@ class
PandasGroupedAggIterUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
for group in data:
series_iter = map(extract_series, group)
- result = self._funcs[0](series_iter)
+ result = self._func(series_iter)
# Drain remaining batches to maintain stream position
for _ in series_iter:
pass
@@ -632,14 +642,17 @@ class
PandasWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
converted back to Arrow. The pandas counterpart of
ArrowWindowAggUDFHandler."""
eval_type = PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF
+ _multi_udf = True
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf"
+ ) -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
self._window_bound_types = runner_conf.window_bound_types
self._return_schema = StructType(
- [StructField("_%d" % i, udf.return_type) for i, udf in
enumerate(udfs)]
+ [StructField("_%d" % i, udf.return_type) for i, udf in
enumerate(udf_descs)]
)
def run(self, split_index: int, data: "Iterator[GroupedBatch]") ->
"Iterator[pa.RecordBatch]":
@@ -659,7 +672,7 @@ class
PandasWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
num_rows = table.num_rows
result_series = []
- for udf_index, (udf, udf_func) in enumerate(zip(self._udfs,
self._funcs)):
+ for udf_index, (udf, udf_func) in enumerate(zip(self._udf_descs,
self._funcs)):
args_offsets, kwargs_offsets = udf.args_offsets,
udf.kwargs_offsets
bound_type = self._window_bound_types[udf_index]
if bound_type == "unbounded":
@@ -725,17 +738,20 @@ class
ArrowBatchedUDFLegacyPandasHandler(BatchEvalTypeHandler["pa.RecordBatch"])
eval_type = PythonEvalType.SQL_ARROW_BATCHED_UDF
variant_conf = "use_legacy_pandas_udf_conversion"
variant_value = True
+ _multi_udf = True
# run wraps the row-by-row evaluation in fail_on_stopiteration itself.
_fail_on_stopiteration = False
- def __init__(self, udfs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf") -> None:
+ def __init__(
+ self, udf_descs: list[UDF], runner_conf: "RunnerConf", eval_conf:
"EvalConf"
+ ) -> None:
require_minimum_pandas_version()
require_minimum_pyarrow_version()
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
# --- UDF preparation ---
self._udf_infos = []
- for udf, func in zip(udfs, self._funcs):
+ for udf, func in zip(udf_descs, self._funcs):
udf_return_type = udf.return_type
wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
func, udf.args_offsets, udf.kwargs_offsets
@@ -758,7 +774,7 @@ class
ArrowBatchedUDFLegacyPandasHandler(BatchEvalTypeHandler["pa.RecordBatch"])
coerce,
)
)
- col_names = [f"_{i}" for i in range(len(udfs))]
+ col_names = [f"_{i}" for i in range(len(udf_descs))]
self._return_schema = StructType(
[StructField(name, info[3]) for name, info in zip(col_names,
self._udf_infos)]
)
diff --git a/python/pyspark/eval_handlers/_stateful.py
b/python/pyspark/eval_handlers/_stateful.py
index 87f38371cf6c..56503ba9eef0 100644
--- a/python/pyspark/eval_handlers/_stateful.py
+++ b/python/pyspark/eval_handlers/_stateful.py
@@ -98,16 +98,14 @@ class
PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
]
)
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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_WITH_STATE UDF
expected here."
+ super().__init__(udf_descs, runner_conf, eval_conf)
# See FlatMapGroupsInPandasWithStateExec for how arg_offsets are used
to
# distinguish between grouping attributes and data attributes.
- parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._value_offsets = parsed_offsets[0][1]
@@ -227,13 +225,13 @@ class
PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
# The number of columns of the result must match the return type,
# but an empty result with no columns at all is acceptable.
if not (
- len(result.columns) == len(self._udfs[0].return_type)
+ len(result.columns) == len(self._udf_desc.return_type)
or (len(result.columns) == 0 and result.empty)
):
raise PySparkRuntimeError(
errorClass="RESULT_COLUMN_SCHEMA_MISMATCH",
messageParameters={
- "expected": str(len(self._udfs[0].return_type)),
+ "expected": str(len(self._udf_desc.return_type)),
"actual": str(len(result.columns)),
},
)
@@ -261,7 +259,7 @@ class
PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
else:
values = (pd.concat(x, axis=1) for x in value_series_gen)
- result_iter = self._funcs[0](key, values, state)
+ result_iter = self._func(key, values, state)
if isinstance(result_iter, pd.DataFrame):
raise PySparkTypeError(
@@ -394,7 +392,7 @@ class
PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
yield construct_record_batch(
pdfs,
pdf_data_cnt,
- self._udfs[0].return_type,
+ self._udf_desc.return_type,
state_pdfs,
state_data_cnt,
)
@@ -411,7 +409,7 @@ class
PandasGroupedMapWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]
# Flush the trailing batch if it has any data or state left.
if pdf_data_cnt > 0 or state_data_cnt > 0:
yield construct_record_batch(
- pdfs, pdf_data_cnt, self._udfs[0].return_type, state_pdfs,
state_data_cnt
+ pdfs, pdf_data_cnt, self._udf_desc.return_type, state_pdfs,
state_data_cnt
)
@@ -423,23 +421,21 @@ class
PandasTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"])
eval_type = PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_UDF
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
+ super().__init__(udf_descs, runner_conf, eval_conf)
# See TransformWithStateInPandasExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes
- parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.args_offsets)
assert len(parsed_offsets) == 1, (
"Expected one pair of offsets for TRANSFORM_WITH_STATE_PANDAS UDF."
)
self._key_offsets = parsed_offsets[0][0]
self._value_offsets = parsed_offsets[0][1]
- self._output_schema = StructType([StructField("_0",
udfs[0].return_type)])
+ self._output_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
self._stateful_processor_api_client = StatefulProcessorApiClient(
eval_conf.state_server_socket_port, # type: ignore[arg-type]
@@ -517,7 +513,7 @@ class
PandasTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"])
def convert_results(result_iter: Iterable[pd.DataFrame]) ->
Iterator[pa.RecordBatch]:
for result in result_iter:
- if isinstance(self._udfs[0].return_type, StructType) and not
isinstance(
+ if isinstance(self._udf_desc.return_type, StructType) and not
isinstance(
result, pd.DataFrame
):
raise PySparkValueError(
@@ -540,7 +536,7 @@ class
PandasTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"])
# This must be a generator expression - do not materialize.
values_gen = (df.iloc[:, self._value_offsets] for _, df in group)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_DATA,
key,
@@ -549,7 +545,7 @@ class
PandasTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"])
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_TIMER,
None,
@@ -558,7 +554,7 @@ class
PandasTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"])
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.COMPLETE,
None,
@@ -575,12 +571,10 @@ class
PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
eval_type = PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
+ super().__init__(udf_descs, runner_conf, eval_conf)
# See TransformWithStateInPandasExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes.
@@ -589,10 +583,10 @@ class
PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
# [groupingKeyOffsets, dedupDataOffsets],
# [initStateGroupingOffsets, dedupInitDataOffsets]
# ]
- parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._init_key_offsets = parsed_offsets[1][0]
- self._output_schema = StructType([StructField("_0",
udfs[0].return_type)])
+ self._output_schema = StructType([StructField("_0",
self._udf_desc.return_type)])
self._stateful_processor_api_client = StatefulProcessorApiClient(
eval_conf.state_server_socket_port, # type: ignore[arg-type]
@@ -729,7 +723,7 @@ class
PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
# TODO(SPARK-49100): add verification that elements in result_iter
are
# indeed of type pd.DataFrame and conform to assigned cols
for result in result_iter:
- if isinstance(self._udfs[0].return_type, StructType) and not
isinstance(
+ if isinstance(self._udf_desc.return_type, StructType) and not
isinstance(
result, pd.DataFrame
):
raise PySparkValueError(
@@ -756,7 +750,7 @@ class
PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
state_values = (data_df for _, data_df, _ in group_data if not
data_df.empty)
init_states = (init_df for _, _, init_df in group_init if not
init_df.empty)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_DATA,
key,
@@ -766,7 +760,7 @@ class
PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_TIMER,
None,
@@ -776,7 +770,7 @@ class
PandasTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.Recor
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.COMPLETE,
None,
@@ -794,15 +788,13 @@ class
RowTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
eval_type = PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
+ super().__init__(udf_descs, runner_conf, eval_conf)
# See TransformWithStateInPySparkExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes.
- parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._stateful_processor_api_client = StatefulProcessorApiClient(
@@ -850,7 +842,7 @@ class
RowTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
# Convert spark type to arrow type
# TODO: we need to make this configurable, currently using default
values.
arrow_type = to_arrow_type(
- self._udfs[0].return_type,
+ self._udf_desc.return_type,
timezone="UTC",
prefers_large_types=False,
)
@@ -865,7 +857,7 @@ class
RowTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
# This must be a generator expression - do not materialize.
values_gen = map(lambda x: x[1], group)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_DATA,
key,
@@ -874,7 +866,7 @@ class
RowTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_TIMER,
None,
@@ -883,7 +875,7 @@ class
RowTransformWithStateUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.COMPLETE,
None,
@@ -900,11 +892,9 @@ class
RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
eval_type =
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF
- def __init__(self, udfs: list[UDF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: 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."
+ super().__init__(udf_descs, runner_conf, eval_conf)
# See TransformWithStateInPandasExec for how arg_offsets are used to
# distinguish between grouping attributes and data attributes.
@@ -913,7 +903,7 @@ class
RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
# [groupingKeyOffsets, dedupDataOffsets],
# [initStateGroupingOffsets, dedupInitDataOffsets]
# ]
- parsed_offsets = extract_key_value_indexes(udfs[0].args_offsets)
+ parsed_offsets = extract_key_value_indexes(self._udf_desc.args_offsets)
self._key_offsets = parsed_offsets[0][0]
self._init_key_offsets = parsed_offsets[1][0]
@@ -1030,7 +1020,7 @@ class
RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
# Convert spark type to arrow type
# TODO: we need to make this configurable, currently using default
values.
arrow_type = to_arrow_type(
- self._udfs[0].return_type,
+ self._udf_desc.return_type,
timezone="UTC",
prefers_large_types=False,
)
@@ -1045,7 +1035,7 @@ class
RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
# These must be generator expressions - do not materialize.
for _, (values_gen, init_states_gen) in group:
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_DATA,
key,
@@ -1055,7 +1045,7 @@ class
RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.PROCESS_TIMER,
None,
@@ -1065,7 +1055,7 @@ class
RowTransformWithStateInitStateUDFHandler(BatchEvalTypeHandler["pa.RecordBa
)
yield from convert_results(
- self._funcs[0](
+ self._func(
self._stateful_processor_api_client,
TransformWithStateInPandasFuncMode.COMPLETE,
None,
diff --git a/python/pyspark/eval_handlers/_typing.py
b/python/pyspark/eval_handlers/_typing.py
index d01180bd5544..833ccd3ceb0a 100644
--- a/python/pyspark/eval_handlers/_typing.py
+++ b/python/pyspark/eval_handlers/_typing.py
@@ -45,10 +45,10 @@ OutputBatch = TypeVar("OutputBatch")
class UDF:
"""One UDF as the worker reads it; each handler takes the fields it
needs."""
- # The chained function before any wrapping: the pickled object itself (an
incremental
- # Aggregator) with its signature. EvalTypeHandler.__init__ keeps it
wrapped for profiling and,
+ # The functions applied in turn to the UDF's arguments, as unpickled.
EvalTypeHandler.__init__
+ # chains them into one as the handler's _raw_funcs, and keeps that wrapped
for profiling and,
# for most eval types, in fail_on_stopiteration, as the handler's _funcs.
- raw_func: Any
+ funcs: list[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
diff --git a/python/pyspark/eval_handlers/_udtf.py
b/python/pyspark/eval_handlers/_udtf.py
index 2b1b67b348e2..e179f4b68a1a 100644
--- a/python/pyspark/eval_handlers/_udtf.py
+++ b/python/pyspark/eval_handlers/_udtf.py
@@ -599,20 +599,23 @@ def build_null_checker(return_type: StructType) ->
Optional[Callable[[Any], None
class UDTFHandlerBase(EvalTypeHandler[InputBatch, OutputBatch]):
- """Base for every UDTF eval type handler; ``__init__`` creates the UDTF
instance as
- ``_udtf``."""
+ """Base for every UDTF eval type handler; ``__init__`` keeps the UDTF
record as
+ ``_udtf_desc`` and creates the UDTF instance as ``_udtf``."""
# Wraps the UDTF when its TABLE argument has PARTITION BY expressions.
partition_wrapper: ClassVar[type[PartitionedUDTF]] = UDTFWithPartitions
- def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def _expected_udf(self) -> str:
+ return "UDTF"
+
+ def __init__(self, udf_descs: list[UDTF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
"""Create, as ``_udtf``, the instance of the UDTF's class (wrapped in
``partition_wrapper`` when its TABLE argument has PARTITION BY
expressions, its
constructor given the AnalyzeResult when it takes one), and check its
'eval' signature
against the call arguments."""
- super().__init__(udfs, runner_conf, eval_conf)
- assert len(udfs) == 1, "One UDTF expected here."
- udtf = udfs[0]
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._udtf_desc: UDTF = udf_descs[0]
+ udtf = self._udtf_desc
if not isinstance(udtf.udtf_cls, type):
raise PySparkRuntimeError(
f"Invalid UDTF handler type. Expected a class (type 'type'),
but "
@@ -685,11 +688,10 @@ class UDTFEvalTypeHandler(UDTFHandlerBase[InputBatch,
OutputBatch]):
``cleanup``. Subclasses implement the two hooks and pick the serializer.
"""
- 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."
+ def __init__(self, udf_descs: list[UDTF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
+ super().__init__(udf_descs, runner_conf, eval_conf)
self._eval_method, self._args_kwargs_offsets = wrap_kwargs_support(
- self._udtf.eval, udfs[0].args_offsets, udfs[0].kwargs_offsets
+ self._udtf.eval, self._udtf_desc.args_offsets,
self._udtf_desc.kwargs_offsets
)
self._terminate_method = getattr(self._udtf, "terminate", None)
self._cleanup_method = getattr(self._udtf, "cleanup", None)
@@ -725,19 +727,20 @@ class TableUDTFHandler(UDTFHandlerBase[Any, tuple]):
eval_type = PythonEvalType.SQL_TABLE_UDF
- 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."
- self._check_output_row_against_schema =
build_null_checker(udfs[0].return_type)
+ def __init__(self, udf_descs: list[UDTF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._check_output_row_against_schema =
build_null_checker(self._udtf_desc.return_type)
eval_func_kwargs_support, self._args_kwargs_offsets =
wrap_kwargs_support(
- self._udtf.eval, udfs[0].args_offsets, udfs[0].kwargs_offsets
+ self._udtf.eval, self._udtf_desc.args_offsets,
self._udtf_desc.kwargs_offsets
)
- self._eval = self._wrap_udtf(eval_func_kwargs_support,
udfs[0].return_type)
+ self._eval = self._wrap_udtf(eval_func_kwargs_support,
self._udtf_desc.return_type)
self._terminate: Optional[Callable[..., tuple]]
if hasattr(self._udtf, "terminate"):
- self._terminate = self._wrap_udtf(getattr(self._udtf,
"terminate"), udfs[0].return_type)
+ self._terminate = self._wrap_udtf(
+ getattr(self._udtf, "terminate"), self._udtf_desc.return_type
+ )
else:
self._terminate = None
@@ -850,16 +853,16 @@ class ArrowUDTFHandler(
eval_type = PythonEvalType.SQL_ARROW_UDTF
partition_wrapper = ArrowUDTFWithPartition
- def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDTF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
import pyarrow as pa
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
arrow_return_type = to_arrow_type(
- self._udfs[0].return_type,
+ self._udtf_desc.return_type,
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
- self._return_type_size = len(self._udfs[0].return_type)
+ self._return_type_size = len(self._udtf_desc.return_type)
self._target_schema = pa.schema(list(arrow_return_type))
self._table_arg_offsets = (
@@ -962,10 +965,10 @@ class ArrowTableUDTFHandler(
``spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled``.
Subclasses implement
``_to_columns`` and ``_convert_result``."""
- 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._udfs[0].return_type)
- self._check_output_row_against_schema =
build_null_checker(self._udfs[0].return_type)
+ def __init__(self, udf_descs: list[UDTF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._return_type_size = len(self._udtf_desc.return_type)
+ self._check_output_row_against_schema =
build_null_checker(self._udtf_desc.return_type)
@abstractmethod
def _to_columns(self, batch: pa.RecordBatch) -> list:
@@ -1014,13 +1017,13 @@ class
ArrowTableUDTFLegacyPandasHandler(ArrowTableUDTFHandler):
variant_conf = "use_legacy_pandas_udtf_conversion"
variant_value = True
- def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDTF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
import pandas # noqa: F401
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
# The output pandas DataFrame is converted as a single struct column
named
# "_0" against this schema.
- self._output_schema = StructType([StructField("_0",
self._udfs[0].return_type)])
+ self._output_schema = StructType([StructField("_0",
self._udtf_desc.return_type)])
def _verify_result(self, result: Any, method_name: str) -> Any:
import pandas as pd
@@ -1053,7 +1056,7 @@ class
ArrowTableUDTFLegacyPandasHandler(ArrowTableUDTFHandler):
# Verify the type and the schema of the result.
verify_pandas_result(
result,
- self._udfs[0].return_type,
+ self._udtf_desc.return_type,
assign_cols_by_name=False,
truncate_return_schema=False,
)
@@ -1127,12 +1130,12 @@ class ArrowTableUDTFArrowHandler(ArrowTableUDTFHandler):
variant_conf = "use_legacy_pandas_udtf_conversion"
variant_value = False
- def __init__(self, udfs: list[UDTF], runner_conf: RunnerConf, eval_conf:
EvalConf) -> None:
+ def __init__(self, udf_descs: list[UDTF], runner_conf: RunnerConf,
eval_conf: EvalConf) -> None:
import pyarrow # noqa: F401
- super().__init__(udfs, runner_conf, eval_conf)
+ super().__init__(udf_descs, runner_conf, eval_conf)
self._arrow_return_type = to_arrow_type(
- self._udfs[0].return_type,
+ self._udtf_desc.return_type,
timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types,
)
@@ -1217,14 +1220,14 @@ class ArrowTableUDTFArrowHandler(ArrowTableUDTFHandler):
errorClass="UDTF_ARROW_DATA_CONVERSION_ERROR",
messageParameters={
"data": str(data),
- "schema": self._udfs[0].return_type.simpleString(),
+ "schema": self._udtf_desc.return_type.simpleString(),
"arrow_schema": str(arrow_return_type),
},
) from original_exception
try:
table = LocalDataToArrowConversion.convert(
- data, self._udfs[0].return_type,
self._runner_conf.use_large_var_types
+ data, self._udtf_desc.return_type,
self._runner_conf.use_large_var_types
)
except PySparkValueError as e:
if e.getErrorClass() == "AXIS_LENGTH_MISMATCH":
diff --git a/python/pyspark/eval_handlers/_vanilla.py
b/python/pyspark/eval_handlers/_vanilla.py
index b20da0f2c839..a76a0cae6e45 100644
--- a/python/pyspark/eval_handlers/_vanilla.py
+++ b/python/pyspark/eval_handlers/_vanilla.py
@@ -42,10 +42,15 @@ class BatchedUDFHandler(EvalTypeHandler[Any, Any]):
row."""
eval_type = PythonEvalType.SQL_BATCHED_UDF
+ _multi_udf = True
- 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, func) for udf, func in
zip(udfs, self._funcs)]
+ def __init__(
+ self, udf_descs: "list[UDF]", runner_conf: "RunnerConf", eval_conf:
"EvalConf"
+ ) -> None:
+ super().__init__(udf_descs, runner_conf, eval_conf)
+ self._udf_infos = [
+ self._prepare_udf(udf, func) for udf, func in zip(udf_descs,
self._funcs)
+ ]
@staticmethod
def _prepare_udf(udf: "UDF", func: Callable[..., Any]) -> tuple[list[int],
Callable[..., Any]]:
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 1c29e79b98d4..834fa7599874 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
@@ -128,7 +128,7 @@ class ArrowEvalTypeHandlerTestCase(PySparkBaseTestCase):
def _udf(func, args_offsets, kwargs_offsets=None, return_type=None):
"""A ``UDF`` record as the worker builds it."""
return UDF(
- raw_func=func,
+ funcs=[func],
args_offsets=args_offsets,
kwargs_offsets=kwargs_offsets or {},
return_type=return_type or LongType(),
@@ -138,10 +138,9 @@ class ArrowEvalTypeHandlerTestCase(PySparkBaseTestCase):
@classmethod
def _grouped_handler(cls, handler_cls, udf, arg_offsets):
"""Build a grouped/cogrouped-map handler with the shared return type.
The function the
- handler wraps loses its signature, so the handler reads it from the
record's
- ``raw_func``."""
+ handler wraps loses its signature, so the handler reads it from its
``_raw_func``."""
return handler_cls(
- udfs=[cls._udf(udf, arg_offsets, return_type=cls._RETURN_TYPE)],
+ udf_descs=[cls._udf(udf, arg_offsets,
return_type=cls._RETURN_TYPE)],
runner_conf=RunnerConf({}),
eval_conf=None,
)
@@ -192,10 +191,10 @@ class
ArrowEvalTypeHandlerRegistrationTests(PySparkBaseTestCase):
class ArrowBatchedUDFArrowHandlerTests(ArrowEvalTypeHandlerTestCase):
_INPUT_TYPE = StructType([StructField("a", LongType()), StructField("b",
LongType())])
- def _handler(self, *udfs, runner_conf=None, input_type=None):
+ def _handler(self, *udf_descs, runner_conf=None, input_type=None):
"""Build the handler from ``(func, args_offsets, kwargs_offsets,
return_type)`` tuples."""
return ArrowBatchedUDFArrowHandler(
- udfs=[self._udf(*u) for u in udfs],
+ udf_descs=[self._udf(*u) for u in udf_descs],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=EvalConf({"input_type": (input_type or
self._INPUT_TYPE).json()}),
)
@@ -383,7 +382,7 @@ class
ArrowScalarUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return pa.array([v.as_py() + 1 for v in col], type=pa.int64())
handler = ArrowScalarUDFHandler(
- udfs=[self._udf(add_one, [0])], runner_conf=RunnerConf({}),
eval_conf=None
+ udf_descs=[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]])
@@ -394,7 +393,7 @@ class
ArrowScalarUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return pa.array([v.as_py() + 1 for v in col], type=pa.int32())
handler = ArrowScalarUDFHandler(
- udfs=[self._udf(add_one, [0])], runner_conf=RunnerConf({}),
eval_conf=None
+ udf_descs=[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())
@@ -409,7 +408,7 @@ class
ArrowScalarIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
yield pa.array([v.as_py() + 1 for v in col], type=pa.int64())
handler = ArrowScalarIterUDFHandler(
- udfs=[self._udf(add_one, [0])], runner_conf=RunnerConf({}),
eval_conf=None
+ udf_descs=[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]])
@@ -421,7 +420,7 @@ class
ArrowScalarIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
yield pa.array(list(range(len(col) + 1)), type=pa.int64())
handler = ArrowScalarIterUDFHandler(
- udfs=[self._udf(too_many, [0])], runner_conf=RunnerConf({}),
eval_conf=None
+ udf_descs=[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])])))
@@ -436,7 +435,7 @@ class ArrowMapUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
yield self._batch(v=[c.as_py() * 2 for c in batch.column("v")])
handler = ArrowMapUDFHandler(
- udfs=[self._udf(double_v, [])], runner_conf=RunnerConf({}),
eval_conf=None
+ udf_descs=[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])
@@ -546,7 +545,7 @@ class
ArrowGroupedAggUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return sum(c.as_py() for c in col)
handler = ArrowGroupedAggUDFHandler(
- udfs=[self._udf(sum_udf, [0])], runner_conf=RunnerConf({}),
eval_conf=None
+ udf_descs=[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])
@@ -560,7 +559,7 @@ class
ArrowGroupedAggIterUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return sum(c.as_py() for col in col_iter for c in col)
handler = ArrowGroupedAggIterUDFHandler(
- udfs=[self._udf(sum_iter_udf, [0])], runner_conf=RunnerConf({}),
eval_conf=None
+ udf_descs=[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])
@@ -574,7 +573,7 @@ class
ArrowWindowAggUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return sum(c.as_py() for c in col)
handler = ArrowWindowAggUDFHandler(
- udfs=[self._udf(sum_udf, [0])],
+ udf_descs=[self._udf(sum_udf, [0])],
runner_conf=RunnerConf({"window_bound_types": "unbounded"}),
eval_conf=None,
)
@@ -587,7 +586,7 @@ class
ArrowWindowAggUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
return sum(c.as_py() for c in col)
handler = ArrowWindowAggUDFHandler(
- udfs=[self._udf(sum_udf, [0, 1, 2])],
+ udf_descs=[self._udf(sum_udf, [0, 1, 2])],
runner_conf=RunnerConf({"window_bound_types": "bounded"}),
eval_conf=None,
)
@@ -605,9 +604,9 @@ class
ArrowGroupedAggIncrementalPartialUDFHandlerTests(ArrowEvalTypeHandlerTestC
_KEY_SCHEMA = StructType([StructField("k", LongType())])
def _handler(self, cap=10000):
- # The handler takes the Aggregator itself from the record's raw_func.
+ # The handler takes the Aggregator itself, unchained and unwrapped, as
its raw function.
return ArrowGroupedAggIncrementalPartialUDFHandler(
- udfs=[self._udf(_SumAggregator(), [1])],
+ udf_descs=[self._udf(_SumAggregator(), [1])],
runner_conf=RunnerConf({"spark.sql.execution.arrow.maxRecordsPerBatch":
str(cap)}),
eval_conf=EvalConf({"grouping_key_schema":
self._KEY_SCHEMA.json()}),
)
@@ -665,7 +664,7 @@ class
ArrowGroupedAggIncrementalPartialUDFHandlerTests(ArrowEvalTypeHandlerTestC
["k", "v"],
)
handler = ArrowGroupedAggIncrementalPartialUDFHandler(
- udfs=[self._udf(_SumAggregator(), [1])],
+ udf_descs=[self._udf(_SumAggregator(), [1])],
runner_conf=RunnerConf({}),
eval_conf=EvalConf(
{"grouping_key_schema": StructType([StructField("k",
DoubleType())]).json()}
@@ -686,9 +685,9 @@ class
ArrowGroupedAggIncrementalFinalUDFHandlerTests(ArrowEvalTypeHandlerTestCas
return pa.RecordBatch.from_arrays([pa.array(buffers,
type=self._BUFFER_TYPE)], ["buf"])
def _handler(self):
- # The handler takes the Aggregator itself from the record's raw_func.
+ # The handler takes the Aggregator itself, unchained and unwrapped, as
its raw function.
return ArrowGroupedAggIncrementalFinalUDFHandler(
- udfs=[self._udf(_SumAggregator(), [0])],
runner_conf=RunnerConf({}), eval_conf=None
+ udf_descs=[self._udf(_SumAggregator(), [0])],
runner_conf=RunnerConf({}), eval_conf=None
)
def test_merges_partial_buffers_per_group(self):
@@ -713,7 +712,7 @@ class
ArrowGroupedAggIncrementalFinalUDFHandlerTests(ArrowEvalTypeHandlerTestCas
class ArrowWindowAggIncrementalUDFHandlerTests(ArrowEvalTypeHandlerTestCase):
def _handler(self, agg, bound_type, args_offsets):
return ArrowWindowAggIncrementalUDFHandler(
- udfs=[self._udf(agg, args_offsets)],
+ udf_descs=[self._udf(agg, args_offsets)],
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 baff8a51e6d7..b3e0438569ce 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,6 +22,7 @@ Handler-flavor tests live alongside their module, e.g. the
Arrow handlers in
"""
import os
+import unittest
from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError
@@ -34,19 +35,23 @@ from pyspark.eval_handlers._base import (
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, StringType
-from pyspark.testing.utils import PySparkBaseTestCase
+from pyspark.sql.types import LongType
+from pyspark.testing.utils import (
+ PySparkBaseTestCase,
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+)
from pyspark.util import PythonEvalType
with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
- from pyspark.worker import unpickle_and_chain
- from pyspark.worker_util import RunnerConf, pickleSer
+ from pyspark.worker_util import RunnerConf
class EvalTypeHandlerTests(PySparkBaseTestCase):
@@ -77,14 +82,20 @@ class EvalTypeHandlerTests(PySparkBaseTestCase):
def test_default_serializer_per_category(self):
class _Batch(BatchEvalTypeHandler["pa.RecordBatch"]):
+ _multi_udf = True
+
def run(self, split_index, data):
return data
class _Grouped(GroupedEvalTypeHandler["pa.RecordBatch"]):
+ _multi_udf = True
+
def run(self, split_index, data):
return data
class _CoGrouped(CoGroupedEvalTypeHandler["pa.RecordBatch"]):
+ _multi_udf = True
+
def run(self, split_index, data):
return data
@@ -98,6 +109,8 @@ class EvalTypeHandlerTests(PySparkBaseTestCase):
def test_run_produces_output(self):
class _Doubler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ _multi_udf = True
+
def run(self, split_index, data):
for item in data:
yield item * 2
@@ -272,6 +285,8 @@ class EvalTypeHandlerTests(PySparkBaseTestCase):
class HandlerUDFsTests(PySparkBaseTestCase):
class _Handler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ _multi_udf = True
+
def run(self, split_index, data):
return data
@@ -279,35 +294,34 @@ class HandlerUDFsTests(PySparkBaseTestCase):
def _stop(x):
raise StopIteration(x)
- def _new(self, handler_cls, runner_conf=None):
- return handler_cls(
- [
- UDF(
- raw_func=self._stop,
- args_offsets=[1],
- kwargs_offsets={"y": 2},
- return_type=LongType(),
- result_id=0,
- )
- ],
- runner_conf or RunnerConf({}),
- None,
+ @staticmethod
+ def _udf(*funcs):
+ return UDF(
+ funcs=list(funcs),
+ args_offsets=[1],
+ kwargs_offsets={"y": 2},
+ return_type=LongType(),
+ result_id=0,
)
+ def _new(self, handler_cls, runner_conf=None):
+ return handler_cls([self._udf(self._stop)], runner_conf or
RunnerConf({}), None)
+
def test_wraps_udfs(self):
handler = self._new(self._Handler)
- udf = handler._udfs[0]
+ udf = handler._udf_descs[0]
self.assertIsInstance(udf, UDF)
self.assertEqual(udf.args_offsets, [1])
self.assertEqual(udf.kwargs_offsets, {"y": 2})
self.assertEqual(udf.return_type, LongType())
- # The handler keeps the function wrapped in fail_on_stopiteration; the
record keeps it
- # unwrapped.
+ # The handler keeps the function wrapped in fail_on_stopiteration; the
record and the
+ # handler's raw function keep it unwrapped.
with self.assertRaises(PySparkRuntimeError):
handler._funcs[0](1)
with self.assertRaises(StopIteration):
- udf.raw_func(1)
+ udf.funcs[0](1)
+ self.assertIs(handler._raw_funcs[0], self._stop)
def test_fail_on_stopiteration_per_handler(self):
not_wrapped = {PythonEvalType.SQL_ARROW_BATCHED_UDF}
@@ -370,19 +384,87 @@ class HandlerUDFsTests(PySparkBaseTestCase):
handler._funcs[0](1)
self.assertEqual(cm.exception.getCondition(),
"STOP_ITERATION_OCCURRED")
-
-class UnpickleAndChainTests(PySparkBaseTestCase):
- def test_chains_in_order(self):
- ser = CloudPickleSerializer()
- func, return_type = unpickle_and_chain(
- pickleSer,
- [
- ser.dumps((lambda x: x + 1, LongType())),
- ser.dumps((lambda x: str(x * 2), StringType())),
- ],
+ def test_chains_funcs_in_order(self):
+ handler = self._Handler(
+ [self._udf(lambda x: x + 1, lambda x: x * 2, str)],
RunnerConf({}), None
)
- self.assertEqual(func(3), "8")
- self.assertEqual(return_type, StringType())
+ self.assertEqual(handler._raw_funcs[0](3), "8")
+ self.assertEqual(handler._funcs[0](3), "8")
+
+ def test_single_func_not_chained(self):
+ # A single function is kept as is, so its signature (and an
incremental Aggregator
+ # object) survives.
+ def f(key, value):
+ return value
+
+ handler = self._Handler([self._udf(f)], RunnerConf({}), None)
+ self.assertIs(handler._raw_funcs[0], f)
+
+ def test_single_udf(self):
+ class _Single(BatchEvalTypeHandler["pa.RecordBatch"]):
+ def run(self, split_index, data):
+ return data
+
+ def f(x):
+ return x + 1
+
+ handler = _Single([self._udf(f)], RunnerConf({}), None)
+ self.assertIs(handler._udf_desc, handler._udf_descs[0])
+ self.assertIs(handler._raw_func, f)
+ self.assertIs(handler._func, handler._funcs[0])
+ self.assertEqual(handler._func(1), 2)
+ for udf_descs in ([], [self._udf(f), self._udf(f)]):
+ with self.subTest(num_udfs=len(udf_descs)):
+ with self.assertRaisesRegex(
+ PySparkRuntimeError,
+ "^One _Single UDF expected, but got
{}.$".format(len(udf_descs)),
+ ):
+ _Single(udf_descs, RunnerConf({}), None)
+
+ handler = self._Handler([self._udf(f), self._udf(f)], RunnerConf({}),
None)
+ self.assertFalse(hasattr(handler, "_udf_desc"))
+
+ @unittest.skipIf(
+ not (have_pandas and have_pyarrow),
+ pandas_requirement_message or pyarrow_requirement_message,
+ )
+ def test_single_udf_count_error_names_eval_type(self):
+ for eval_type, expected in (
+ (PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF, "COGROUPED_MAP_ARROW
UDF"),
+ (
+ PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE,
+ "GROUPED_MAP_PANDAS_UDF_WITH_STATE UDF",
+ ),
+ (PythonEvalType.SQL_TABLE_UDF, "UDTF"),
+ ):
+ handler_cls = _eval_type_handlers[eval_type][1][None]
+ with self.subTest(handler=handler_cls.__name__):
+ with self.assertRaisesRegex(
+ PySparkRuntimeError, "^One {} expected, but got
2.$".format(expected)
+ ):
+ handler_cls([self._udf(str), self._udf(str)],
RunnerConf({}), None)
+
+ def test_multi_udf_per_handler(self):
+ multi = {
+ PythonEvalType.SQL_BATCHED_UDF,
+ PythonEvalType.SQL_ARROW_BATCHED_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_UDF,
+ PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
+ }
+ for _, variants in list(_eval_type_handlers.values()):
+ for handler_cls in variants.values():
+ with self.subTest(handler=handler_cls.__name__):
+ self.assertEqual(handler_cls._multi_udf,
handler_cls.eval_type in multi)
if __name__ == "__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 69a2301e7a64..0542478ff7f0 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
@@ -93,22 +93,22 @@ class
ElementwiseEvalTypeHandlerTestCase(PySparkBaseTestCase):
["a", "b", "c"],
)
- def _handler(self, handler_cls, *udfs, nesting=None, input_type=None,
runner_conf=None):
+ def _handler(self, handler_cls, *udf_descs, nesting=None, input_type=None,
runner_conf=None):
"""Build ``handler_cls`` from ``(func, args_offsets, kwargs_offsets,
return_type)``
tuples, with ``nesting`` as the per-UDF ``elementwise_nesting``."""
eval_conf = {"input_type": (input_type or self._INPUT_TYPE).json()}
if nesting is not None:
eval_conf["elementwise_nesting"] = ",".join(map(str, nesting))
return handler_cls(
- udfs=[
+ udf_descs=[
UDF(
- raw_func=func,
+ funcs=[func],
args_offsets=args_offsets,
kwargs_offsets=kwargs_offsets,
return_type=return_type,
result_id=0,
)
- for func, args_offsets, kwargs_offsets, return_type in udfs
+ for func, args_offsets, kwargs_offsets, return_type in
udf_descs
],
runner_conf=RunnerConf(runner_conf or
{"spark.sql.session.timeZone": "UTC"}),
eval_conf=EvalConf(eval_conf),
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 72c101ff901f..d95be4788de6 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
@@ -93,11 +93,11 @@ class PandasEvalTypeHandlerTestCase(PySparkBaseTestCase):
@classmethod
def _grouped_handler(cls, handler_cls, udf, arg_offsets, runner_conf=None):
"""Build a grouped/cogrouped-map handler. The function the handler
wraps loses its
- signature, so the handler reads it from the record's ``raw_func``."""
+ signature, so the handler reads it from its ``_raw_func``."""
return handler_cls(
- udfs=[
+ udf_descs=[
UDF(
- raw_func=udf,
+ funcs=[udf],
args_offsets=arg_offsets,
kwargs_offsets={},
return_type=cls._GROUPED_RETURN_TYPE,
@@ -117,7 +117,7 @@ class PandasEvalTypeHandlerTestCase(PySparkBaseTestCase):
def _udf(func, return_type=None, args=(0,), kwargs=None):
"""A ``UDF`` record for a pandas UDF, as the worker builds it."""
return UDF(
- raw_func=func,
+ funcs=[func],
args_offsets=list(args),
kwargs_offsets=dict(kwargs or {}),
return_type=return_type or LongType(),
@@ -125,10 +125,10 @@ class PandasEvalTypeHandlerTestCase(PySparkBaseTestCase):
)
@staticmethod
- def _agg_handler(handler_cls, *udfs, runner_conf=None):
+ def _agg_handler(handler_cls, *udf_descs, runner_conf=None):
"""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
+ udf_descs=list(udf_descs), runner_conf=runner_conf or
RunnerConf({}), eval_conf=None
)
@staticmethod
@@ -172,10 +172,10 @@ class
PandasEvalTypeHandlerRegistrationTests(PySparkBaseTestCase):
class ArrowBatchedUDFLegacyPandasHandlerTests(PandasEvalTypeHandlerTestCase):
_INPUT_TYPE = StructType([StructField("a", LongType()), StructField("b",
LongType())])
- def _handler(self, *udfs, runner_conf=None, input_type=None):
+ def _handler(self, *udf_descs, runner_conf=None, input_type=None):
"""Build the handler from one or more ``_udf`` records."""
return ArrowBatchedUDFLegacyPandasHandler(
- udfs=list(udfs),
+ udf_descs=list(udf_descs),
runner_conf=runner_conf or RunnerConf({}),
eval_conf=EvalConf({"input_type": (input_type or
self._INPUT_TYPE).json()}),
)
@@ -232,10 +232,10 @@ class
ArrowBatchedUDFLegacyPandasHandlerTests(PandasEvalTypeHandlerTestCase):
@unittest.skipIf(not (have_pandas and have_pyarrow), _missing_message)
class PandasScalarUDFHandlerTests(PandasEvalTypeHandlerTestCase):
@staticmethod
- def _handler(*udfs, runner_conf=None):
+ def _handler(*udf_descs, runner_conf=None):
"""Build a PandasScalarUDFHandler from one or more ``_udf`` records."""
return PandasScalarUDFHandler(
- udfs=list(udfs), runner_conf=runner_conf or RunnerConf({}),
eval_conf=None
+ udf_descs=list(udf_descs), runner_conf=runner_conf or
RunnerConf({}), eval_conf=None
)
def test_invokes_udf_per_batch(self):
@@ -320,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=[PandasEvalTypeHandlerTestCase._udf(udf, return_type, args)],
+ udf_descs=[PandasEvalTypeHandlerTestCase._udf(udf, return_type,
args)],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=None,
)
@@ -413,7 +413,7 @@ class
PandasMapUDFHandlerTests(PandasEvalTypeHandlerTestCase):
def _handler(udf, return_type, runner_conf=None):
"""Build a PandasMapUDFHandler; map UDFs take no arg offsets."""
return PandasMapUDFHandler(
- udfs=[PandasEvalTypeHandlerTestCase._udf(udf, return_type,
args=())],
+ udf_descs=[PandasEvalTypeHandlerTestCase._udf(udf, return_type,
args=())],
runner_conf=runner_conf or RunnerConf({}),
eval_conf=None,
)
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 678561e64d77..ad8160cec554 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
@@ -117,9 +117,9 @@ class
TransformWithStateHandlerTestCase(PySparkBaseTestCase):
@classmethod
def _handler(cls, handler_cls, udf, arg_offsets, runner_conf=None):
return handler_cls(
- udfs=[
+ udf_descs=[
UDF(
- raw_func=udf,
+ funcs=[udf],
args_offsets=arg_offsets,
kwargs_offsets={},
return_type=cls._RETURN_TYPE,
@@ -235,9 +235,9 @@ class
PandasGroupedMapWithStateUDFHandlerTests(PySparkBaseTestCase):
@classmethod
def _handler(cls, udf, runner_conf=None):
return PandasGroupedMapWithStateUDFHandler(
- udfs=[
+ udf_descs=[
UDF(
- raw_func=udf,
+ funcs=[udf],
args_offsets=cls._ARG_OFFSETS,
kwargs_offsets={},
return_type=cls._RETURN_TYPE,
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 7d9760921418..f4c165bf9b3d 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
@@ -162,7 +162,9 @@ class TableUDTFHandlerTests(UDTFEvalTypeHandlerTestCase):
def _handler(self, udtf_cls, args=(0,), kwargs=None, return_type=None):
udtf = self._udtf_record(udtf_cls, return_type or
self._ROW_RETURN_TYPE, args, kwargs)
- return TableUDTFHandler(udfs=[udtf], runner_conf=RunnerConf({}),
eval_conf=EvalConf({}))
+ return TableUDTFHandler(
+ udf_descs=[udtf], runner_conf=RunnerConf({}),
eval_conf=EvalConf({})
+ )
def _run(self, handler, *rows):
# Each input row is a list of argument values; each output is the
tuple of rows one
@@ -271,7 +273,9 @@ class ArrowUDTFHandlerTests(UDTFEvalTypeHandlerTestCase):
):
conf = {} if table_arg_offsets is None else {"table_arg_offsets":
table_arg_offsets}
udtf = self._udtf_record(udtf_cls, self._RETURN_TYPE, args, kwargs,
partition_child_indexes)
- return ArrowUDTFHandler(udfs=[udtf], runner_conf=RunnerConf({}),
eval_conf=EvalConf(conf))
+ return ArrowUDTFHandler(
+ udf_descs=[udtf], runner_conf=RunnerConf({}),
eval_conf=EvalConf(conf)
+ )
def test_serializer(self):
self.assertIsInstance(self._handler(_RecordingUDTF).serializer,
ArrowStreamSerializer)
@@ -439,7 +443,7 @@ class ArrowTableUDTFHandlerTestsMixin:
def _handler(self, udtf_cls, args=(0,), kwargs=None, return_type=None):
udtf = self._udtf_record(udtf_cls, return_type or
self._ROW_RETURN_TYPE, args, kwargs)
return self.handler_cls(
- udfs=[udtf],
+ udf_descs=[udtf],
runner_conf=RunnerConf({}),
eval_conf=EvalConf({"input_type": self._ROW_INPUT_TYPE.json()}),
)
@@ -563,8 +567,10 @@ class UDTFHandlerInitTests(UDTFEvalTypeHandlerTestCase):
handler = self._new(TableUDTFHandler, Udtf)
self.assertIsInstance(handler._udtf, Udtf)
- udtf = handler._udfs[0]
+ udtf = handler._udf_descs[0]
self.assertIsInstance(udtf, UDTF)
+ self.assertIs(handler._udtf_desc, udtf)
+ self.assertFalse(hasattr(handler, "_udf_desc"))
self.assertEqual(udtf.args_offsets, [0])
self.assertEqual(udtf.kwargs_offsets, {})
self.assertEqual(udtf.return_type, self._RETURN_TYPE)
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 fe6aafb75143..def01b7ca8a2 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
@@ -40,18 +40,18 @@ with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME":
"PYTHON_WORKER"}):
class BatchedUDFHandlerTests(PySparkBaseTestCase):
@staticmethod
- def _handler(*udfs, kwargs_offsets=None, return_type=None):
+ def _handler(*udf_descs, kwargs_offsets=None, return_type=None):
"""Build the handler from ``(arg_offsets, func)`` pairs."""
return BatchedUDFHandler(
- udfs=[
+ udf_descs=[
UDF(
- raw_func=f,
+ funcs=[f],
args_offsets=arg_offsets,
kwargs_offsets=kwargs_offsets or {},
return_type=return_type or LongType(),
result_id=0,
)
- for arg_offsets, f in udfs
+ for arg_offsets, f in udf_descs
],
runner_conf=RunnerConf({}),
eval_conf=EvalConf({}),
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 563b4d9dfe3d..dc20ee355c7d 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -27,7 +27,7 @@ import json
import os
import sys
import warnings
-from typing import Any, BinaryIO, Callable
+from typing import BinaryIO
from pyspark import shuffle
from pyspark.accumulators import _accumulatorRegistry
@@ -41,12 +41,11 @@ from pyspark.messages import (
SparkSocketMessageReceiver,
)
from pyspark.serializers import (
- FramedSerializer,
SpecialLengths,
write_int,
write_with_length,
)
-from pyspark.sql.types import DataType, StructType, _parse_datatype_json_string
+from pyspark.sql.types import StructType, _parse_datatype_json_string
from pyspark.taskcontext import BarrierTaskContext, TaskContext
from pyspark.util import (
PythonEvalType,
@@ -77,23 +76,6 @@ def report_metrics(outfile: BinaryIO, metrics: dict[str,
int]) -> None:
write_with_length(payload, outfile)
-def chain(f: Callable[..., Any], g: Callable[..., Any]) -> Callable[..., Any]:
- """chain two functions together"""
- return lambda *a: g(f(*a))
-
-
-def unpickle_and_chain(
- pickleSer: FramedSerializer, pickled_udfs: list[memoryview]
-) -> tuple[Callable[..., Any], DataType]:
- """Unpickle the functions applied in turn to one column and chain them
into one; returns the
- chained function and the return type of the last one."""
- chained_func, return_type = read_command(pickleSer, pickled_udfs[0])
- for udf in pickled_udfs[1:]:
- f, return_type = read_command(pickleSer, udf)
- chained_func = chain(chained_func, f)
- return chained_func, return_type
-
-
def invoke_udf(message_receiver: SparkMessageReceiver, outfile: BinaryIO):
"""
This function is the main processing function for worker.py.
@@ -170,22 +152,28 @@ def invoke_udf(message_receiver: SparkMessageReceiver,
outfile: BinaryIO):
analyze_result=analyze_result,
name=init_info.udf_info.name,
)
- handler = handler_cls(udfs=[udtf], runner_conf=runner_conf,
eval_conf=eval_conf)
+ handler = handler_cls(
+ udf_descs=[udtf], runner_conf=runner_conf,
eval_conf=eval_conf
+ )
else:
assert isinstance(init_info.udf_info, list)
- udfs = []
+ udf_descs = []
for udf_info in init_info.udf_info:
- chained_func, return_type = unpickle_and_chain(pickleSer,
udf_info.udfs)
- udfs.append(
+ # The functions applied in turn to one column; the last
one's return type
+ # is the UDF's.
+ funcs, return_types = zip(*(read_command(pickleSer, c) for
c in udf_info.udfs))
+ udf_descs.append(
UDF(
- raw_func=chained_func,
+ funcs=list(funcs),
args_offsets=udf_info.args,
kwargs_offsets=udf_info.kwargs,
- return_type=return_type,
+ return_type=return_types[-1],
result_id=udf_info.result_id,
)
)
- handler = handler_cls(udfs=udfs, runner_conf=runner_conf,
eval_conf=eval_conf)
+ handler = handler_cls(
+ udf_descs=udf_descs, runner_conf=runner_conf,
eval_conf=eval_conf
+ )
func, serializer = handler.run, handler.serializer
deserializer = serializer
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
