gaogaotiantian commented on code in PR #57260:
URL: https://github.com/apache/spark/pull/57260#discussion_r3599666840
##########
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 conform to assigned cols
+ for result in result_iter:
+ if isinstance(return_type, StructType) and not
isinstance(result, pd.DataFrame):
+ raise PySparkValueError(
+ "Invalid return type. Please make sure that the
UDF returns a "
+ "pandas.DataFrame when the specified return type
is StructType."
+ )
+ yield PandasToArrowConversion.convert(
+ [result],
+ output_schema,
+ timezone=runner_conf.timezone,
+ safecheck=runner_conf.safecheck,
+ arrow_cast=True,
+ assign_cols_by_name=runner_conf.assign_cols_by_name,
+
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
+ )
+
+ for key, group in itertools.groupby(generate_data_batches(),
key=lambda x: x[0]):
+ # These must be generator expressions - do not materialize. The
+ # UDF receives the data and init-state DataFrames as two
+ # separate iterators, with empty chunks filtered out.
+ group_data, group_init = itertools.tee(group, 2)
Review Comment:
This is so crazy LOL
##########
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:
This is the reason I once asked "how long is the per eval type specific
code". I understand that this could be temporary, but it's very difficult to
have a 100-line code block in our function. We eventually just need to define
an eval-type based class and put all the eval type specific code there.
I assume because we moved the logic from the serializer to this function so
it's so large and we did not remove the serializer because it was used by some
other streaming eval types?
--
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]
