From c75b0232a242d89d2c4961f243f9e73d60d333d6 Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Fri, 21 Aug 2026 22:49:49 +0000 Subject: [PATCH] refactor: reorder arrow collect batches inside ArrowCollectSerializer --- 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 d0b89354e1005..48998a6a239e3 100644 --- a/python/pyspark/sql/pandas/conversion.py +++ b/python/pyspark/sql/pandas/conversion.py @@ -533,7 +533,7 @@ def _collect_as_arrow( 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 @@ def _collect_as_arrow( # 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 3b2bb187ee4dc..2ee42a9df8297 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 @@ def __repr__(self) -> str: 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