[Datasets] [WIP] Prototype wrapping operation-based lazy compute model.

python/ray/data/dataset.py

@@ -28,6 +28,7 @@
     import tensorflow as tf
     from ray.data.dataset_pipeline import DatasetPipeline
     from ray.data.grouped_dataset import GroupedDataset
+    from ray.data.lazy_dataset import LazyDataset
 
 import collections
 import itertools
@@ -2493,6 +2494,19 @@ def __iter__(self):
         it = Iterable(self._blocks, self._epoch)
         return DatasetPipeline(it, length=len(it._splits))
 
+    def lazy(self) -> "LazyDataset":
+        """
+        Convert this Dataset into a LazyDataset, where all subsequent
+        operations won't be applied until .compute() or a consuming
+        function (such as .iter_batches(), .to_torch(), etc.) is called.
+
+        Returns:
+            A LazyDataset.
+        """
+        from ray.data.lazy_dataset import LazyDataset
+
+        return LazyDataset(lambda: self)
+
     @DeveloperAPI
     def get_internal_block_refs(self) -> List[ObjectRef[Block]]:
         """Get a list of references to the underlying blocks of this dataset.

python/ray/data/dataset_pipeline.py

@@ -25,6 +25,7 @@
 
 if TYPE_CHECKING:
     import pyarrow
+    from ray.data.lazy_dataset import LazyDatasetPipeline
 
 # Operations that can be naively applied per dataset row in the pipeline.
 PER_DATASET_OPS = ["map", "map_batches", "flat_map", "filter"]
@@ -563,6 +564,19 @@ def __iter__(self):
 
         return EpochDelimitedIterator(self)
 
+    def lazy(self) -> "LazyDatasetPipeline":
+        """
+        Convert this DatasetPipeline into a LazyDatasetPipeline, where all subsequent
+        operations won't be applied until .compute() or a consuming
+        function (such as .iter_batches(), .to_torch(), etc.) is called.
+
+        Returns:
+            A LazyDatasetPipeline.
+        """
+        from ray.data.lazy_dataset import LazyDatasetPipeline
+
+        return LazyDatasetPipeline(lambda: self)
+
     @DeveloperAPI
     def iter_datasets(self) -> Iterator[Dataset[T]]:
         """Iterate over the output datasets of this pipeline.