Yicong-Huang opened a new pull request, #58221: URL: https://github.com/apache/spark/pull/58221
### What changes were proposed in this pull request? This PR moves the batch-reordering logic for Arrow-based collect from the `_collect_as_arrow` caller into `ArrowCollectSerializer`. `Dataset.collectAsArrowToPython` (Spark classic) streams Arrow record batches to Python out of order (as partition tasks finish, without buffering on the driver) and appends batch-order indices at the end of the stream, expecting Python to reorder them. Previously `ArrowCollectSerializer` wrapped `ArrowStreamSerializer` by composition, yielded the raw batches followed by the order list, and left `_collect_as_arrow` to split the order list off the results and reorder with `[batches[i] for i in batch_order]`. Now `ArrowCollectSerializer` extends `ArrowStreamSerializer` and owns the whole protocol: its `load_stream` reads all batches, reads the order indices, and yields the batches already in the correct order. The caller no longer sees the trailing order list, so `_collect_as_arrow` drops the results slicing (`results[:-1]` / `results[-1]`), the reorder comprehension, and the `isinstance` check that used to distinguish batches from the order list. The unused `dump_stream` delegate is removed (collect is one-directional: the JVM serializes, Python deserializes). Net -14 lines. To keep the `spark.sql.execution.arrow.pyspark.selfDestruct.enabled` optimization effective, the serializer drops its reference to each batch as it yields it (`batches[i] = None`); the indices form a permutation, so each batch is yielded exactly once. This lets the caller's reallocated (self-destruct) copy become the only reference so the original batch memory is freed immediately, keeping peak memory at ~1x as before rather than pinning all batches until the stream is fully consumed. ### Why are the changes needed? The reordering and the on-the-wire order protocol are an implementation detail of the Arrow collect stream. Keeping them inside `ArrowCollectSerializer` lets the caller work with a plain ordered batch iterator, removes the duplicated/awkward reordering logic from `_collect_as_arrow`, and drops the dead `dump_stream` delegate, with no behavior change. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Existing tests in `python/pyspark/sql/tests/arrow/test_arrow.py` exercise the `toPandas` / `toArrow` Arrow paths and `_collect_as_arrow(split_batches=True)` for self-destruct, and run in CI. Additionally validated the serializer's reorder, JVM-error propagation (`num == -1`), empty-stream, and `repr` behaviors locally against a simulated collect stream. ### Was this patch authored or co-authored using generative AI tooling? No. -- 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]
