code: def test_prefetch_batches_locally(
      num_batches_to_prefetch, batch_size
  ):
      class DummyPrefetcher(BlockPrefetcher):
          def __init__(self):
              self.windows = []
  
          def prefetch_blocks(self, blocks: List[Block]):
              if batch_size is None:
                  assert len(blocks) == num_batches_to_prefetch
              else:
                  assert (
                      sum(len(block) for block in blocks)
                      >= batch_size * num_batches_to_prefetch
                  )
              self.windows.append(blocks)
  
      num_blocks = 10
      num_rows = 2
      prefetcher = DummyPrefetcher()
      blocks = list(
          block_generator(num_blocks=num_blocks, num_rows=num_rows)
      )
      prefetch_block_iter = prefetch_batches_locally(
          iter(blocks),
          prefetcher=prefetcher,
          num_batches_to_prefetch=num_batches_to_prefetch,
          batch_size=batch_size,
      )
  
      block_count = 0
      prefetched_blocks = []
      previous_num_windows = 1
  
      for block in prefetch_block_iter:
          prefetched_blocks.append(block)
          block_count += 1
          remaining_rows = (num_blocks - block_count) * num_rows
          if (
              batch_size is None
              and block_count < num_blocks - num_batches_to_prefetch
          ):
              # Test that we are actually prefetching in advance if this is not the last
              # block.
              assert len(prefetcher.windows) == previous_num_windows + 1
              previous_num_windows = len(prefetcher.windows)
          elif (
              batch_size is not None
              and remaining_rows > batch_size * num_batches_to_prefetch
          ):
              # Test that we are actually prefetching in advance if this is not the last
              # batch.
              assert len(prefetcher.windows) == previous_num_windows + 1
              previous_num_windows = len(prefetcher.windows)
  
      # Test that original blocks are unchanged.
      assert prefetched_blocks == [block for block, metadata, in blocks]
  
file_path: ray/python/ray/data/tests/block_batching/test_iter_batches.py
imports: import itertools
  import queue
  import threading
  import time
  from typing import Iterator, List, Tuple
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data._internal.block_batching.interfaces import Batch, BlockPrefetcher
  from ray.data._internal.block_batching.iter_batches import iter_batches, prefetch_batches_locally, restore_original_order
  from ray.data.block import Block, BlockMetadata
label_desc: Use image node_id 2 for calling a global function with example usage: test_prefetch_batches_locally(num_batches_to_prefetch, batch_size) without return types