HyukjinKwon commented on code in PR #57260:
URL: https://github.com/apache/spark/pull/57260#discussion_r3593026035
##########
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
+ total_rows = 0
+ average_arrow_row_size = 0.0
- # This must be generator comprehension - do not materialize.
- return f(stateful_processor_api_client, mode, key,
values_gen())
- else:
- # mode == PROCESS_TIMER or mode == COMPLETE
- return f(stateful_processor_api_client, mode, None, iter([]))
+ def flatten_columns(cur_batch: "pa.RecordBatch", col_name: str) ->
"pa.Table":
+ struct_column =
cur_batch.column(cur_batch.schema.get_field_index(col_name))
+ # Check if the entire column is null: an empty table (no
columns)
+ # signals the struct is absent from this batch.
+ if struct_column.null_count == len(struct_column):
+ return pa.Table.from_arrays([], names=[])
+ field_names = [
+ struct_column.type[i].name for i in
range(struct_column.type.num_fields)
+ ]
+ field_arrays = [
+ struct_column.field(i) for i in
range(struct_column.type.num_fields)
+ ]
+ return pa.Table.from_arrays(field_arrays, names=field_names)
- elif eval_type == PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF:
+ def to_pandas(table: "pa.Table") -> list:
+ return ArrowBatchTransformer.to_pandas(
+ table,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+
+ def row_stream() -> Iterator[tuple]:
+ nonlocal total_bytes, total_rows, average_arrow_row_size
+ for batch in data:
+ # Short circuit batch size stats if the batch size is
+ # unlimited as computing batch size is computationally
+ # expensive.
+ if arrow_max_bytes_per_batch != 2**31 - 1 and
batch.num_rows > 0:
+ total_bytes += sum(
+ buf.size
+ for col in batch.columns
+ for buf in col.buffers()
+ if buf is not None
+ )
+ total_rows += batch.num_rows
+ average_arrow_row_size = total_bytes / total_rows
+
+ data_table = flatten_columns(batch, "inputData")
+ init_table = flatten_columns(batch, "initState")
+
+ # Empty table has no columns. Each batch carries either
+ # input data or init state, never both.
+ has_data = data_table.num_columns > 0
+ has_init = init_table.num_columns > 0
+ assert not (has_data and has_init)
+
+ if has_data:
+ for row in pd.concat(to_pandas(data_table),
axis=1).itertuples(index=False):
+ batch_key = tuple(row[o] for o in key_offsets)
+ yield (batch_key, row, None)
+ elif has_init:
+ for row in pd.concat(to_pandas(init_table),
axis=1).itertuples(index=False):
+ batch_key = tuple(row[o] for o in init_key_offsets)
+ yield (batch_key, None, row)
+
+ empty_dataframe = pd.DataFrame()
+
+ def generate_data_batches() -> Iterator[tuple]:
+ """
+ Deserialize ArrowRecordBatches and return a generator of
+ (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
+ from the same grouping key should appear sequentially.
+ """
+ for batch_key, group_rows in itertools.groupby(row_stream(),
key=lambda x: x[0]):
+ rows = []
+ init_state_rows = []
+ for _, row, init_state_row in group_rows:
+ if row is not None:
+ rows.append(row)
+ if init_state_row is not None:
+ init_state_rows.append(init_state_row)
+
+ total_len = len(rows) + len(init_state_rows)
+ if (
+ total_len >= arrow_max_records_per_batch
+ or total_len * average_arrow_row_size >=
arrow_max_bytes_per_batch
+ ):
+ yield (
+ batch_key,
+ pd.DataFrame(rows) if rows else
empty_dataframe.copy(),
+ (
+ pd.DataFrame(init_state_rows)
+ if init_state_rows
+ else empty_dataframe.copy()
+ ),
+ )
+ rows = []
+ init_state_rows = []
+ if rows or init_state_rows:
+ yield (
+ batch_key,
+ pd.DataFrame(rows) if rows else
empty_dataframe.copy(),
+ (
+ pd.DataFrame(init_state_rows)
+ if init_state_rows
+ else empty_dataframe.copy()
+ ),
+ )
+
+ def convert_results(
+ result_iter: Iterable["pd.DataFrame"],
+ ) -> Iterator["pa.RecordBatch"]:
+ # TODO(SPARK-49100): add verification that elements in
result_iter are
+ # indeed of type pd.DataFrame and confirm to assigned cols
Review Comment:
Nit: typo `confirm to` should be `conform to`. This is pre-existing text
(also at lines 508 and 527) carried verbatim onto a new line, so purely
optional — but easy to fix here since the line is new.
```suggestion
# indeed of type pd.DataFrame and conform to assigned cols
```
--
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]
