chore: import upstream snapshot with attribution

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

@@ -0,0 +1,99 @@
+from typing import TYPE_CHECKING, List, Tuple
+
+from ray.data._internal.planner import create_planner
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.execution_callback import ExecutionCallback
+
+from .ruleset import Ruleset
+from ray.data._internal.logical.interfaces import (
+    LogicalPlan,
+    Optimizer,
+    PhysicalPlan,
+    Rule,
+)
+from ray.data._internal.logical.rules import (
+    CombineShuffles,
+    CommonSubExprElimination,
+    ConfigureMapTaskMemoryUsingOutputSize,
+    FuseOperators,
+    InheritTargetMaxBlockSizeRule,
+    LimitPushdownRule,
+    PredicatePushdown,
+    ProjectionPushdown,
+    SetReadParallelismRule,
+)
+from ray.util.annotations import DeveloperAPI
+
+_LOGICAL_RULESET = Ruleset(
+    [
+        LimitPushdownRule,
+        ProjectionPushdown,
+        PredicatePushdown,
+        CombineShuffles,
+    ]
+)
+
+_PHYSICAL_RULESET = Ruleset(
+    [
+        InheritTargetMaxBlockSizeRule,
+        SetReadParallelismRule,
+        FuseOperators,
+        ConfigureMapTaskMemoryUsingOutputSize,
+    ]
+)
+
+
+@DeveloperAPI
+def get_logical_ruleset() -> Ruleset:
+    return _LOGICAL_RULESET
+
+
+@DeveloperAPI
+def get_physical_ruleset() -> Ruleset:
+    return _PHYSICAL_RULESET
+
+
+class LogicalOptimizer(Optimizer):
+    """The optimizer for logical operators."""
+
+    @property
+    def rules(self) -> List[Rule]:
+        return [rule_cls() for rule_cls in get_logical_ruleset()]
+
+    def _post_optimize(self, plan: LogicalPlan) -> LogicalPlan:
+        # CommonSubExprElimination is only supposed to run once
+        # isolated from the optimizer rule loop as it applies to
+        # a single Projection operator not a chain of operators.
+        return CommonSubExprElimination().apply(plan)
+
+
+class PhysicalOptimizer(Optimizer):
+    """The optimizer for physical operators."""
+
+    @property
+    def rules(self) -> List[Rule]:
+        return [rule_cls() for rule_cls in get_physical_ruleset()]
+
+
+def get_execution_plan(
+    logical_plan: LogicalPlan,
+) -> Tuple[PhysicalPlan, List["ExecutionCallback"]]:
+    """Get the physical execution plan for the provided logical plan.
+
+    This process has 3 steps:
+    (1) logical optimization: optimize logical operators.
+    (2) planning: convert logical to physical operators.
+    (3) physical optimization: optimize physical operators.
+    """
+    # 1. Logical -> Logical (Optimized)
+    optimized_logical_plan = LogicalOptimizer().optimize(logical_plan)
+
+    # 2. Rewire Logical -> Logical (Optimized)
+    logical_plan._dag = optimized_logical_plan.dag
+
+    # 3. Logical (Optimized) -> Physical
+    physical_plan, callbacks = create_planner().plan(optimized_logical_plan)
+
+    # 4. Physical (Optimized) -> Physical
+    return PhysicalOptimizer().optimize(physical_plan), callbacks