This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 9cdd8a8fc073 [SPARK-58939][PYTHON] Move Arrow collect batch reordering
into ArrowCollectSerializer
9cdd8a8fc073 is described below
commit 9cdd8a8fc0733228712da267a1c38af97a21031d
Author: Yicong Huang <[email protected]>
AuthorDate: Mon Aug 24 18:20:42 2026 +0000
[SPARK-58939][PYTHON] Move Arrow collect batch reordering into
ArrowCollectSerializer
### What changes were proposed in this pull request?
Move the batch-reordering logic for Arrow-based collect from the
`_collect_as_arrow` caller into `ArrowCollectSerializer`.
`Dataset.collectAsArrowToPython` (Spark classic) streams Arrow batches to
Python out of order and appends batch-order indices at the end. Previously the
serializer yielded raw batches plus the order list, and `_collect_as_arrow`
split off the order list and reordered. Now `ArrowCollectSerializer` extends
`ArrowStreamSerializer` and yields the batches already ordered, so the caller
drops the results slicing, the reorder comprehension, and the `isinstance`
check. The unused `dump_stream [...]
To keep `spark.sql.execution.arrow.pyspark.selfDestruct.enabled` effective,
the serializer drops its reference to each batch as it yields it (`batches[i] =
None`; the indices are a permutation), so peak memory stays ~1x.
### Why are the changes needed?
The ordering protocol is an implementation detail of the collect stream;
owning it inside the serializer removes duplicated logic from
`_collect_as_arrow` and drops dead code, with no behavior change.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing `python/pyspark/sql/tests/arrow/test_arrow.py` (toPandas/toArrow +
self-destruct) in CI; also validated reorder, JVM-error, empty-stream, and repr
against a simulated stream.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #58221 from Yicong-Huang/simplify-arrow-collect-serializer.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit 0eb8ef1314955e7b76c8baa3b02e3d94acbc1cb4)
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/sql/pandas/conversion.py | 31 +++++-------
python/pyspark/sql/pandas/serializers.py | 81 +++++++++++++++-----------------
2 files changed, 49 insertions(+), 63 deletions(-)
diff --git a/python/pyspark/sql/pandas/conversion.py
b/python/pyspark/sql/pandas/conversion.py
index d0b89354e100..48998a6a239e 100644
--- a/python/pyspark/sql/pandas/conversion.py
+++ b/python/pyspark/sql/pandas/conversion.py
@@ -533,7 +533,7 @@ class PandasConversionMixin:
jsocket_auth_server,
) = self._jdf.collectAsArrowToPython()
- # Collect list of un-ordered batches where last element is a list of
correct order indices
+ # Collect the batches already reordered by ArrowCollectSerializer.
try:
with _load_from_socket((port, auth_secret),
ArrowCollectSerializer()) as batch_stream:
if split_batches:
@@ -545,32 +545,23 @@ class PandasConversionMixin:
# converted.
import pyarrow as pa
- results = []
- for batch_or_indices in batch_stream:
- if isinstance(batch_or_indices, pa.RecordBatch):
- batch_or_indices = pa.RecordBatch.from_arrays(
- [
- # This call actually reallocates the array
- pa.concat_arrays([array])
- for array in batch_or_indices
- ],
- schema=batch_or_indices.schema,
- )
- results.append(batch_or_indices)
+ batches = [
+ pa.RecordBatch.from_arrays(
+ # This call actually reallocates the array
+ [pa.concat_arrays([array]) for array in batch],
+ schema=batch.schema,
+ )
+ for batch in batch_stream
+ ]
else:
- results = list(batch_stream)
+ batches = list(batch_stream)
finally:
with unwrap_spark_exception():
# Join serving thread and raise any exceptions from
collectAsArrowToPython
jsocket_auth_server.getResult()
- # Separate RecordBatches from batch order indices in results
- batches = results[:-1]
- batch_order = results[-1]
-
if len(batches) or empty_list_if_zero_records:
- # Re-order the batch list using the correct order
- return [batches[i] for i in batch_order]
+ return batches
else:
from pyspark.sql.pandas.types import to_arrow_schema
import pyarrow as pa
diff --git a/python/pyspark/sql/pandas/serializers.py
b/python/pyspark/sql/pandas/serializers.py
index 3b2bb187ee4d..2ee42a9df829 100644
--- a/python/pyspark/sql/pandas/serializers.py
+++ b/python/pyspark/sql/pandas/serializers.py
@@ -42,49 +42,6 @@ class SpecialLengths:
START_ARROW_STREAM = -6
-class ArrowCollectSerializer(Serializer):
- """
- Deserialize a stream of batches followed by batch order information. Used
in
- PandasConversionMixin._collect_as_arrow() after invoking
Dataset.collectAsArrowToPython()
- in the JVM.
- """
-
- def __init__(self):
- self.serializer = ArrowStreamSerializer()
-
- def dump_stream(self, iterator, stream):
- return self.serializer.dump_stream(iterator, stream)
-
- def load_stream(self, stream):
- """
- Load a stream of un-ordered Arrow RecordBatches, where the last
iteration yields
- a list of indices that can be used to put the RecordBatches in the
correct order.
- """
- # load the batches
- for batch in self.serializer.load_stream(stream):
- yield batch
-
- # load the batch order indices or propagate any error that occurred in
the JVM
- num = read_int(stream)
- if num == -1:
- error_msg = UTF8Deserializer().loads(stream)
- raise PySparkRuntimeError(
- errorClass="ERROR_OCCURRED_WHILE_CALLING",
- messageParameters={
- "func_name": "ArrowCollectSerializer.load_stream",
- "error_msg": error_msg,
- },
- )
- batch_order = []
- for i in range(num):
- index = read_int(stream)
- batch_order.append(index)
- yield batch_order
-
- def __repr__(self):
- return "ArrowCollectSerializer(%s)" % self.serializer
-
-
class ArrowStreamSerializer(Serializer):
"""
Serializes Arrow record batches as a plain stream.
@@ -149,6 +106,44 @@ class ArrowStreamSerializer(Serializer):
return "ArrowStreamSerializer(write_start_stream=%s)" %
self._write_start_stream
+class ArrowCollectSerializer(ArrowStreamSerializer):
+ """
+ Extends :class:`ArrowStreamSerializer` to load Arrow RecordBatches that
the JVM
+ sends out of order, followed by the indices giving their correct order,
and yields
+ the batches already reordered. Used in
PandasConversionMixin._collect_as_arrow()
+ after invoking Dataset.collectAsArrowToPython() in the JVM.
+ """
+
+ def load_stream(self, stream: IO[bytes]) -> Iterator["pa.RecordBatch"]:
+ """Load the out-of-order batches, then yield them in the correct
order."""
+ batches = list(super().load_stream(stream))
+
+ # Load the batch order indices, or propagate any error that occurred
in the JVM.
+ num = read_int(stream)
+ if num == -1:
+ error_msg = UTF8Deserializer().loads(stream)
+ raise PySparkRuntimeError(
+ errorClass="ERROR_OCCURRED_WHILE_CALLING",
+ messageParameters={
+ "func_name": "ArrowCollectSerializer.load_stream",
+ "error_msg": error_msg,
+ },
+ )
+ # Yield the batches in order, dropping our reference to each as it
goes so
+ # that, when selfDestruct is enabled, the caller's reallocated copy is
the
+ # only remaining reference and the original batch can be freed
immediately
+ # rather than staying pinned here until the stream is fully consumed.
The
+ # indices are a permutation, so each batch is yielded exactly once.
+ for _ in range(num):
+ i = read_int(stream)
+ batch = batches[i]
+ batches[i] = None
+ yield batch
+
+ def __repr__(self) -> str:
+ return "ArrowCollectSerializer()"
+
+
class ArrowStreamGroupSerializer(ArrowStreamSerializer):
"""
Extends :class:`ArrowStreamSerializer` with group-count protocol for
loading
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
