code: def test_split_operator(ray_start_regular_shared, equal, chunk_size):
      num_input_blocks = 100
      num_splits = 3
      # Add this many input blocks each time.
      # Make sure it is greater than num_splits * 2,
      # so we can test the output order of `OutputSplitter.get_next`.
      num_add_input_blocks = 10
      input_op = InputDataBuffer(
          make_ref_bundles(
              [[i] * chunk_size for i in range(num_input_blocks)]
          )
      )
      op = OutputSplitter(input_op, num_splits, equal=equal)
  
      # Feed data and implement streaming exec.
      output_splits = [[] for _ in range(num_splits)]
      op.start(ExecutionOptions())
      while input_op.has_next():
          for _ in range(num_add_input_blocks):
              if not input_op.has_next():
                  break
              op.add_input(input_op.get_next(), 0)
          while op.has_next():
              ref = op.get_next()
              assert ref.owns_blocks, ref
              for block, _ in ref.blocks:
                  assert ref.output_split_idx is not None
                  output_splits[ref.output_split_idx].extend(
                      list(ray.get(block)["id"])
                  )
      op.all_inputs_done()
  
      expected_splits = [[] for _ in range(num_splits)]
      for i in range(num_splits):
          for j in range(i, num_input_blocks, num_splits):
              expected_splits[i].extend([j] * chunk_size)
      if equal:
          min_len = min(
              len(expected_splits[i]) for i in range(num_splits)
          )
          for i in range(num_splits):
              expected_splits[i] = expected_splits[i][:min_len]
      for i in range(num_splits):
          assert output_splits[i] == expected_splits[i], (
              output_splits[i],
              expected_splits[i],
          )
  
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 9 for calling a global function with example usage: test_split_operator(ray_start_regular_shared, equal, chunk_size) without return types