[Data] Cleaning up Operator Fusion rule (#52566)

<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?

1. Minor clean up of `FuseOperators` rule
2. Abstracted operator fusion rules into `test_operator_fusion`

## Related issue number

<!-- For example: "Closes #1234" -->

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Author: alexeykudinkin

python/ray/data/BUILD

@@ -704,6 +704,20 @@ py_test(
     ],
 )
 
+py_test(
+    name = "test_operator_fusion",
+    size = "medium",
+    srcs = ["tests/test_operator_fusion.py"],
+    tags = [
+        "exclusive",
+        "team:data",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)
+
 py_test(
     name = "test_execution_optimizer",
     size = "medium",

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

@@ -13,7 +13,7 @@
 from ray.data._internal.logical.rules.inherit_target_max_block_size import (
     InheritTargetMaxBlockSizeRule,
 )
-from ray.data._internal.logical.rules.operator_fusion import OperatorFusionRule
+from ray.data._internal.logical.rules.operator_fusion import FuseOperators
 from ray.data._internal.logical.rules.randomize_blocks import ReorderRandomizeBlocksRule
 from ray.data._internal.logical.rules.set_read_parallelism import SetReadParallelismRule
 from ray.data._internal.logical.rules.zero_copy_map_fusion import (
@@ -35,7 +35,7 @@
     [
         InheritTargetMaxBlockSizeRule,
         SetReadParallelismRule,
-        OperatorFusionRule,
+        FuseOperators,
         EliminateBuildOutputBlocks,
         ConfigureMapTaskMemoryUsingOutputSize,
     ]

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

@@ -1,4 +1,4 @@
-from ray.data._internal.logical.rules.operator_fusion import OperatorFusionRule
+from ray.data._internal.logical.rules.operator_fusion import FuseOperators
 from ray.data._internal.logical.rules.randomize_blocks import ReorderRandomizeBlocksRule
 
-__all__ = ["ReorderRandomizeBlocksRule", "OperatorFusionRule"]
+__all__ = ["ReorderRandomizeBlocksRule", "FuseOperators"]

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

@@ -38,7 +38,7 @@
 INHERITABLE_REMOTE_ARGS = ["scheduling_strategy"]
 
 
-class OperatorFusionRule(Rule):
+class FuseOperators(Rule):
     """Fuses linear chains of compatible physical operators."""
 
     def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
@@ -54,21 +54,21 @@ def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
         # Update output dependencies after fusion.
         # TODO(hchen): Instead of updating the depdencies manually,
         # we need a better abstraction for manipulating the DAG.
-        self._remove_output_depes(fused_dag)
-        self._update_output_depes(fused_dag)
+        self._remove_output_deps(fused_dag)
+        self._update_output_deps(fused_dag)
 
         new_plan = PhysicalPlan(fused_dag, self._op_map, plan.context)
         return new_plan
 
-    def _remove_output_depes(self, op: PhysicalOperator) -> None:
+    def _remove_output_deps(self, op: PhysicalOperator) -> None:
         for input in op._input_dependencies:
             input._output_dependencies = []
-            self._remove_output_depes(input)
+            self._remove_output_deps(input)
 
-    def _update_output_depes(self, op: PhysicalOperator) -> None:
+    def _update_output_deps(self, op: PhysicalOperator) -> None:
         for input in op._input_dependencies:
             input._output_dependencies.append(op)
-            self._update_output_depes(input)
+            self._update_output_deps(input)
 
     def _fuse_map_operators_in_dag(self, dag: PhysicalOperator) -> MapOperator:
         """Starting at the given operator, traverses up the DAG of operators
@@ -305,6 +305,7 @@ def _get_fused_map_operator(
         # 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

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

@@ -10,7 +10,7 @@
 )
 from ray.data._internal.logical.interfaces.optimizer import Rule
 from ray.data._internal.logical.interfaces.physical_plan import PhysicalPlan
-from ray.data._internal.logical.rules.operator_fusion import OperatorFusionRule
+from ray.data._internal.logical.rules.operator_fusion import FuseOperators
 
 
 class ZeroCopyMapFusionRule(Rule):
@@ -27,7 +27,7 @@ class ZeroCopyMapFusionRule(Rule):
 
     @classmethod
     def dependencies(cls) -> List[Type[Rule]]:
-        return [OperatorFusionRule]
+        return [FuseOperators]
 
     def apply(self, plan: PhysicalPlan) -> PhysicalPlan:
         self._traverse(plan.dag)