[Datasets] Add logical operator for map()

python/ray/data/_internal/logical/operators/map_operator.py

@@ -7,13 +7,14 @@
 from ray.data._internal.execution.operators.map_operator import MapOperator
 from ray.data._internal.logical.interfaces import LogicalOperator
 from ray.data._internal.compute import (
+    UDF,
     get_compute,
     CallableClass,
     ComputeStrategy,
     TaskPoolStrategy,
     ActorPoolStrategy,
 )
-from ray.data.block import BatchUDF, Block
+from ray.data.block import BatchUDF, Block, RowUDF
 
 
 if sys.version_info >= (3, 8):
@@ -22,44 +23,99 @@
     from typing_extensions import Literal
 
 
-class MapBatches(LogicalOperator):
-    """Logical operator for map_batches."""
+class BaseMapLike(LogicalOperator):
+    """Abstract class for logical operators should be converted to physical
+    MapOperator.
+    """
 
     def __init__(
         self,
+        name: str,
         input_op: LogicalOperator,
         block_fn: BlockTransform,
-        fn: BatchUDF,
-        batch_size: Optional[Union[int, Literal["default"]]] = "default",
         compute: Optional[Union[str, ComputeStrategy]] = None,
-        batch_format: Literal["default", "pandas", "pyarrow", "numpy"] = "default",
-        zero_copy_batch: bool = False,
         target_block_size: Optional[int] = None,
+        fn: Optional[UDF] = None,
         fn_args: Optional[Iterable[Any]] = None,
         fn_kwargs: Optional[Dict[str, Any]] = None,
         fn_constructor_args: Optional[Iterable[Any]] = None,
         fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
         ray_remote_args: Optional[Dict[str, Any]] = None,
     ):
-        super().__init__("MapBatches", [input_op])
+        super().__init__(name, [input_op])
         self._block_fn = block_fn
-        self._fn = fn
-        self._batch_size = batch_size
         self._compute = compute or "tasks"
-        self._batch_format = batch_format
-        self._zero_copy_batch = zero_copy_batch
         self._target_block_size = target_block_size
+        self._fn = fn
         self._fn_args = fn_args
         self._fn_kwargs = fn_kwargs
         self._fn_constructor_args = fn_constructor_args
         self._fn_constructor_kwargs = fn_constructor_kwargs
         self._ray_remote_args = ray_remote_args or {}
 
 
-def plan_map_batches_op(
-    op: MapBatches, input_physical_dag: PhysicalOperator
-) -> PhysicalOperator:
-    """Get the corresponding DAG of physical operators for MapBatches."""
+class MapBatches(BaseMapLike):
+    """Logical operator for map_batches."""
+
+    def __init__(
+        self,
+        input_op: LogicalOperator,
+        block_fn: BlockTransform,
+        fn: BatchUDF,
+        batch_size: Optional[Union[int, Literal["default"]]] = "default",
+        compute: Optional[Union[str, ComputeStrategy]] = None,
+        batch_format: Literal["default", "pandas", "pyarrow", "numpy"] = "default",
+        zero_copy_batch: bool = False,
+        target_block_size: Optional[int] = None,
+        fn_args: Optional[Iterable[Any]] = None,
+        fn_kwargs: Optional[Dict[str, Any]] = None,
+        fn_constructor_args: Optional[Iterable[Any]] = None,
+        fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(
+            "MapBatches",
+            input_op,
+            block_fn,
+            compute=compute,
+            target_block_size=target_block_size,
+            fn=fn,
+            fn_args=fn_args,
+            fn_kwargs=fn_kwargs,
+            fn_constructor_args=fn_constructor_args,
+            fn_constructor_kwargs=fn_constructor_kwargs,
+            ray_remote_args=ray_remote_args,
+        )
+        self._batch_size = batch_size
+        self._batch_format = batch_format
+        self._zero_copy_batch = zero_copy_batch
+
+
+class Map(BaseMapLike):
+    """Logical operator for map."""
+
+    def __init__(
+        self,
+        input_op: LogicalOperator,
+        block_fn: BlockTransform,
+        fn: RowUDF,
+        compute: Optional[Union[str, ComputeStrategy]] = None,
+        ray_remote_args: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(
+            "Map",
+            input_op,
+            block_fn,
+            compute=compute,
+            fn=fn,
+            ray_remote_args=ray_remote_args,
+        )
+
+
+def plan_map_like_op(
+    op: BaseMapLike, input_physical_dag: PhysicalOperator
+) -> MapOperator:
+    """Get the corresponding physical operators DAG for map-like operators."""
     compute = get_compute(op._compute)
     block_fn = op._block_fn
 

python/ray/data/_internal/logical/planner.py

@@ -2,8 +2,8 @@
 from ray.data._internal.logical.interfaces import LogicalOperator
 from ray.data._internal.logical.operators.read_operator import Read, plan_read_op
 from ray.data._internal.logical.operators.map_operator import (
-    MapBatches,
-    plan_map_batches_op,
+    BaseMapLike,
+    plan_map_like_op,
 )
 
 
@@ -24,9 +24,9 @@ def plan(self, logical_dag: LogicalOperator) -> PhysicalOperator:
         if isinstance(logical_dag, Read):
             assert not physical_children
             physical_dag = plan_read_op(logical_dag)
-        elif isinstance(logical_dag, MapBatches):
+        elif isinstance(logical_dag, BaseMapLike):
             assert len(physical_children) == 1
-            physical_dag = plan_map_batches_op(logical_dag, physical_children[0])
+            physical_dag = plan_map_like_op(logical_dag, physical_children[0])
         else:
             raise ValueError(
                 f"Found unknown logical operator during planning: {logical_dag}"

python/ray/data/dataset.py

@@ -31,7 +31,10 @@
 from ray.air.constants import TENSOR_COLUMN_NAME
 from ray.air.util.data_batch_conversion import BlockFormat
 from ray.data._internal.logical.optimizers import LogicalPlan
-from ray.data._internal.logical.operators.map_operator import MapBatches
+from ray.data._internal.logical.operators.map_operator import (
+    Map,
+    MapBatches,
+)
 from ray.data.dataset_iterator import DatasetIterator
 from ray.data._internal.block_batching import batch_block_refs, batch_blocks
 from ray.data._internal.block_list import BlockList
@@ -334,7 +337,18 @@ def transform(block: Block, fn: RowUDF[T, U]) -> Iterable[Block]:
                 fn=fn,
             )
         )
-        return Dataset(plan, self._epoch, self._lazy)
+
+        logical_plan = self._logical_plan
+        if logical_plan is not None:
+            map_op = Map(
+                logical_plan.dag,
+                transform,
+                fn,
+                compute=compute,
+                ray_remote_args=ray_remote_args,
+            )
+            logical_plan = LogicalPlan(map_op)
+        return Dataset(plan, self._epoch, self._lazy, logical_plan)
 
     def map_batches(
         self,
@@ -676,16 +690,16 @@ def process_next_batch(batch: DataBatch) -> Iterator[Block]:
                 logical_plan.dag,
                 transform,
                 fn,
-                batch_size,
-                compute,
-                batch_format,
-                zero_copy_batch,
-                target_block_size,
-                fn_args,
-                fn_kwargs,
-                fn_constructor_args,
-                fn_constructor_kwargs,
-                ray_remote_args,
+                batch_size=batch_size,
+                compute=compute,
+                batch_format=batch_format,
+                zero_copy_batch=zero_copy_batch,
+                target_block_size=target_block_size,
+                fn_args=fn_args,
+                fn_kwargs=fn_kwargs,
+                fn_constructor_args=fn_constructor_args,
+                fn_constructor_kwargs=fn_constructor_kwargs,
+                ray_remote_args=ray_remote_args,
             )
             logical_plan = LogicalPlan(map_batches_op)
 

python/ray/data/tests/test_execution_optimizer.py

@@ -8,8 +8,8 @@ def test_e2e_optimizer_sanity(ray_start_cluster_enabled):
     ctx = DatasetContext.get_current()
     ctx.new_execution_backend = True
     ctx.optimizer_enabled = True
-    ds = ray.data.range(5).map_batches(lambda x: x)
-    assert ds.take_all() == [0, 1, 2, 3, 4], ds
+    ds = ray.data.range(5).map_batches(lambda x: x).map(lambda x: x + 1)
+    assert ds.take_all() == [1, 2, 3, 4, 5], ds
 
 
 if __name__ == "__main__":