[Data] Fix Union operator to avoid blocking when preserve order

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

@@ -35,15 +35,13 @@ def __init__(
         self._preserve_order = False
 
         # Intermediary buffers used to store blocks from each input dependency.
-        # Only used when `self._prserve_order` is True.
+        # Only used when `self._preserve_order` is True.
         self._input_buffers: List[BundleQueue] = [
             FIFOBundleQueue() for _ in range(len(input_ops))
         ]
 
-        # The index of the input dependency that is currently the source of
-        # the output buffer. New inputs from this input dependency will be added
-        # directly to the output buffer. Only used when `self._preserve_order` is True.
-        self._input_idx_to_output = 0
+        self._input_done_flags: List[bool] = [False] * len(input_ops)
+        self._next_input_to_flush = 0
 
         self._output_buffer: collections.deque[RefBundle] = collections.deque()
         self._stats: StatsDict = {"Union": []}
@@ -102,26 +100,32 @@ def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
         assert not self.has_completed()
         assert 0 <= input_index <= len(self._input_dependencies), input_index
 
-        if not self._preserve_order:
+        if not self._preserve_order or input_index == self._next_input_to_flush:
+            if self._preserve_order:
+                while self._input_buffers[input_index]:
+                    buffered = self._input_buffers[input_index].get_next()
+                    self._metrics.on_input_dequeued(buffered)
+                    self._output_buffer.append(buffered)
+                    self._metrics.on_output_queued(buffered)
             self._output_buffer.append(refs)
             self._metrics.on_output_queued(refs)
         else:
             self._input_buffers[input_index].add(refs)
             self._metrics.on_input_queued(refs)
 
+    def input_done(self, input_index: int) -> None:
+        self._input_done_flags[input_index] = True
+        if self._preserve_order:
+            self._try_advance_and_flush()
+
     def all_inputs_done(self) -> None:
         super().all_inputs_done()
 
         if not self._preserve_order:
             return
 
-        assert len(self._output_buffer) == 0, len(self._output_buffer)
-        for input_buffer in self._input_buffers:
-            while input_buffer:
-                refs = input_buffer.get_next()
-                self._metrics.on_input_dequeued(refs)
-                self._output_buffer.append(refs)
-                self._metrics.on_output_queued(refs)
+        self._try_advance_and_flush()
+        assert all(not ref for ref in self._input_buffers)
 
     def has_next(self) -> bool:
         # Check if the output buffer still contains at least one block.
@@ -134,3 +138,16 @@ def _get_next_inner(self) -> RefBundle:
 
     def get_stats(self) -> StatsDict:
         return self._stats
+
+    def _try_advance_and_flush(self) -> None:
+
+        while self._next_input_to_flush < len(self._input_buffers):
+            if not self._input_done_flags[self._next_input_to_flush]:
+                break
+            buffer_to_flush = self._input_buffers[self._next_input_to_flush]
+            while buffer_to_flush:
+                refs = buffer_to_flush.get_next()
+                self._metrics.on_input_dequeued(refs)
+                self._output_buffer.append(refs)
+                self._metrics.on_output_queued(refs)
+            self._next_input_to_flush += 1