[Data] Fixing FuseOperators rule to properly handle the case of transformations drastically changing size of the dataset (ray-project#52570)

These changes are needed to make sure `FuseOperators` is appropriately
handling potential impacts of transformations on the dataset sizes and
whether fusion should occur in that case.

For ex, consider following scenarios:

```
ds.filter(...).map_batches(..., bathc_size=1024)
```

Could not be fused as fusing it could potentially violate batching
semantic -- fused operator gonna first gonna collect 1024 rows, then
apply filter and subsequent transformation (which might expect exactly
1024 rows to be provided in a batch).

```
read_parquet(...).map_batches(..., bathc_size=1024)
```

Also could not be fused in that case, as fusing these 2 operations could
lead to drastic reduction of the parallelism of the read operation:
fused operator first gonna batch 1024 rows, then apply combined
read->map transformation.

This change:

1. Cleans up `can_modify_num_rows` method
2. Makes sure `Read` overrides `can_modify_num_rows` as well
3. Avoids fusion with ops that could be drastically modifying dataset
size
4. Cleaning up `FuseOperators` rule
5. Adding telemetry

---------

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>
Signed-off-by: zhaoch23 <c233zhao@uwaterloo.ca>

Author: alexeykudinkin

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

@@ -564,6 +564,10 @@ def __init__(self, min_rows_per_bundle: Optional[int]):
                 bundle up to this target, but only exceed it if not doing so would
                 result in an empty bundle.
         """
+        assert (
+            min_rows_per_bundle is None or min_rows_per_bundle >= 0
+        ), "Min rows per bundle has to be non-negative"
+
         self._min_rows_per_bundle = min_rows_per_bundle
         self._bundle_buffer: List[RefBundle] = []
         self._bundle_buffer_size = 0

python/ray/data/_internal/logical/interfaces/logical_operator.py

@@ -82,3 +82,7 @@ def is_lineage_serializable(self) -> bool:
         objects aren't available on the deserialized machine.
         """
         return True
+
+    @classmethod
+    def is_read_op(cls):
+        return False

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

@@ -37,8 +37,8 @@ def __init__(
                 inspecting the logical plan of a Dataset.
             input_op: The operator preceding this operator in the plan DAG. The outputs
                 of `input_op` will be the inputs to this operator.
-            min_rows_per_bundled_input: The target number of rows to pass to
-                ``MapOperator._add_bundled_input()``.
+            min_rows_per_bundled_input: Min number of rows a single bundle of blocks
+                passed on to the task must possess.
             ray_remote_args: Args to provide to :func:`ray.remote`.
             ray_remote_args_fn: A function that returns a dictionary of remote args
                 passed to each map worker. The purpose of this argument is to generate
@@ -177,7 +177,6 @@ def __init__(
         self._batch_format = batch_format
         self._zero_copy_batch = zero_copy_batch
 
-    @property
     def can_modify_num_rows(self) -> bool:
         return False
 
@@ -210,7 +209,6 @@ def __init__(
             ray_remote_args=ray_remote_args,
         )
 
-    @property
     def can_modify_num_rows(self) -> bool:
         return False
 
@@ -249,7 +247,6 @@ def __init__(
             ray_remote_args=ray_remote_args,
         )
 
-    @property
     def can_modify_num_rows(self) -> bool:
         return True
 
@@ -285,7 +282,6 @@ def cols(self) -> Optional[List[str]]:
     def cols_rename(self) -> Optional[Dict[str, str]]:
         return self._cols_rename
 
-    @property
     def can_modify_num_rows(self) -> bool:
         return False
 
@@ -318,7 +314,6 @@ def __init__(
             ray_remote_args=ray_remote_args,
         )
 
-    @property
     def can_modify_num_rows(self) -> bool:
         return True
 
@@ -342,6 +337,5 @@ def __init__(
     def target_num_rows_per_block(self) -> int:
         return self._target_num_rows_per_block
 
-    @property
     def can_modify_num_rows(self) -> bool:
         return False

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

@@ -1,4 +1,3 @@
-import abc
 from typing import Optional
 
 from ray.data._internal.logical.interfaces import LogicalOperator
@@ -29,11 +28,10 @@ def __init__(
     def input_dependency(self) -> LogicalOperator:
         return self._input_dependencies[0]
 
-    @property
-    @abc.abstractmethod
     def can_modify_num_rows(self) -> bool:
         """Whether this operator can modify the number of rows,
         i.e. number of input rows != number of output rows."""
+        ...
 
 
 class Limit(AbstractOneToOne):
@@ -50,7 +48,6 @@ def __init__(
         )
         self._limit = limit
 
-    @property
     def can_modify_num_rows(self) -> bool:
         return True
 

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

@@ -93,3 +93,12 @@ def _cached_output_metadata(self) -> BlockMetadata:
             input_files=input_files,
             exec_stats=None,
         )
+
+    @classmethod
+    def is_read_op(cls):
+        return True
+
+    def can_modify_num_rows(self) -> bool:
+        # NOTE: Returns true, since most of the readers expands its input
+        #       and produce many rows for every single row of the input
+        return True

python/ray/data/_internal/logical/rules/limit_pushdown.py

@@ -61,7 +61,7 @@ def _apply_limit_pushdown(self, op: LogicalOperator) -> LogicalOperator:
                 while (
                     isinstance(new_input_into_limit, AbstractOneToOne)
                     and not isinstance(new_input_into_limit, Read)
-                    and not getattr(new_input_into_limit, "can_modify_num_rows", False)
+                    and not new_input_into_limit.can_modify_num_rows()
                 ):
                     new_input_into_limit_copy = copy.copy(new_input_into_limit)
                     ops_between_new_input_and_limit.append(new_input_into_limit_copy)

python/ray/data/_internal/logical/rules/operator_fusion.py

@@ -1,4 +1,5 @@
 import itertools
+import logging
 from typing import List, Optional, Tuple
 
 from ray.data._internal.compute import (
@@ -38,6 +39,9 @@
 INHERITABLE_REMOTE_ARGS = ["scheduling_strategy"]
 
 
+logger = logging.getLogger(__name__)
+
+
 class FuseOperators(Rule):
     """Fuses linear chains of compatible physical operators."""
 
@@ -175,35 +179,22 @@ def _can_fuse(self, down_op: PhysicalOperator, up_op: PhysicalOperator) -> bool:
             (
                 isinstance(up_logical_op, AbstractMap)
                 and isinstance(down_logical_op, AbstractMap)
+                and self._can_fuse_map_ops(up_logical_op, down_logical_op)
             )
             or (
                 isinstance(up_logical_op, AbstractMap)
                 and isinstance(down_logical_op, RandomShuffle)
             )
+            # Do not fuse Repartition operator if shuffle is disabled
+            # (i.e. using split shuffle).
             or (
                 isinstance(up_logical_op, AbstractMap)
                 and isinstance(down_logical_op, Repartition)
+                and down_logical_op._shuffle
             )
         ):
             return False
 
-        # Do not fuse Repartition operator if shuffle is disabled
-        # (i.e. using split shuffle).
-        if isinstance(down_logical_op, Repartition) and not down_logical_op._shuffle:
-            return False
-
-        if isinstance(down_logical_op, AbstractMap) and isinstance(
-            up_logical_op, AbstractMap
-        ):
-            if (
-                self._fuse_compute_strategy(
-                    up_logical_op._compute,
-                    down_logical_op._compute,
-                )
-                is None
-            ):
-                return False
-
         # Only fuse if the ops' remote arguments are compatible.
         if not _are_remote_args_compatible(
             getattr(up_logical_op, "_ray_remote_args", {}),
@@ -228,8 +219,9 @@ def _can_fuse(self, down_op: PhysicalOperator, up_op: PhysicalOperator) -> bool:
         # Otherwise, ops are compatible for fusion.
         return True
 
+    @classmethod
     def _fuse_compute_strategy(
-        self, up_compute: ComputeStrategy, down_compute: ComputeStrategy
+        cls, up_compute: ComputeStrategy, down_compute: ComputeStrategy
     ) -> Optional[ComputeStrategy]:
         """Fuse the compute strategies of the upstream and downstream operators.
         Returns None if they are not compatible.
@@ -302,21 +294,10 @@ def _get_fused_map_operator(
         assert isinstance(down_logical_op, AbstractMap)
         assert isinstance(up_logical_op, AbstractMap)
 
-        # Merge minimum block sizes.
-        down_min_rows_per_bundled_input = down_logical_op._min_rows_per_bundled_input
-        up_min_rows_per_bundled_input = up_logical_op._min_rows_per_bundled_input
-
-        if (
-            down_min_rows_per_bundled_input is not None
-            and up_min_rows_per_bundled_input is not None
-        ):
-            min_rows_per_bundled_input = max(
-                down_min_rows_per_bundled_input, up_min_rows_per_bundled_input
-            )
-        elif up_min_rows_per_bundled_input is not None:
-            min_rows_per_bundled_input = up_min_rows_per_bundled_input
-        else:
-            min_rows_per_bundled_input = down_min_rows_per_bundled_input
+        # Derive min num rows per input bundle
+        min_rows_per_bundled_input = self._derive_bundle_min_num_rows(
+            down_logical_op, up_logical_op
+        )
 
         target_max_block_size = self._get_merged_target_max_block_size(
             up_op.target_max_block_size, down_op.target_max_block_size
@@ -389,6 +370,27 @@ def _get_fused_map_operator(
         # Return the fused physical operator.
         return op
 
+    @classmethod
+    def _derive_bundle_min_num_rows(
+        cls,
+        down_logical_op: AbstractMap,
+        up_logical_op: AbstractMap,
+    ) -> Optional[int]:
+        us_bundle_min_rows_req = up_logical_op._min_rows_per_bundled_input
+        ds_bundle_min_rows_req = down_logical_op._min_rows_per_bundled_input
+
+        # In case neither of the ops specify `min_rows_per_bundled_input`,
+        # return None
+        if us_bundle_min_rows_req is None and ds_bundle_min_rows_req is None:
+            return None
+
+        # Target min bundle size is selected as max of upstream and downstream ones
+        # such that it could satisfy both of their requirements
+        return max(
+            ds_bundle_min_rows_req or 0,
+            us_bundle_min_rows_req or 0,
+        )
+
     def _get_fused_all_to_all_operator(
         self, down_op: AllToAllOperator, up_op: MapOperator
     ) -> AllToAllOperator:
@@ -460,6 +462,50 @@ def fused_all_to_all_transform_fn(
         # Return the fused physical operator.
         return op
 
+    @classmethod
+    def _can_fuse_map_ops(
+        cls,
+        upstream_op: AbstractMap,
+        downstream_op: AbstractMap,
+    ) -> bool:
+        if (
+            cls._fuse_compute_strategy(
+                upstream_op._compute,
+                downstream_op._compute,
+            )
+            is None
+        ):
+            return False
+
+        # Do not fuse Map operators in case:
+        #
+        #   - Upstream could (potentially) drastically modify number of rows, while
+        #   - Downstream has `min_rows_per_input_bundle` specified
+        #
+        # Fusing such transformations is not desirable as it could
+        #
+        #   - Drastically reduce parallelism for the upstream up (for ex, if
+        #   fusing ``Read->MapBatches(batch_size=...)`` with large enough batch-size
+        #   could drastically reduce parallelism level of the Read op)
+        #
+        #   - Potentially violate batching semantic by fusing
+        #   ``Filter->MapBatches(batch_size=...)``
+        #
+        if (
+            upstream_op.can_modify_num_rows()
+            and downstream_op._min_rows_per_bundled_input is not None
+        ):
+            logger.debug(
+                f"Upstream operator '{upstream_op}' could be modifying # of input "
+                f"rows, while downstream operator '{downstream_op}' expects at least "
+                f"{downstream_op._min_rows_per_bundled_input} rows in a batch. "
+                f"Skipping fusion"
+            )
+
+            return False
+
+        return True
+
 
 def _are_remote_args_compatible(prev_args, next_args):
     """Check if Ray remote arguments are compatible for merging."""

python/ray/data/_internal/plan.py

@@ -10,8 +10,6 @@
 from ray.data._internal.execution.interfaces import RefBundle
 from ray.data._internal.logical.interfaces.logical_operator import LogicalOperator
 from ray.data._internal.logical.interfaces.logical_plan import LogicalPlan
-from ray.data._internal.logical.operators.from_operators import AbstractFrom
-from ray.data._internal.logical.operators.input_data_operator import InputData
 from ray.data._internal.logical.operators.read_operator import Read
 from ray.data._internal.stats import DatasetStats
 from ray.data._internal.util import unify_block_metadata_schema
@@ -136,7 +134,7 @@ def generate_logical_plan_string(
             ):
                 """Traverse (DFS) the LogicalPlan DAG and
                 return a string representation of the operators."""
-                if isinstance(op, (Read, InputData, AbstractFrom)):
+                if not op.input_dependencies or op.is_read_op():
                     return curr_str, depth
 
                 curr_max_depth = depth
@@ -168,19 +166,27 @@ def generate_logical_plan_string(
                 count = self._snapshot_metadata.num_rows
             else:
                 # This plan hasn't executed any operators.
-                sources = self._logical_plan.sources()
+                has_n_ary_operator = False
+                dag = self._logical_plan.dag
+
+                while not dag.is_read_op() and dag.input_dependencies:
+                    if len(dag.input_dependencies) > 1:
+                        has_n_ary_operator = True
+                        break
+
+                    dag = dag.input_dependencies[0]
+
                 # TODO(@bveeramani): Handle schemas for n-ary operators like `Union`.
-                if len(sources) > 1:
-                    # Multiple sources, cannot determine schema.
+                if has_n_ary_operator:
                     schema = None
                     count = None
                 else:
-                    assert len(sources) == 1
+                    assert dag.is_read_op() or not dag.input_dependencies, dag
                     plan = ExecutionPlan(
                         DatasetStats(metadata={}, parent=None),
                         self._context,
                     )
-                    plan.link_logical_plan(LogicalPlan(sources[0], plan._context))
+                    plan.link_logical_plan(LogicalPlan(dag, plan._context))
                     schema = plan.schema()
                     count = plan.meta_count()
         else:
@@ -587,7 +593,8 @@ def is_read_only(self, root_op: Optional[LogicalOperator] = None) -> bool:
         the LogicalPlan is used."""
         if root_op is None:
             root_op = self._logical_plan.dag
-        return isinstance(root_op, Read) and len(root_op.input_dependencies) == 0
+
+        return root_op.is_read_op()
 
     def has_computed_output(self) -> bool:
         """Whether this plan has a computed snapshot for the final operator, i.e. for