code: def iter_batches(
      ds: Dataset,
      block_format: Literal["pandas", "pyarrow", "simple"] = "pyarrow",
      batch_size: Optional[int] = None,
      batch_format: Literal[
          "default", "pandas", "pyarrow", "numpy"
      ] = "default",
      local_shuffle_buffer_size: Optional[int] = None,
      use_default_params: bool = False,
  ) -> Dataset:
      num_batches = 0
      if use_default_params:
          for batch in ds.iter_batches():
              num_batches += 1
      else:
          for batch in ds.iter_batches(
              batch_format=batch_format,
              batch_size=batch_size,
              local_shuffle_buffer_size=local_shuffle_buffer_size,
          ):
              num_batches += 1
      print(
          "iter_batches done, block_format:",
          block_format,
          "batch_format:",
          batch_format,
          "num_rows:",
          ds.count(),
          "num_blocks:",
          ds.num_blocks(),
          "num_batches:",
          num_batches,
      )
      return ds
  
imports: from typing import Optional
  import ray
  from ray.data.dataset import Dataset
  from benchmark import Benchmark
  from typing import Literal
label_desc: Use image node_id 1 for calling a global function with example usage: iter_batches(ds, block_format, batch_size, batch_format, local_shuffle_buffer_size, use_default_params) and returns: ds