code: def test_limit_operator(ray_start_regular_shared):
      """Test basic functionalities of LimitOperator."""
      num_refs = 3
      num_rows_per_block = 3
      total_rows = num_refs * num_rows_per_block
      # Test limits with different values, from 0 to more than input size.
      limits = list(range(0, total_rows + 2))
      for limit in limits:
          refs = make_ref_bundles(
              [[i] * num_rows_per_block for i in range(num_refs)]
          )
          input_op = InputDataBuffer(refs)
          limit_op = LimitOperator(limit, input_op)
          limit_op.all_inputs_done = MagicMock(
              wraps=limit_op.all_inputs_done
          )
          if limit == 0:
              # If the limit is 0, the operator should be completed immediately.
              assert limit_op.completed()
              assert limit_op._limit_reached()
          cur_rows = 0
          loop_count = 0
          while input_op.has_next() and not limit_op._limit_reached():
              loop_count += 1
              assert not limit_op.completed(), limit
              assert limit_op.need_more_inputs(), limit
              limit_op.add_input(input_op.get_next(), 0)
              while limit_op.has_next():
                  # Drain the outputs. So the limit operator
                  # will be completed when the limit is reached.
                  limit_op.get_next()
              cur_rows += num_rows_per_block
              if cur_rows >= limit:
                  assert limit_op.all_inputs_done.call_count == 1, limit
                  assert limit_op.completed(), limit
                  assert limit_op._limit_reached(), limit
                  assert not limit_op.need_more_inputs(), limit
              else:
                  assert limit_op.all_inputs_done.call_count == 0, limit
                  assert not limit_op.completed(), limit
                  assert not limit_op._limit_reached(), limit
                  assert limit_op.need_more_inputs(), limit
          limit_op.all_inputs_done()
          # After inputs done, the number of output bundles
          # should be the same as the number of `add_input`s.
          assert limit_op.num_outputs_total() == loop_count, limit
          assert limit_op.completed(), limit
  
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 20 for calling a global function with example usage: test_limit_operator(ray_start_regular_shared) without return types