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 d0d0adb938da [SPARK-58191][PYTHON] Refactor 
SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF and 
SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF
d0d0adb938da is described below

commit d0d0adb938da475613ea6c876fb73e1158cace97
Author: Yicong Huang <[email protected]>
AuthorDate: Wed Jul 22 16:23:28 2026 +0000

    [SPARK-58191][PYTHON] Refactor SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF and 
SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF
    
    ### What changes were proposed in this pull request?
    
    This PR refactors `SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF` and 
`SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF` so that the worker uses 
the plain `ArrowStreamSerializer` for pure Arrow stream I/O, moving the 
per-eval-type logic from `TransformWithStateInPySparkRowSerializer` / 
`TransformWithStateInPySparkRowInitStateSerializer` into `read_udfs()` in 
`worker.py`:
    
    - **Row variant**: reading each input batch row by row into `Row` objects, 
regrouping by grouping key, invoking the UDF once per key with a lazy iterator 
of `Row`s, and converting the yielded `(iterator of Row, Spark type)` result 
back into a struct-wrapped Arrow `RecordBatch`.
    - **Row init-state variant**: flattening the `inputData` / `initState` 
struct columns, regrouping rows by grouping key, re-chunking into `Row` lists 
bounded by `arrow_max_records_per_batch`, splitting each chunk into a data 
iterator and an init-state iterator, and converting results back to Arrow.
    
    This mirrors the earlier refactor of the Pandas variants 
([SPARK-57400](https://issues.apache.org/jira/browse/SPARK-57400), 
[SPARK-58128](https://issues.apache.org/jira/browse/SPARK-58128)). Both 
variants are done together in one PR because they share a serializer class 
hierarchy (`...RowInitStateSerializer` extends `...RowSerializer`) and are 
deleted together in the follow-up cleanup. The two serializer classes are kept 
for now and deleted in that cleanup.
    
    ### Why are the changes needed?
    
    Part of [SPARK-55388](https://issues.apache.org/jira/browse/SPARK-55388). 
Keeping serializers as pure Arrow stream I/O and concentrating 
eval-type-specific logic in `worker.py` makes the per-eval-type data flow 
explicit.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Existing tests. No behavior change.
    
    ASV comparison (`bench_eval_type.TransformWithStateRowUDFTimeBench` / 
`TransformWithStateRowInitStateUDFTimeBench`, `-a repeat=3`, each combo run in 
a fresh Python process to match production): before = `upstream/master`, after 
= this PR. Values are from one representative run per side, keeping ASV's 
native `+-` intervals; the conclusion is consistent across runs. All deltas are 
within run-to-run noise, showing no regression. (A handful of combos flagged 
larger deltas only when the be [...]
    
    ```text
    time_worker -- SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF
    
    scenario        udf           before       after        diff
    --------------  ------------  -----------  -----------  -----
    few_groups_lg   count_udf     3.97+-0.05s  4.01+-0.02s  +1.0%
    few_groups_lg   identity_udf  4.95+-0.02s  4.83+-0.05s  -2.4%
    few_groups_lg   rebuild_udf   5.65+-0.07s  5.63+-0.03s  -0.4%
    few_groups_sm   count_udf     399+-3ms     421+-5ms     +5.6%
    few_groups_sm   identity_udf  528+-17ms    503+-2ms     -4.7%
    few_groups_sm   rebuild_udf   606+-30ms    600+-14ms    -1.1%
    many_groups_lg  count_udf     1.65+-0.02s  1.65+-0.03s  +0.4%
    many_groups_lg  identity_udf  2.18+-0.03s  2.10+-0.05s  -3.6%
    many_groups_lg  rebuild_udf   2.45+-0.01s  2.45+-0.04s  -0.1%
    many_groups_sm  count_udf     1.80+-0.02s  1.87+-0.02s  +4.0%
    many_groups_sm  identity_udf  2.32+-0.03s  2.42+-0.05s  +4.3%
    many_groups_sm  rebuild_udf   2.81+-0.05s  2.76+-0.02s  -1.6%
    mixed_cols      count_udf     1.69+-0.03s  1.71+-0.04s  +1.1%
    mixed_cols      identity_udf  2.11+-0.04s  2.05+-0.07s  -2.8%
    mixed_cols      rebuild_udf   2.42+-0.04s  2.41+-0.05s  -0.5%
    nested_struct   count_udf     2.51+-0.03s  2.58+-0.01s  +2.8%
    nested_struct   identity_udf  3.34+-0.02s  3.31+-0.06s  -0.8%
    nested_struct   rebuild_udf   3.74+-0.05s  3.74+-0.03s  -0.1%
    wide_cols       count_udf     4.54+-0.11s  4.55+-0.04s  +0.2%
    wide_cols       identity_udf  5.39+-0.03s  5.52+-0.16s  +2.3%
    wide_cols       rebuild_udf   6.27+-0.13s  6.23+-0.04s  -0.7%
    ```
    
    ```text
    time_worker -- SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF
    
    scenario        udf           before       after        diff
    --------------  ------------  -----------  -----------  -----
    few_groups_lg   count_udf     3.65+-0.05s  3.75+-0.08s  +2.7%
    few_groups_lg   identity_udf  4.65+-0.03s  4.70+-0.06s  +1.2%
    few_groups_lg   rebuild_udf   5.25+-0.10s  5.21+-0.06s  -0.8%
    few_groups_sm   count_udf     425+-9ms     407+-14ms    -4.1%
    few_groups_sm   identity_udf  500+-13ms    529+-16ms    +5.8%
    few_groups_sm   rebuild_udf   593+-44ms    607+-30ms    +2.4%
    many_groups_lg  count_udf     2.10+-0.10s  2.17+-0.06s  +3.2%
    many_groups_lg  identity_udf  2.60+-0.05s  2.62+-0.07s  +0.6%
    many_groups_lg  rebuild_udf   2.83+-0.06s  2.82+-0.02s  -0.3%
    many_groups_sm  count_udf     4.54+-0.06s  4.62+-0.08s  +1.6%
    many_groups_sm  identity_udf  5.05+-0.10s  5.03+-0.13s  -0.4%
    many_groups_sm  rebuild_udf   5.44+-0.13s  5.43+-0.04s  -0.2%
    mixed_cols      count_udf     1.74+-0.05s  1.76+-0.04s  +1.4%
    mixed_cols      identity_udf  2.16+-0.03s  2.16+-0.03s  -0.1%
    mixed_cols      rebuild_udf   2.40+-0.05s  2.35+-0.02s  -1.8%
    nested_struct   count_udf     2.79+-0.01s  2.76+-0.02s  -1.3%
    nested_struct   identity_udf  3.63+-0.03s  3.60+-0.08s  -0.8%
    nested_struct   rebuild_udf   3.94+-0.09s  3.88+-0.09s  -1.5%
    wide_cols       count_udf     4.38+-0.05s  4.49+-0.03s  +2.5%
    wide_cols       identity_udf  5.63+-0.10s  5.51+-0.09s  -2.2%
    wide_cols       rebuild_udf   5.97+-0.00s  5.83+-0.12s  -2.4%
    ```
    
    Closes #57337 from Yicong-Huang/refactor-tws-row-init.
    
    Authored-by: Yicong Huang <[email protected]>
    Signed-off-by: Yicong-Huang <[email protected]>
    (cherry picked from commit 64a30fd5b78d3afdd7eb0fc8287dce22e69a94ec)
    Signed-off-by: Yicong-Huang <[email protected]>
---
 python/pyspark/worker.py | 338 +++++++++++++++++++++++++++++++++++------------
 1 file changed, 257 insertions(+), 81 deletions(-)

diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 558a3d53cd89..cd5b0fb73c32 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -78,8 +78,6 @@ from pyspark.sql.pandas.serializers import (
     ArrowStreamSerializer,
     ArrowStreamGroupSerializer,
     ArrowStreamCoGroupSerializer,
-    TransformWithStateInPySparkRowSerializer,
-    TransformWithStateInPySparkRowInitStateSerializer,
 )
 from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type
 from pyspark.sql.types import (
@@ -500,37 +498,6 @@ def verify_arrow_result(
         )
 
 
-def wrap_grouped_transform_with_state_udf(f, return_type, runner_conf):
-    def wrapped(stateful_processor_api_client, mode, key, values):
-        result_iter = f(stateful_processor_api_client, mode, key, values)
-
-        # TODO(SPARK-XXXXX): add verification that elements in result_iter are
-        # indeed of type Row and confirm to assigned cols
-
-        return result_iter
-
-    return lambda p, m, k, v: [(wrapped(p, m, k, v), return_type)]
-
-
-def wrap_grouped_transform_with_state_init_state_udf(f, return_type, 
runner_conf):
-    def wrapped(stateful_processor_api_client, mode, key, values):
-        if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
-            values_gen = values[0]
-            init_states_gen = values[1]
-        else:
-            values_gen = iter([])
-            init_states_gen = iter([])
-
-        result_iter = f(stateful_processor_api_client, mode, key, values_gen, 
init_states_gen)
-
-        # TODO(SPARK-XXXXX): add verification that elements in result_iter are
-        # indeed of type pd.DataFrame and confirm to assigned cols
-
-        return result_iter
-
-    return lambda p, m, k, v: [(wrapped(p, m, k, v), return_type)]
-
-
 def wrap_kwargs_support(f, args_offsets, kwargs_offsets):
     if len(kwargs_offsets):
         keys = list(kwargs_offsets.keys())
@@ -708,11 +675,9 @@ def read_single_udf(pickleSer, udf_info, eval_type, 
runner_conf, udf_index):
     elif eval_type == 
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF:
         return func, args_offsets, return_type
     elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF:
-        return args_offsets, wrap_grouped_transform_with_state_udf(func, 
return_type, runner_conf)
+        return func, args_offsets, return_type
     elif eval_type == 
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF:
-        return args_offsets, wrap_grouped_transform_with_state_init_state_udf(
-            func, return_type, runner_conf
-        )
+        return func, args_offsets, return_type
     elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
         argspec = inspect.getfullargspec(chained_func)  # signature was lost 
when wrapping it
         return func, args_offsets, return_type, len(argspec.args)
@@ -1928,14 +1893,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type, 
runner_conf, eval_conf):
             ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
         elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
             ser = ArrowStreamCoGroupSerializer(write_start_stream=True)
-        elif eval_type == 
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF:
-            ser = TransformWithStateInPySparkRowSerializer(
-                
arrow_max_records_per_batch=runner_conf.arrow_max_records_per_batch
-            )
-        elif eval_type == 
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF:
-            ser = TransformWithStateInPySparkRowInitStateSerializer(
-                
arrow_max_records_per_batch=runner_conf.arrow_max_records_per_batch
-            )
         else:
             ser = ArrowStreamSerializer(write_start_stream=True)
     else:
@@ -3303,7 +3260,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, 
runner_conf, eval_conf):
                 (grouping key, pandas.DataFrame) chunks.
 
                 This function must avoid materializing multiple Arrow
-                RecordBatches into memory at the same time. And data chunks
+                RecordBatches into memory at the same time, and data chunks
                 from the same grouping key should appear sequentially.
                 """
                 for batch_key, group_rows in itertools.groupby(row_stream(), 
key=lambda x: x[0]):
@@ -3480,7 +3437,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type, 
runner_conf, eval_conf):
                 (grouping key, data DataFrame, init-state DataFrame) chunks.
 
                 This function must avoid materializing multiple Arrow
-                RecordBatches into memory at the same time. And data chunks
+                RecordBatches into memory at the same time, and data chunks
                 from the same grouping key should appear sequentially.
                 """
                 for batch_key, group_rows in itertools.groupby(row_stream(), 
key=lambda x: x[0]):
@@ -3914,64 +3871,283 @@ def read_udfs(pickleSer, udf_info_list, eval_type, 
runner_conf, eval_conf):
         return func, None, ser, ser
 
     if eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF:
-        # We assume there is only one UDF here because grouped map doesn't
-        # support combining multiple UDFs.
-        assert num_udfs == 1
+        import pyarrow as pa
+
+        assert num_udfs == 1, "One TRANSFORM_WITH_STATE_PYTHON_ROW UDF 
expected here."
+        udf, arg_offsets, return_type = udfs[0]
 
         # See TransformWithStateInPySparkExec for how arg_offsets are used to
-        # distinguish between grouping attributes and data attributes
-        arg_offsets, f = udfs[0]
+        # distinguish between grouping attributes and data attributes.
         parsed_offsets = extract_key_value_indexes(arg_offsets)
-        ser.key_offsets = parsed_offsets[0][0]
+        key_offsets = parsed_offsets[0][0]
+
         stateful_processor_api_client = StatefulProcessorApiClient(
             eval_conf.state_server_socket_port, eval_conf.grouping_key_schema
         )
 
-        def mapper(a):
-            mode = a[0]
+        def func(
+            split_index: int,
+            data: Iterator[pa.RecordBatch],
+        ) -> Iterator[pa.RecordBatch]:
+            """Apply transformWithStateInPySpark UDF over Row objects.
+
+            Input batches are read row by row without materializing whole
+            batches. Rows carrying the same grouping key appear sequentially, 
so
+            they are regrouped by key and the UDF is invoked once per key with 
a
+            lazy iterator of Row objects, then once for PROCESS_TIMER and once
+            for COMPLETE. The UDF yields (iterator of Row, Spark type) pairs 
that
+            are converted back into Arrow RecordBatches wrapped in a single
+            struct column for the output stream.
+            """
 
-            if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
-                key = a[1]
-                values = a[2]
+            def generate_data_batches() -> Iterator[Tuple[Any, Any]]:
+                """
+                Deserialize ArrowRecordBatches and return a generator of
+                (grouping key, Row) tuples.
 
-                # This must be generator comprehension - do not materialize.
-                return f(stateful_processor_api_client, mode, key, values)
-            else:
-                # mode == PROCESS_TIMER or mode == COMPLETE
-                return f(stateful_processor_api_client, mode, None, iter([]))
+                This function must avoid materializing multiple Arrow
+                RecordBatches into memory at the same time, and data chunks
+                from the same grouping key should appear sequentially.
+                """
+                for batch in data:
+                    DataRow = Row(*batch.schema.names)
+
+                    # Iterate row by row without converting the whole batch.
+                    num_cols = batch.num_columns
+                    for row_idx in range(batch.num_rows):
+                        row_key = tuple(batch[o][row_idx].as_py() for o in 
key_offsets)
+                        row = DataRow(*(batch.column(i)[row_idx].as_py() for i 
in range(num_cols)))
+                        yield row_key, row
+
+            def convert_results(result_rows: Iterable[Any]) -> 
Iterator["pa.RecordBatch"]:
+                # TODO(SPARK-XXXXX): add verification that elements in 
result_rows
+                # are indeed of type Row and conform to assigned cols
+
+                # Convert spark type to arrow type
+                # TODO: we need to make this configurable, currently using 
default values.
+                arrow_type = to_arrow_type(
+                    return_type,
+                    timezone="UTC",
+                    prefers_large_types=False,
+                )
 
-    elif eval_type == 
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF:
-        # We assume there is only one UDF here because grouped map doesn't
-        # support combining multiple UDFs.
-        assert num_udfs == 1
+                rows_as_dict = [row.asDict(True) for row in result_rows]
+
+                pdf_schema = pa.schema(list(arrow_type))
+                record_batch = pa.RecordBatch.from_pylist(rows_as_dict, 
schema=pdf_schema)
+                yield ArrowBatchTransformer.wrap_struct(record_batch)
+
+            for key, group in itertools.groupby(generate_data_batches(), 
key=lambda x: x[0]):
+                # This must be a generator expression - do not materialize.
+                values_gen = map(lambda x: x[1], group)
+                yield from convert_results(
+                    udf(
+                        stateful_processor_api_client,
+                        TransformWithStateInPandasFuncMode.PROCESS_DATA,
+                        key,
+                        values_gen,
+                    )
+                )
+
+            yield from convert_results(
+                udf(
+                    stateful_processor_api_client,
+                    TransformWithStateInPandasFuncMode.PROCESS_TIMER,
+                    None,
+                    iter([]),
+                )
+            )
+
+            yield from convert_results(
+                udf(
+                    stateful_processor_api_client,
+                    TransformWithStateInPandasFuncMode.COMPLETE,
+                    None,
+                    iter([]),
+                )
+            )
+
+        # profiling is not supported for UDF
+        return func, None, ser, ser
+
+    if eval_type == 
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF:
+        import pyarrow as pa
+
+        assert num_udfs == 1, "One TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE 
UDF expected here."
+        udf, arg_offsets, return_type = udfs[0]
 
         # See TransformWithStateInPandasExec for how arg_offsets are used to
-        # distinguish between grouping attributes and data attributes
-        arg_offsets, f = udfs[0]
+        # distinguish between grouping attributes and data attributes.
         # parsed offsets:
         # [
         #     [groupingKeyOffsets, dedupDataOffsets],
         #     [initStateGroupingOffsets, dedupInitDataOffsets]
         # ]
         parsed_offsets = extract_key_value_indexes(arg_offsets)
-        ser.key_offsets = parsed_offsets[0][0]
-        ser.init_key_offsets = parsed_offsets[1][0]
+        key_offsets = parsed_offsets[0][0]
+        init_key_offsets = parsed_offsets[1][0]
+
         stateful_processor_api_client = StatefulProcessorApiClient(
             eval_conf.state_server_socket_port, eval_conf.grouping_key_schema
         )
 
-        def mapper(a):
-            mode = a[0]
+        arrow_max_records_per_batch = runner_conf.arrow_max_records_per_batch
+        arrow_max_records_per_batch = (
+            arrow_max_records_per_batch if arrow_max_records_per_batch > 0 
else 2**31 - 1
+        )
 
-            if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
-                key = a[1]
-                values = a[2]
+        def func(
+            split_index: int,
+            data: Iterator[pa.RecordBatch],
+        ) -> Iterator[pa.RecordBatch]:
+            """Apply transformWithStateInPySpark UDF with initial state over 
Rows.
 
-                # This must be generator comprehension - do not materialize.
-                return f(stateful_processor_api_client, mode, key, values)
-            else:
-                # mode == PROCESS_TIMER or mode == COMPLETE
-                return f(stateful_processor_api_client, mode, None, iter([]))
+            The input batches carry two struct columns, ``inputData`` and
+            ``initState``; each batch holds one or the other but never both.
+            Rows are flattened out of whichever struct is present, regrouped by
+            grouping key, and re-chunked into Row lists bounded by
+            arrow_max_records_per_batch. The UDF is invoked once per chunk with
+            two separate iterators (data Rows and init-state Rows), then once 
for
+            PROCESS_TIMER and once for COMPLETE. The UDF yields (iterator of 
Row,
+            Spark type) pairs that are converted back into Arrow RecordBatches
+            wrapped in a single struct column for the output stream.
+            """
+
+            def extract_rows(
+                cur_batch: "pa.RecordBatch", col_name: str, offsets: list
+            ) -> Optional[Iterator[Tuple[Any, Any]]]:
+                data_column = 
cur_batch.column(cur_batch.schema.get_field_index(col_name))
+
+                # Check if the entire column is null.
+                if data_column.null_count == len(data_column):
+                    return None
+
+                data_field_names = [
+                    data_column.type[i].name for i in 
range(data_column.type.num_fields)
+                ]
+                data_field_arrays = [
+                    data_column.field(i) for i in 
range(data_column.type.num_fields)
+                ]
+
+                DataRow = Row(*data_field_names)
+
+                table = pa.Table.from_arrays(data_field_arrays, 
names=data_field_names)
+
+                if table.num_rows == 0:
+                    return None
+
+                def row_iterator() -> Iterator[Tuple[Any, Any]]:
+                    for row_idx in range(table.num_rows):
+                        key = tuple(table.column(o)[row_idx].as_py() for o in 
offsets)
+                        row = DataRow(
+                            *(table.column(i)[row_idx].as_py() for i in 
range(table.num_columns))
+                        )
+                        yield (key, row)
+
+                return row_iterator()
+
+            def row_stream() -> Iterator[Tuple[Any, Optional[Any], 
Optional[Any]]]:
+                # The arrow batch is written in the schema:
+                # schema: StructType = new StructType()
+                #     .add("inputData", dataSchema)
+                #     .add("initState", initStateSchema)
+                # We parse each batch into tuples of (key, inputData, 
initState).
+                # Each batch will have either init_data or input_data, not 
both.
+                for batch in data:
+                    input_result = extract_rows(batch, "inputData", 
key_offsets)
+                    init_result = extract_rows(batch, "initState", 
init_key_offsets)
+
+                    assert not (input_result is not None and init_result is 
not None)
+
+                    if input_result is not None:
+                        for key, input_data_row in input_result:
+                            yield (key, input_data_row, None)
+                    elif init_result is not None:
+                        for key, init_state_row in init_result:
+                            yield (key, None, init_state_row)
+
+            def generate_data_batches() -> Iterator[Tuple[Any, Tuple[Any, 
Any]]]:
+                """
+                Deserialize ArrowRecordBatches and return a generator of
+                (grouping key, (data Rows iterator, init-state Rows iterator))
+                chunks bounded by arrow_max_records_per_batch.
+
+                This function must avoid materializing multiple Arrow
+                RecordBatches into memory at the same time, and data chunks
+                from the same grouping key should appear sequentially.
+                """
+                for k, group_rows in itertools.groupby(row_stream(), 
key=lambda x: x[0]):
+                    input_rows: list = []
+                    init_rows: list = []
+
+                    for _, input_row, init_row in group_rows:
+                        if input_row is not None:
+                            input_rows.append(input_row)
+                        if init_row is not None:
+                            init_rows.append(init_row)
+
+                        total_len = len(input_rows) + len(init_rows)
+                        if total_len >= arrow_max_records_per_batch:
+                            yield (k, (iter(input_rows), iter(init_rows)))
+                            input_rows = []
+                            init_rows = []
+
+                    if input_rows or init_rows:
+                        yield (k, (iter(input_rows), iter(init_rows)))
+
+            def convert_results(result_rows: Iterable[Any]) -> 
Iterator["pa.RecordBatch"]:
+                # TODO(SPARK-XXXXX): add verification that elements in 
result_rows
+                # are indeed of type Row and conform to assigned cols
+
+                # Convert spark type to arrow type
+                # TODO: we need to make this configurable, currently using 
default values.
+                arrow_type = to_arrow_type(
+                    return_type,
+                    timezone="UTC",
+                    prefers_large_types=False,
+                )
+
+                rows_as_dict = [row.asDict(True) for row in result_rows]
+
+                pdf_schema = pa.schema(list(arrow_type))
+                record_batch = pa.RecordBatch.from_pylist(rows_as_dict, 
schema=pdf_schema)
+                yield ArrowBatchTransformer.wrap_struct(record_batch)
+
+            for key, group in itertools.groupby(generate_data_batches(), 
key=lambda x: x[0]):
+                # These must be generator expressions - do not materialize.
+                for _, (values_gen, init_states_gen) in group:
+                    yield from convert_results(
+                        udf(
+                            stateful_processor_api_client,
+                            TransformWithStateInPandasFuncMode.PROCESS_DATA,
+                            key,
+                            values_gen,
+                            init_states_gen,
+                        )
+                    )
+
+            yield from convert_results(
+                udf(
+                    stateful_processor_api_client,
+                    TransformWithStateInPandasFuncMode.PROCESS_TIMER,
+                    None,
+                    iter([]),
+                    iter([]),
+                )
+            )
+
+            yield from convert_results(
+                udf(
+                    stateful_processor_api_client,
+                    TransformWithStateInPandasFuncMode.COMPLETE,
+                    None,
+                    iter([]),
+                    iter([]),
+                )
+            )
+
+        # profiling is not supported for UDF
+        return func, None, ser, ser
 
     else:
 


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to