code: def test_union_operator(ray_start_regular_shared, preserve_order):
      """Test basic functionalities of UnionOperator."""
      execution_options = ExecutionOptions(
          preserve_order=preserve_order
      )
      ctx = ray.data.DataContext.get_current()
      ctx.execution_options = execution_options
  
      num_rows_per_block = 3
      data0 = make_ref_bundles(
          [[i] * num_rows_per_block for i in range(3)]
      )
      data1 = make_ref_bundles(
          [[i] * num_rows_per_block for i in range(2)]
      )
      data2 = make_ref_bundles(
          [[i] * num_rows_per_block for i in range(1)]
      )
  
      op0 = InputDataBuffer(data0)
      op1 = InputDataBuffer(data1)
      op2 = InputDataBuffer(data2)
      union_op = UnionOperator(op0, op1, op2)
      union_op.start(execution_options)
  
      assert not union_op.has_next()
      union_op.add_input(op0.get_next(), 0)
      assert union_op.has_next()
  
      assert union_op.get_next() == data0[0]
      assert not union_op.has_next()
  
      union_op.add_input(op0.get_next(), 0)
      union_op.add_input(op0.get_next(), 0)
      assert union_op.get_next() == data0[1]
      assert union_op.get_next() == data0[2]
  
      union_op.input_done(0)
      assert not union_op.completed()
      if preserve_order:
          assert union_op._input_idx_to_output == 1
  
      if preserve_order:
          union_op.add_input(op1.get_next(), 1)
          union_op.add_input(op2.get_next(), 2)
          assert union_op._input_idx_to_output == 1
  
          assert union_op.get_next() == data1[0]
          assert not union_op.has_next()
  
          # Check the case where an input op which is not the op
          # corresponding to _input_idx_to_output finishes first.
          union_op.input_done(2)
          assert union_op._input_idx_to_output == 1
  
          union_op.add_input(op1.get_next(), 1)
          assert union_op.has_next()
          assert union_op.get_next() == data1[1]
          assert not union_op.has_next()
          # Marking the current output buffer source op will
          # increment _input_idx_to_output to the next source.
          union_op.input_done(1)
          assert union_op._input_idx_to_output == 2
          assert union_op.has_next()
          assert union_op.get_next() == data2[0]
      else:
          union_op.add_input(op1.get_next(), 1)
          union_op.add_input(op2.get_next(), 2)
          union_op.add_input(op1.get_next(), 1)
          # The output will be in the same order as the inputs
          # were added with `add_input()`.
          assert union_op.get_next() == data1[0]
          assert union_op.get_next() == data2[0]
          assert union_op.get_next() == data1[1]
  
      assert all([len(b) == 0 for b in union_op._input_buffers])
  
      _take_outputs(union_op)
      union_op.all_inputs_done()
      assert union_op.completed()
  
imports: import collections
  import random
  import time
  from typing import Any, Iterable, List
  from unittest.mock import MagicMock
  import numpy
  import pandas
  import pytest
  import ray
  from ray._private.test_utils import wait_for_condition
  from ray.data._internal.compute import ActorPoolStrategy, TaskPoolStrategy
  from ray.data._internal.execution.interfaces import ExecutionOptions, PhysicalOperator, RefBundle
  from ray.data._internal.execution.operators.actor_pool_map_operator import ActorPoolMapOperator
  from ray.data._internal.execution.operators.base_physical_operator import AllToAllOperator
  from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
  from ray.data._internal.execution.operators.limit_operator import LimitOperator
  from ray.data._internal.execution.operators.map_operator import MapOperator, _BlockRefBundler
  from ray.data._internal.execution.operators.map_transformer import create_map_transformer_from_block_fn
  from ray.data._internal.execution.operators.output_splitter import OutputSplitter
  from ray.data._internal.execution.operators.task_pool_map_operator import TaskPoolMapOperator
  from ray.data._internal.execution.operators.union_operator import UnionOperator
  from ray.data._internal.execution.util import make_ref_bundles
  from ray.data.block import Block
  from ray.data.context import DataContext
  from ray.data.tests.util import run_one_op_task, run_op_tasks_sync
  from ray.tests.conftest import *
label_desc: Use image node_id 22 for calling a global function with example usage: test_union_operator(ray_start_regular_shared, preserve_order) without return types