Yicong-Huang opened a new pull request, #57337:
URL: https://github.com/apache/spark/pull/57337

   ### 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 and are tracked/deleted together in the cleanup step. The 
`TransformWithStateInPySparkRowSerializer` and 
`TransformWithStateInPySparkRowInitStateSerializer` classes are kept for now 
and deleted in a follow-up 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`): before = 
`upstream/master`, after = this PR. Values are from one representative run per 
side; the conclusion is consistent across runs.
   
   (ASV before/after comparison table will be added once benchmark runs 
complete.)
   


-- 
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]

Reply via email to