Yicong-Huang commented on code in PR #57260:
URL: https://github.com/apache/spark/pull/57260#discussion_r3590070188
##########
python/pyspark/worker.py:
##########
@@ -3499,44 +3466,215 @@ def convert_results(result_iter):
return transform_with_state_func, None, ser, ser
if eval_type ==
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_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
+ import pandas as pd
+
+ assert num_udfs == 1, "One TRANSFORM_WITH_STATE_PANDAS_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]
+ output_schema = StructType([StructField("_0", return_type)])
+
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
+ )
+ arrow_max_bytes_per_batch = runner_conf.arrow_max_bytes_per_batch
- if mode == TransformWithStateInPandasFuncMode.PROCESS_DATA:
- key = a[1]
+ def transform_with_state_init_state_func(
+ split_index: int,
+ batches: Iterator[pa.RecordBatch],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply transformWithStateInPandas UDF with initial state.
- def values_gen():
- for x in a[2]:
- retVal = x[1]
- initVal = x[2]
- yield retVal, initVal
+ 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 pandas DataFrames bounded by
+ arrow_max_records_per_batch and arrow_max_bytes_per_batch. The UDF
+ is invoked once per grouping key with a lazy iterator of
Review Comment:
right, that was not accurate. I have updated it. Thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
