single output buffer

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

Author: iamjustinhsu

python/ray/data/_internal/execution/operators/hash_shuffle.py

@@ -1560,14 +1560,13 @@ def __init__(
             aggregator_id, target_partition_ids
         )
         # One buffer per partition to enable concurrent finalization
-        self._output_buffers: Dict[int, BlockOutputBuffer] = {
-            partition_id: BlockOutputBuffer(
-                output_block_size_option=OutputBlockSizeOption(
-                    target_max_block_size=data_context.target_max_block_size
-                )
+        self._output_buffer_lock = threading.Lock()
+        self._output_buffer = BlockOutputBuffer(
+            output_block_size_option=OutputBlockSizeOption(
+                target_max_block_size=data_context.target_max_block_size
             )
-            for partition_id in target_partition_ids
-        }
+        )
+        self._partition_ids_to_finalize: Set[int] = set(target_partition_ids)
 
     def submit(self, input_seq_id: int, partition_id: int, partition_shard: Block):
         with self._lock:
@@ -1581,19 +1580,22 @@ def finalize(
             exec_stats_builder = BlockExecStats.builder()
             # Finalize given partition id
             block = self._agg.finalize(partition_id)
-            exec_stats = exec_stats_builder.build()
+
             # Clear any remaining state (to release resources)
             self._agg.clear(partition_id)
-
-        # No lock needed - each partition has its own buffer
-        output_buffer = self._output_buffers[partition_id]
-
-        output_buffer.add_block(block)
-        output_buffer.finalize()
-        while output_buffer.has_next():
-            block = output_buffer.next()
-            yield block
-            yield BlockMetadataWithSchema.from_block(block, stats=exec_stats)
+            self._partition_ids_to_finalize.remove(partition_id)
+
+        with self._output_buffer_lock:
+            self._output_buffer.add_block(block)
+            if len(self._partition_ids_to_finalize) == 0:
+                # this is the last finalized partition_id
+                self._output_buffer.finalize()
+            while self._output_buffer.has_next():
+                block = self._output_buffer.next()
+                yield block
+                yield BlockMetadataWithSchema.from_block(
+                    block, stats=exec_stats_builder.build()
+                )
 
 
 def _get_total_cluster_resources() -> ExecutionResources: