[SPARK-58939][PYTHON] Move Arrow collect batch reordering into ArrowCollectSerializer

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

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