[Data] Make iter_batches an Iterable (ray-project#37881)

Make iter_batches, iter_torch_batches, iter_rows and iter_tf_batches return an Iterable instead of a single use Iterator. --------- Signed-off-by: amogkam <[email protected]>
sihanwang41 · Aug 1, 2023 · 0bded3b · 0bded3b
1 parent 844c2cc
commit 0bded3b
Show file tree

Hide file tree

Showing 13 changed files with 199 additions and 435 deletions.
diff --git a/doc/source/data/doc_code/tensor.py b/doc/source/data/doc_code/tensor.py
diff --git a/doc/source/train/doc_code/key_concepts.py b/doc/source/train/doc_code/key_concepts.py
@@ -37,7 +37,7 @@ def train_fn(config):
  # Local worker rank on the current machine
  "local_rank": session.get_local_rank(),
  # Data
- "data_shard": next(dataset_shard.iter_batches(batch_format="pandas")),
+ "data_shard": next(iter(dataset_shard.iter_batches(batch_format="pandas"))),
  }
  )
 

diff --git a/python/ray/data/_internal/iterator/pipelined_iterator.py b/python/ray/data/_internal/iterator/pipelined_iterator.py
@@ -1,4 +1,13 @@
-from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional, Tuple, Union
+from typing import (
+ TYPE_CHECKING,
+ Any,
+ Callable,
+ Iterable,
+ Iterator,
+ Optional,
+ Tuple,
+ Union,
+)
 
 from ray.data._internal.stats import DatasetStats
 from ray.data.block import Block, BlockMetadata, DataBatch
@@ -70,7 +79,7 @@ def iter_batches(
  _finalize_fn: Optional[Callable[[Any], Any]] = None,
  # Deprecated.
  prefetch_blocks: int = 0,
- ) -> Iterator[DataBatch]:
+ ) -> Iterable[DataBatch]:
  # Set prefetch_batches to default of 0 for DatasetPipeline.
  return super().iter_batches(
  prefetch_batches=prefetch_batches,

diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
@@ -2296,8 +2296,12 @@ def take_batch(
  batch_format = _apply_strict_mode_batch_format(batch_format)
  try:
  res = next(
- self.iter_batches(
- batch_size=batch_size, prefetch_batches=0, batch_format=batch_format
+ iter(
+ self.iter_batches(
+ batch_size=batch_size,
+ prefetch_batches=0,
+ batch_format=batch_format,
+ )
  )
  )
  except StopIteration: