You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/10/14 07:27:53 UTC
(spark) branch branch-4.x updated: [SPARK-60218][PYTHON] Migrate SQL_TABLE_UDF to the EvalTypeHandler pipeline and remove read_udtf
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 0c7ffe0aa812 [SPARK-60218][PYTHON] Migrate SQL_TABLE_UDF to the EvalTypeHandler pipeline and remove read_udtf
0c7ffe0aa812 is described below
commit 0c7ffe0aa812f2ad3a8bec2cbcefbee588578086
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Wed Oct 14 07:27:27 2026 +0000
[SPARK-60218][PYTHON] Migrate SQL_TABLE_UDF to the EvalTypeHandler pipeline and remove read_udtf
### What changes were proposed in this pull request?
- Add `TableUDTFHandler` for `SQL_TABLE_UDF` (the pickled Python UDTF). It extends `EvalTypeHandler` directly and keeps the existing serializer, `BatchedSerializer(CPickleSerializer(), 1)`. The `wrap_udtf` / `verify_and_convert_result` / `mapper` logic from `read_udtf` moves into the handler unchanged.
- Dispatch every UDF and UDTF eval type through `get_eval_type_handler` in `invoke_udf`: a `UDTFInfo` is read with `read_single_udtf`, a list of `UDFInfo` with `read_single_udf`. Eval types without a handler still go through `read_udfs`.
- Remove `read_udtf`.
### Why are the changes needed?
Part of SPARK-59415. All three UDTF eval types now run through handlers, so `read_udtf` is no longer needed.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- New unit tests for `TableUDTFHandler` in `test_udtf_eval_type_handlers.py`.
- Existing tests: `pyspark.sql.tests.test_udtf`, `pyspark.sql.tests.arrow.test_arrow_udtf`, `pyspark.sql.tests.connect.test_parity_udtf`, `pyspark.tests.test_worker`.
- Worker microbenchmark for `SQL_TABLE_UDF` (50k rows, 2 args, `worker.main` on in-memory input, 10 repeats). Output bytes are identical before and after.
| UDTF | before (ms) | after (ms) |
|---|---|---|
| identity | 206.6 +- 8.0 | 195.0 +- 7.0 |
| explode x3 | 241.0 +- 5.5 | 232.2 +- 8.0 |
| count (terminate only) | 136.2 +- 13.2 | 135.8 +- 7.2 |
| stringify | 204.0 +- 5.8 | 196.8 +- 8.8 |
### 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 #59432 from Yicong-Huang/udtf-remove-read-udtf.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit 18b5254a51f654e808f45ffeddd835bae37859fe)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/eval_handlers/__init__.py | 4 +-
python/pyspark/eval_handlers/_udtf.py | 127 +++++++++++++++
.../tests/test_udtf_eval_type_handlers.py | 109 ++++++++++++-
python/pyspark/worker.py | 178 ++++-----------------
4 files changed, 264 insertions(+), 154 deletions(-)
diff --git a/python/pyspark/eval_handlers/__init__.py b/python/pyspark/eval_handlers/__init__.py
index 281cda66a535..be4e5a3d8b1a 100644
--- a/python/pyspark/eval_handlers/__init__.py
+++ b/python/pyspark/eval_handlers/__init__.py
@@ -20,8 +20,8 @@
Not a supported extension API: the registry is per executor process and filled at
import time, so handlers must be defined in this package rather than by user code.
Each eval type handled here is an ``EvalTypeHandler`` subclass (in ``_base``) that
-declares its ``eval_type`` and self-registers at class definition, which
-``read_udfs`` and ``read_udtf`` look up via ``get_eval_type_handler``. An eval type may have
+declares its ``eval_type`` and self-registers at class definition, which the worker
+looks up via ``get_eval_type_handler``. An eval type may have
several handlers, each declaring the ``variant_value`` of a ``variant_conf`` ``RunnerConf``
property that selects it at lookup. Importing this package imports the concrete handler
submodules so they register.
diff --git a/python/pyspark/eval_handlers/_udtf.py b/python/pyspark/eval_handlers/_udtf.py
index 79dbd42c24ab..e470b2c69db4 100644
--- a/python/pyspark/eval_handlers/_udtf.py
+++ b/python/pyspark/eval_handlers/_udtf.py
@@ -37,6 +37,7 @@ from pyspark.eval_handlers._base import BatchEvalTypeHandler, EvalTypeHandler
from pyspark.eval_handlers._typing import InputBatch, OutputBatch
from pyspark.eval_handlers.utils import wrap_kwargs_support
from pyspark.eval_handlers.verification import verify_arrow_result, verify_pandas_result
+from pyspark.serializers import BatchedSerializer, CPickleSerializer, Serializer
from pyspark.sql.conversion import (
ArrowBatchTransformer,
ArrowTableToRowsConversion,
@@ -638,6 +639,132 @@ class UDTFEvalTypeHandler(EvalTypeHandler[InputBatch, OutputBatch]):
cleanup()
+class TableUDTFHandler(EvalTypeHandler[Any, tuple]):
+ """SQL_TABLE_UDF, the Python UDTF: input rows arrive pickled one at a time and the UDTF's
+ ``eval`` is called once per row; the results of each call are converted to the internal
+ representation and returned to the executor as one tuple."""
+
+ eval_type = PythonEvalType.SQL_TABLE_UDF
+ # Wraps the UDTF when its TABLE argument has PARTITION BY expressions; see read_single_udtf.
+ partition_wrapper: ClassVar[type[PartitionedUDTF]] = UDTFWithPartitions
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf: EvalConf
+ ) -> None:
+ super().__init__(udfs, runner_conf, eval_conf)
+ assert len(udfs) == 1, "One UDTF expected here."
+ udtf, args_offsets, kwargs_offsets, return_type = udfs[0]
+ self._check_output_row_against_schema = build_null_checker(return_type)
+
+ eval_func_kwargs_support, self._args_kwargs_offsets = wrap_kwargs_support(
+ getattr(udtf, "eval"), args_offsets, kwargs_offsets
+ )
+ self._eval = self._wrap_udtf(eval_func_kwargs_support, return_type)
+
+ self._terminate: Optional[Callable[..., tuple]]
+ if hasattr(udtf, "terminate"):
+ self._terminate = self._wrap_udtf(getattr(udtf, "terminate"), return_type)
+ else:
+ self._terminate = None
+
+ self._cleanup = getattr(udtf, "cleanup") if hasattr(udtf, "cleanup") else None
+
+ @property
+ def serializer(self) -> Serializer:
+ # Each row is a group so do not batch but send one by one.
+ return BatchedSerializer(CPickleSerializer(), 1)
+
+ def _wrap_udtf(self, f: Callable, return_type: StructType) -> Callable[..., tuple]:
+ assert return_type.needConversion()
+ toInternal = return_type.toInternal
+ return_type_size = len(return_type)
+ check_output_row_against_schema = self._check_output_row_against_schema
+
+ def verify_and_convert_result(result: Any) -> Any:
+ if result is not None:
+ if hasattr(result, "__UDT__"):
+ # UDT object should not be returned directly.
+ raise PySparkRuntimeError(
+ errorClass="UDTF_INVALID_OUTPUT_ROW_TYPE",
+ messageParameters={
+ "type": type(result).__name__,
+ "func": f.__name__,
+ },
+ )
+
+ if hasattr(result, "__len__") and len(result) != return_type_size:
+ raise PySparkRuntimeError(
+ errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
+ messageParameters={
+ "expected": str(return_type_size),
+ "actual": str(len(result)),
+ "func": f.__name__,
+ },
+ )
+
+ if not (isinstance(result, (list, dict, tuple)) or hasattr(result, "__dict__")):
+ raise PySparkRuntimeError(
+ errorClass="UDTF_INVALID_OUTPUT_ROW_TYPE",
+ messageParameters={
+ "type": type(result).__name__,
+ "func": f.__name__,
+ },
+ )
+ if check_output_row_against_schema is not None:
+ check_output_row_against_schema(result)
+ return toInternal(result)
+
+ # Evaluate the function and return a tuple back to the executor.
+ def evaluate(*a: Any) -> tuple:
+ try:
+ res = f(*a)
+ except SkipRestOfInputTableException:
+ raise
+ except Exception as e:
+ raise PySparkRuntimeError(
+ errorClass="UDTF_EXEC_ERROR",
+ messageParameters={"method_name": f.__name__, "error": str(e)},
+ )
+
+ if res is None:
+ # If the function returns None or does not have an explicit return statement,
+ # an empty tuple is returned to the executor.
+ # This is because directly constructing tuple(None) results in an exception.
+ return tuple()
+
+ if not isinstance(res, Iterable):
+ raise PySparkRuntimeError(
+ errorClass="UDTF_RETURN_NOT_ITERABLE",
+ messageParameters={
+ "type": type(res).__name__,
+ "func": f.__name__,
+ },
+ )
+
+ # If the function returns a result, we map it to the internal representation and
+ # returns the results as a tuple.
+ return tuple(map(verify_and_convert_result, res))
+
+ return evaluate
+
+ def run(self, split_index: int, data: Iterator[Any]) -> Iterator[tuple]:
+ eval = self._eval
+ terminate = self._terminate
+ cleanup = self._cleanup
+ args_kwargs_offsets = self._args_kwargs_offsets
+ try:
+ for a in data:
+ yield eval(*[a[o] for o in args_kwargs_offsets])
+ if terminate is not None:
+ yield terminate()
+ except SkipRestOfInputTableException:
+ if terminate is not None:
+ yield terminate()
+ finally:
+ if cleanup is not None:
+ cleanup()
+
+
class ArrowUDTFHandler(
BatchEvalTypeHandler["pa.RecordBatch"],
UDTFEvalTypeHandler["pa.RecordBatch", "pa.RecordBatch"],
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 3570b9ec3f93..684d596e42e8 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
@@ -28,6 +28,8 @@ from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers._udtf import TableUDTFHandler, UDTFWithPartitions
+from pyspark.serializers import BatchedSerializer, CPickleSerializer
from pyspark.sql.functions import SkipRestOfInputTableException
from pyspark.sql.pandas.serializers import ArrowStreamSerializer
from pyspark.sql.types import LongType, StructField, StructType
@@ -52,7 +54,6 @@ if have_pyarrow:
ArrowUDTFHandler,
ArrowUDTFWithPartition,
UDTFEvalTypeHandler,
- UDTFWithPartitions,
)
from pyspark.sql.conversion import ArrowBatchTransformer
@@ -112,6 +113,8 @@ class UDTFEvalTypeHandlerRegistrationTests(PySparkBaseTestCase):
# Each UDTF eval type dispatches to its handler by lookup; SQL_ARROW_TABLE_UDF picks
# its variant by the legacy pandas conversion conf.
for eval_type, conf, handler_cls in (
+ (PythonEvalType.SQL_TABLE_UDF, {}, TableUDTFHandler),
+ (PythonEvalType.SQL_TABLE_UDF, {legacy_conf: "true"}, TableUDTFHandler),
(PythonEvalType.SQL_ARROW_UDTF, {}, ArrowUDTFHandler),
(PythonEvalType.SQL_ARROW_TABLE_UDF, {}, ArrowTableUDTFArrowHandler),
(
@@ -129,6 +132,110 @@ class UDTFEvalTypeHandlerRegistrationTests(PySparkBaseTestCase):
self.assertIs(get_eval_type_handler(eval_type, RunnerConf(conf)), handler_cls)
+class TableUDTFHandlerTests(UDTFEvalTypeHandlerTestCase):
+ _ROW_RETURN_TYPE = StructType([StructField("x", LongType()), StructField("y", LongType())])
+
+ def _handler(self, udtf, args=(0,), kwargs=None, return_type=None):
+ return TableUDTFHandler(
+ udfs=[(udtf, list(args), kwargs or {}, return_type or self._ROW_RETURN_TYPE)],
+ 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
+ # eval/terminate call returned.
+ return list(handler.run(0, iter([list(r) for r in rows])))
+
+ def test_serializer_and_partition_wrapper(self):
+ ser = self._handler(_RowUDTF()).serializer
+ self.assertIsInstance(ser, BatchedSerializer)
+ self.assertIsInstance(ser.serializer, CPickleSerializer)
+ self.assertEqual(ser.batchSize, 1)
+ self.assertIs(TableUDTFHandler.partition_wrapper, UDTFWithPartitions)
+
+ def test_eval_terminate_cleanup(self):
+ udtf = _RowUDTF()
+ out = self._run(self._handler(udtf), [1], [2])
+ self.assertEqual(out, [((1, 2),), ((2, 4),), ((-1, -2),)])
+ self.assertEqual(udtf.calls, ["eval", "eval", "terminate", "cleanup"])
+
+ def test_no_terminate_or_cleanup(self):
+ class EvalOnly:
+ def eval(self, a):
+ yield a, a
+
+ self.assertEqual(self._run(self._handler(EvalOnly()), [1]), [((1, 1),)])
+
+ def test_kwargs(self):
+ class KwargsUDTF:
+ def eval(self, *, a):
+ yield a, -a
+
+ out = self._run(self._handler(KwargsUDTF(), args=(), kwargs={"a": 0}), [7])
+ self.assertEqual(out, [((7, -7),)])
+
+ def test_skip_rest_of_input_table(self):
+ udtf = _RowUDTF(skip_after=1)
+ out = self._run(self._handler(udtf), [1], [2], [3])
+ self.assertEqual(out, [((1, 2),), ((-1, -2),)])
+ self.assertEqual(udtf.calls, ["eval", "eval", "terminate", "cleanup"])
+
+ def test_none_result(self):
+ class NoneUDTF:
+ def eval(self, a):
+ return None
+
+ self.assertEqual(self._run(self._handler(NoneUDTF()), [1]), [()])
+
+ def test_non_nullable_column(self):
+ class NullUDTF:
+ def eval(self, a):
+ yield None, a
+
+ return_type = StructType(
+ [StructField("x", LongType(), nullable=False), StructField("y", LongType())]
+ )
+ with self.assertRaises(PySparkRuntimeError) as ctx:
+ self._run(self._handler(NullUDTF(), return_type=return_type), [1])
+ self.assertEqual(ctx.exception.getCondition(), "UDTF_EXEC_ERROR")
+
+ def test_errors(self):
+ class Raises:
+ def eval(self, a):
+ raise ValueError("boom")
+
+ class NotIterable:
+ def eval(self, a):
+ return 1
+
+ class WrongColumns:
+ def eval(self, a):
+ yield a, a, a
+
+ class NotARow:
+ def eval(self, a):
+ yield a
+
+ for udtf, error_class in [
+ (Raises(), "UDTF_EXEC_ERROR"),
+ (NotIterable(), "UDTF_RETURN_NOT_ITERABLE"),
+ (WrongColumns(), "UDTF_RETURN_SCHEMA_MISMATCH"),
+ (NotARow(), "UDTF_INVALID_OUTPUT_ROW_TYPE"),
+ ]:
+ with self.subTest(error_class=error_class):
+ with self.assertRaises(PySparkRuntimeError) as ctx:
+ self._run(self._handler(udtf), [1])
+ self.assertEqual(ctx.exception.getCondition(), error_class)
+
+ def test_cleanup_on_error(self):
+ udtf = _RowUDTF()
+ udtf.eval = lambda a: 1 # not iterable
+ with self.assertRaises(PySparkRuntimeError):
+ self._run(self._handler(udtf), [1])
+ self.assertEqual(udtf.calls, ["cleanup"])
+
+
@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
class ArrowUDTFHandlerTests(UDTFEvalTypeHandlerTestCase):
def _handler(self, udtf, args=(0,), kwargs=None, table_arg_offsets="0"):
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 20c4ac132a71..9d87ca0fab24 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -51,11 +51,7 @@ from pyspark.accumulators import (
)
from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError
from pyspark.eval_handlers._base import get_eval_type_handler
-from pyspark.eval_handlers._udtf import (
- UDTFEvalTypeHandler,
- UDTFWithPartitions,
- build_null_checker,
-)
+from pyspark.eval_handlers._udtf import TableUDTFHandler, UDTFEvalTypeHandler
from pyspark.eval_handlers.utils import extract_key_value_indexes, wrap_kwargs_support
from pyspark.eval_handlers.verification import (
verify_iterator_exhausted,
@@ -83,7 +79,6 @@ from pyspark.sql.conversion import (
LocalDataToArrowConversion,
PandasToArrowConversion,
)
-from pyspark.sql.functions import SkipRestOfInputTableException
from pyspark.sql.pandas.serializers import ArrowStreamSerializer
from pyspark.sql.pandas.types import to_arrow_type
from pyspark.sql.streaming.stateful_processor_api_client import StatefulProcessorApiClient
@@ -108,7 +103,7 @@ from pyspark.util import (
start_faulthandler_periodic_traceback,
with_faulthandler,
)
-from pyspark.worker_message import WorkerInitInfo
+from pyspark.worker_message import UDTFInfo, WorkerInitInfo
from pyspark.worker_util import (
EvalConf,
RunnerConf,
@@ -431,130 +426,6 @@ def read_single_udtf(pickleSer, udtf_info, partition_wrapper):
return udtf, udtf_info.args, udtf_info.kwargs, return_type
-# Read and process a serialized user-defined table function (UDTF) from a socket.
-# It expects the UDTF to be in a specific format and performs various checks to
-# ensure the UDTF is valid. This function also prepares a mapper function for applying
-# the UDTF logic to input rows.
-def read_udtf(pickleSer, udtf_info, eval_type, runner_conf, eval_conf):
- # If an eval type has a registered handler, dispatch through it: the handler
- # provides both the function and the serializer.
- handler_cls = get_eval_type_handler(eval_type, runner_conf)
- if handler_cls is not None:
- assert issubclass(handler_cls, UDTFEvalTypeHandler)
- udfs = [read_single_udtf(pickleSer, udtf_info, handler_cls.partition_wrapper)]
- handler = handler_cls(udfs=udfs, runner_conf=runner_conf, eval_conf=eval_conf)
- return handler.run, handler.serializer
-
- # Each row is a group so do not batch but send one by one.
- ser = BatchedSerializer(CPickleSerializer(), 1)
-
- udtf, args_offsets, kwargs_offsets, return_type = read_single_udtf(
- pickleSer, udtf_info, UDTFWithPartitions
- )
-
- check_output_row_against_schema = build_null_checker(return_type)
-
- def wrap_udtf(f, return_type):
- assert return_type.needConversion()
- toInternal = return_type.toInternal
- return_type_size = len(return_type)
-
- def verify_and_convert_result(result):
- if result is not None:
- if hasattr(result, "__UDT__"):
- # UDT object should not be returned directly.
- raise PySparkRuntimeError(
- errorClass="UDTF_INVALID_OUTPUT_ROW_TYPE",
- messageParameters={
- "type": type(result).__name__,
- "func": f.__name__,
- },
- )
-
- if hasattr(result, "__len__") and len(result) != return_type_size:
- raise PySparkRuntimeError(
- errorClass="UDTF_RETURN_SCHEMA_MISMATCH",
- messageParameters={
- "expected": str(return_type_size),
- "actual": str(len(result)),
- "func": f.__name__,
- },
- )
-
- if not (isinstance(result, (list, dict, tuple)) or hasattr(result, "__dict__")):
- raise PySparkRuntimeError(
- errorClass="UDTF_INVALID_OUTPUT_ROW_TYPE",
- messageParameters={
- "type": type(result).__name__,
- "func": f.__name__,
- },
- )
- if check_output_row_against_schema is not None:
- check_output_row_against_schema(result)
- return toInternal(result)
-
- # Evaluate the function and return a tuple back to the executor.
- def evaluate(*a) -> tuple:
- try:
- res = f(*a)
- except SkipRestOfInputTableException:
- raise
- except Exception as e:
- raise PySparkRuntimeError(
- errorClass="UDTF_EXEC_ERROR",
- messageParameters={"method_name": f.__name__, "error": str(e)},
- )
-
- if res is None:
- # If the function returns None or does not have an explicit return statement,
- # an empty tuple is returned to the executor.
- # This is because directly constructing tuple(None) results in an exception.
- return tuple()
-
- if not isinstance(res, Iterable):
- raise PySparkRuntimeError(
- errorClass="UDTF_RETURN_NOT_ITERABLE",
- messageParameters={
- "type": type(res).__name__,
- "func": f.__name__,
- },
- )
-
- # If the function returns a result, we map it to the internal representation and
- # returns the results as a tuple.
- return tuple(map(verify_and_convert_result, res))
-
- return evaluate
-
- eval_func_kwargs_support, args_kwargs_offsets = wrap_kwargs_support(
- getattr(udtf, "eval"), args_offsets, kwargs_offsets
- )
- eval = wrap_udtf(eval_func_kwargs_support, return_type)
-
- if hasattr(udtf, "terminate"):
- terminate = wrap_udtf(getattr(udtf, "terminate"), return_type)
- else:
- terminate = None
-
- cleanup = getattr(udtf, "cleanup") if hasattr(udtf, "cleanup") else None
-
- # Return an iterator of iterators.
- def mapper(_, it):
- try:
- for a in it:
- yield eval(*[a[o] for o in args_kwargs_offsets])
- if terminate is not None:
- yield terminate()
- except SkipRestOfInputTableException:
- if terminate is not None:
- yield terminate()
- finally:
- if cleanup is not None:
- cleanup()
-
- return mapper, ser
-
-
def _elementwise_udf_input_type(data_type: DataType, depth: int) -> DataType:
"""The element type ``depth`` ``ArrayType`` levels below ``data_type``.
@@ -734,17 +605,6 @@ def _elementwise_pandas_or_arrow_udf_output_to_flat_batch(
def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
- # If an eval type has a registered handler, dispatch through it: the handler
- # provides both the function and the serializer.
- handler_cls = get_eval_type_handler(eval_type, runner_conf)
- if handler_cls is not None:
- udfs = [
- read_single_udf(pickleSer, udf_info, eval_type, runner_conf, udf_index=udf_index)
- for udf_index, udf_info in enumerate(udf_info_list)
- ]
- handler = handler_cls(udfs=udfs, runner_conf=runner_conf, eval_conf=eval_conf)
- return handler.run, handler.serializer
-
if eval_type in (
PythonEvalType.SQL_ARROW_BATCHED_UDF,
PythonEvalType.SQL_ARROW_ELEMENTWISE_UDF,
@@ -2418,15 +2278,31 @@ def invoke_udf(message_receiver: SparkMessageReceiver, outfile: BinaryIO):
func, profiler, deserializer, serializer = read_command(pickleSer, init_info.udf_info)
else:
# UDF and UDTF runners read input and write output through a single serializer.
- is_udtf = eval_type in (
- PythonEvalType.SQL_TABLE_UDF,
- PythonEvalType.SQL_ARROW_TABLE_UDF,
- PythonEvalType.SQL_ARROW_UDTF,
- )
- read = read_udtf if is_udtf else read_udfs
- func, serializer = read(
- pickleSer, init_info.udf_info, eval_type, runner_conf, eval_conf
- )
+ # If an eval type has a registered handler, dispatch through it: the handler
+ # provides both the function and the serializer.
+ handler_cls = get_eval_type_handler(eval_type, runner_conf)
+ if handler_cls is not None:
+ if isinstance(init_info.udf_info, UDTFInfo):
+ assert issubclass(handler_cls, (UDTFEvalTypeHandler, TableUDTFHandler))
+ udfs = [
+ read_single_udtf(
+ pickleSer, init_info.udf_info, handler_cls.partition_wrapper
+ )
+ ]
+ else:
+ assert isinstance(init_info.udf_info, list)
+ udfs = [
+ read_single_udf(
+ pickleSer, udf_info, eval_type, runner_conf, udf_index=udf_index
+ )
+ for udf_index, udf_info in enumerate(init_info.udf_info)
+ ]
+ handler = handler_cls(udfs=udfs, runner_conf=runner_conf, eval_conf=eval_conf)
+ func, serializer = handler.run, handler.serializer
+ else:
+ func, serializer = read_udfs(
+ pickleSer, init_info.udf_info, eval_type, runner_conf, eval_conf
+ )
deserializer = serializer
split_index = init_info.split_index
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org