code: def restore_original_order(
      batch_iter: Iterator[Batch],
  ) -> Iterator[Batch]:
      """Restores the original order of the provided `batch_iter`
  
      This function will yield items from `base_iterator` in the correct order based on
      each batch's batch_idx. All indexes are expected to be unique.
  
      `batch_iter` is expected to not have any missing indexes. All indexes from 0 to len
      (base_iterator) must be present.
      """
      next_index_required = 0
      buffer: Dict[int, Batch] = {}
      for batch in batch_iter:
          assert batch.batch_idx not in buffer
          buffer[batch.batch_idx] = batch
          while next_index_required in buffer:
              yield buffer.pop(next_index_required)
              next_index_required += 1
  
      while next_index_required in buffer:
          yield buffer.pop(next_index_required)
          next_index_required += 1
  
file_path: ray/python/ray/data/_internal/block_batching/iter_batches.py
imports: import collections
  from contextlib import nullcontext
  from typing import Any, Callable, Dict, Iterator, Optional, Tuple
  import ray
  from ray.data._internal.block_batching.interfaces import Batch, BlockPrefetcher
  from ray.data._internal.block_batching.util import ActorBlockPrefetcher, WaitBlockPrefetcher, blocks_to_batches, collate, extract_data_from_batch, finalize_batches, format_batches, resolve_block_refs
  from ray.data._internal.memory_tracing import trace_deallocation
  from ray.data._internal.stats import DatasetStats
  from ray.data._internal.util import make_async_gen
  from ray.data.block import Block, BlockMetadata, DataBatch
  from ray.data.context import DataContext
  from ray.types import ObjectRef
label_desc: Use image node_id 4 for calling a global function with example usage: restore_original_order(batch_iter) without return types