Repository: ray-project/ray
File: python/ray/data/_internal/execution/operators/union_operator.py (Python)

Code before:
    def all_inputs_done(self) -> None:
        super().all_inputs_done()

        if not self._preserve_order:
            return

        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.
        return len(self._output_buffer) > 0

    def _get_next_inner(self) -> RefBundle:
        refs = self._output_buffer.popleft()
        self._metrics.on_output_dequeued(refs)
        return refs

    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


Review comment (suggestion):
```suggestion
        while self._next_input_to_flush < len(self._input_buffers) and self._input_done_flags[self._next_input_to_flush]:
```

Code after:
        self._try_round_robin()
        assert all(not buffer.has_next() for buffer in self._input_buffers)

    def has_next(self) -> bool:
        # Check if the output buffer still contains at least one block.
        return len(self._output_buffer) > 0

    def _get_next_inner(self) -> RefBundle:
        refs = self._output_buffer.popleft()
        self._metrics.on_output_dequeued(refs)
        return refs

    def get_stats(self) -> StatsDict:
        return self._stats

    def _try_round_robin(self) -> None:
        """Try to move blocks from input buffers to output in round-robin order.

        Pulls one block from the current input, then advances to the next.
        If the current input's buffer is empty but not done, we return
        without advancing to the next input so the scheduling won't be blocked.

        This ensures deterministic ordering of output blocks:
        - We iterate through inputs in a fixed order (0, 1, 2, ..., 0, 1, ...).
        - We only advance to the next input after consuming exactly one block
          from the current input (or if the current input is exhausted).
        - If an input is not ready (empty but not done), we return
          rather than skipping it, preserving the round-robin sequence.
        """
        num_inputs = len(self._input_buffers)

        while True:
            buffer = self._input_buffers[self._current_input_index]

            if buffer.has_next():
                refs = buffer.get_next()
                self._metrics.on_input_dequeued(refs)
                self._output_buffer.append(refs)
                self._metrics.on_output_queued(refs)
            elif not self._input_done_flags[self._current_input_index] or all(
                not buffer.has_next() for buffer in self._input_buffers
            ):
                return

            self._current_input_index = (self._current_input_index + 1) % num_inputs