[Data] Restructure the user-provided optimizer rules file (#38829)

* Disable EliminateBuildOutputBlocks rule temporarily Signed-off-by: Cheng Su <[email protected]> * Change to use default and user-provided rules Signed-off-by: Cheng Su <[email protected]> * Rename to _user_provided_optimizer_rules.py Signed-off-by: Cheng Su <[email protected]> * Disable EliminateBuildOutputBlocks temporarily Signed-off-by: Cheng Su <[email protected]> --------- Signed-off-by: Cheng Su <[email protected]>
ray-project · Aug 25, 2023 · cea281c · cea281c
1 parent 35738b8
commit cea281c
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 22 deletions.
diff --git a/python/ray/data/_internal/logical/optimizers.py b/python/ray/data/_internal/logical/optimizers.py
@@ -6,27 +6,39 @@
  PhysicalPlan,
  Rule,
 )
-from ray.data._internal.logical.rules._default_optimizer_rules import (
- get_logical_optimizer_rules,
- get_physical_optimizer_rules,
+from ray.data._internal.logical.rules._user_provided_optimizer_rules import (
+ USER_PROVIDED_LOGICAL_RULES,
+ USER_PROVIDED_PHYSICAL_RULES,
 )
+from ray.data._internal.logical.rules.operator_fusion import OperatorFusionRule
+from ray.data._internal.logical.rules.randomize_blocks import ReorderRandomizeBlocksRule
 from ray.data._internal.planner.planner import Planner
 
+DEFAULT_LOGICAL_RULES = [
+ ReorderRandomizeBlocksRule,
+]
+
+DEFAULT_PHYSICAL_RULES = [
+ OperatorFusionRule,
+]
+
 
 class LogicalOptimizer(Optimizer):
  """The optimizer for logical operators."""
 
  @property
  def rules(self) -> List[Rule]:
- return [rule_cls() for rule_cls in get_logical_optimizer_rules()]
+ rules = DEFAULT_LOGICAL_RULES + USER_PROVIDED_LOGICAL_RULES
+ return [rule_cls() for rule_cls in rules]
 
 
 class PhysicalOptimizer(Optimizer):
  """The optimizer for physical operators."""
 
  @property
  def rules(self) -> List["Rule"]:
- return [rule_cls() for rule_cls in get_physical_optimizer_rules()]
+ rules = DEFAULT_PHYSICAL_RULES + USER_PROVIDED_PHYSICAL_RULES
+ return [rule_cls() for rule_cls in rules]
 
 
 def get_execution_plan(logical_plan: LogicalPlan) -> PhysicalPlan:

diff --git a/python/ray/data/_internal/logical/rules/_default_optimizer_rules.py b/python/ray/data/_internal/logical/rules/_default_optimizer_rules.py
diff --git a/python/ray/data/_internal/logical/rules/_user_provided_optimizer_rules.py b/python/ray/data/_internal/logical/rules/_user_provided_optimizer_rules.py
@@ -0,0 +1,7 @@
+# Users can provide extra logical optimization rules here
+# to be used in `LogicalOptimizer`.
+USER_PROVIDED_LOGICAL_RULES = []
+
+# Users can provide extra physical optimization rules here
+# to be used in `PhysicalOptimizer`.
+USER_PROVIDED_PHYSICAL_RULES = []
diff --git a/python/ray/data/tests/test_dynamic_block_split.py b/python/ray/data/tests/test_dynamic_block_split.py
@@ -325,6 +325,7 @@ def test_lazy_block_list(shutdown_only, target_max_block_size):
  assert block_metadata.schema is not None
 
 
+@pytest.mark.skip("Needs zero-copy optimization for read->map_batches.")
 def test_read_large_data(ray_start_cluster):
  # Test 20G input with single task
  num_blocks_per_task = 20

diff --git a/python/ray/data/tests/test_execution_optimizer.py b/python/ray/data/tests/test_execution_optimizer.py
@@ -1484,6 +1484,7 @@ def check_transform_fns(op, expected_types):
  assert isinstance(transform_fn, expected_types[i]), transform_fn
 
 
+@pytest.mark.skip("Needs zero-copy optimization for read->map_batches.")
 def test_zero_copy_fusion_eliminate_build_output_blocks(
  ray_start_regular_shared, enable_optimizer
 ):