chore: import upstream snapshot with attribution

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

@@ -0,0 +1,52 @@
+from typing import List, Type
+
+from .plan import Plan
+
+
+class Rule:
+    """Abstract class for optimization rule."""
+
+    def apply(self, plan: Plan) -> Plan:
+        """Apply the optimization rule to the execution plan."""
+        raise NotImplementedError
+
+    @classmethod
+    def dependencies(cls) -> List[Type["Rule"]]:
+        """List of rules that must be applied before this rule."""
+        return []
+
+    @classmethod
+    def dependents(cls) -> List[Type["Rule"]]:
+        """List of rules that must be applied after this rule."""
+        return []
+
+
+class Optimizer:
+    """Abstract class for optimizers.
+
+    An optimizers transforms a DAG of operators with a list of predefined rules.
+    """
+
+    @property
+    def rules(self) -> List[Rule]:
+        """List of predefined rules for this optimizer."""
+        raise NotImplementedError
+
+    def optimize(self, plan: Plan) -> Plan:
+        """Optimize operators with a list of rules."""
+        # Apply rules until the plan is not changed
+        previous_plan = plan
+        while True:
+            for rule in self.rules:
+                plan = rule.apply(plan)
+            # TODO: Eventually we should implement proper equality.
+            # Using str to check equality seems brittle
+            if plan.dag.dag_str == previous_plan.dag.dag_str:
+                break
+            previous_plan = plan
+        return self._post_optimize(plan)
+
+    def _post_optimize(self, plan: Plan) -> Plan:
+        """Post optimize is used for rules or other post-processing
+        that needs to be executed only once after the `optimize` loop."""
+        return plan