Yicong-Huang commented on code in PR #57260:
URL: https://github.com/apache/spark/pull/57260#discussion_r3599770846
##########
python/pyspark/worker.py:
##########
@@ -3499,44 +3466,217 @@ 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 func(
+ split_index: int,
+ data: 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 two separate lazy iterators
+ (data DataFrames and init-state DataFrames), then once for
+ PROCESS_TIMER and once for COMPLETE.
+ """
+ total_bytes = 0
Review Comment:
yes your assumption is correct, the serializer code logic was moved here. we
will clean up the serializers after all TWS eval types are moved.
We will soon to work on the eval type class to encapsulate this. The
logic is indeed complex but that is needed for now to keep the same behavior.
many logic can be further reduced to use existing helpers, but would introduce
some behavior change (e.g., more strict type check). to be done separately.
--
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]
