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

Code before:
from ray.data._internal.stats import StatsDict
from ray.data.context import DataContext


class UnionOperator(InternalQueueOperatorMixin, NAryOperator):
    """An operator that combines output blocks from
    two or more input operators into a single output."""

    def __init__(
        self,
        data_context: DataContext,
        *input_ops: PhysicalOperator,
    ):
        """Create a UnionOperator.

        Args:
            input_ops: Operators generating input data for this operator to union.
        """

        # By default, union does not preserve the order of output blocks.
        # To preserve the order, configure ExecutionOptions accordingly.
        self._preserve_order = False

        # Intermediary buffers used to store blocks from each input dependency.
        # Only used when `self._prserve_order` is True.
        self._input_buffers: List["BaseBundleQueue"] = [
            HashLinkedQueue() 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._output_buffer: collections.deque[RefBundle] = collections.deque()
        self._stats: StatsDict = {"Union": []}
        super().__init__(data_context, *input_ops)

    def start(self, options: ExecutionOptions):
        # Whether to preserve the order of the input data (both the
        # order of the input operators and the order of the blocks within).
        self._preserve_order = options.preserve_order
        super().start(options)

    def num_outputs_total(self) -> Optional[int]:
        num_outputs = 0
        for input_op in self.input_dependencies:
            input_num_outputs = input_op.num_outputs_total()
            if input_num_outputs is None:
                return None
            num_outputs += input_num_outputs

Review comment (suggestion):
```suggestion
        self._input_buffers: List[HashLinkedQueue] = [
```

Code after:
from ray.data._internal.stats import StatsDict
from ray.data.context import DataContext


class UnionOperator(InternalQueueOperatorMixin, NAryOperator):
    """An operator that combines output blocks from
    two or more input operators into a single output."""

    def __init__(
        self,
        data_context: DataContext,
        *input_ops: PhysicalOperator,
    ):
        """Create a UnionOperator.

        Args:
            input_ops: Operators generating input data for this operator to union.
        """

        # By default, union does not preserve the order of output blocks.
        # To preserve the order, configure ExecutionOptions accordingly.
        self._preserve_order = False

        # Intermediary buffers used to store blocks from each input dependency.
        # Only used when `self._prserve_order` is True.
        self._input_buffers: List[HashLinkedQueue] = [
            HashLinkedQueue() 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._output_buffer: collections.deque[RefBundle] = collections.deque()
        self._stats: StatsDict = {"Union": []}
        super().__init__(data_context, *input_ops)

    def start(self, options: ExecutionOptions):
        # Whether to preserve the order of the input data (both the
        # order of the input operators and the order of the blocks within).
        self._preserve_order = options.preserve_order
        super().start(options)

    def num_outputs_total(self) -> Optional[int]:
        num_outputs = 0
        for input_op in self.input_dependencies:
            input_num_outputs = input_op.num_outputs_total()
            if input_num_outputs is None:
                return None
            num_outputs += input_num_outputs