chore: import upstream snapshot with attribution

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

@@ -0,0 +1,136 @@
+import collections
+from dataclasses import dataclass, field
+from typing import Dict, Iterator, List, Optional, Set, Tuple, Type
+
+from ray.data._internal.logical.interfaces import Rule
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class Ruleset:
+    """A collection of rules to apply to a plan.
+
+    This is a utility class to ensure that, if rules depend on each other, they're
+    applied in a correct order.
+    """
+
+    @dataclass(frozen=True)
+    class _Node:
+        rule: Type[Rule]
+        dependents: List["Ruleset._Node"] = field(default_factory=list)
+
+    def __init__(self, rules: Optional[List[Type[Rule]]] = None):
+        if rules is None:
+            rules = []
+
+        self._rules = list(rules)
+
+    def add(self, rule: Type[Rule]):
+        if rule in self._rules:
+            raise ValueError(f"Rule {rule} already in ruleset")
+
+        self._rules.append(rule)
+
+        if self._contains_cycle():
+            raise ValueError("Cannot add rule that would create a cycle")
+
+    def remove(self, rule: Type[Rule]):
+        if rule not in self._rules:
+            raise ValueError(f"Rule {rule} not found in ruleset")
+
+        self._rules.remove(rule)
+
+    def __iter__(self) -> Iterator[Type[Rule]]:
+        """Iterate over the rules in this ruleset.
+
+        This method yields rules in dependency order. For example, if B depends on A,
+        then this method yields A before B. Each rule is yielded exactly once, and a
+        rule is only yielded once *all* of its dependencies have been yielded (so a
+        rule that several others must precede is not emitted early or duplicated).
+        Insertion order breaks ties among rules that are ready at the same time.
+        """
+        order, _ = self._topological_order()
+        for node in order:
+            yield node.rule
+
+    def _topological_order(self) -> Tuple[List["Ruleset._Node"], int]:
+        """Order the nodes by dependency using Kahn's algorithm.
+
+        Returns the topologically-ordered nodes and the total node count.
+        A node is enqueued the moment its in-degree (count of not-yet-emitted
+        dependencies) hits zero; since an in-degree only decreases and we
+        enqueue solely on the zero-crossing, each node is emitted exactly once.
+        Insertion order breaks ties among nodes that are ready together.
+
+        Nodes that participate in a cycle never reach in-degree zero, so they
+        are absent from the result -- i.e. ``len(order) < total`` exactly when
+        the graph contains a cycle.
+        """
+        nodes, indegree = self._build_graph()
+        queue = collections.deque(n for n in nodes if indegree[id(n)] == 0)
+        order: List["Ruleset._Node"] = []
+        while queue:
+            node = queue.popleft()
+            order.append(node)
+            for dep in node.dependents:
+                indegree[id(dep)] -= 1
+                if indegree[id(dep)] == 0:
+                    queue.append(dep)
+        return order, len(nodes)
+
+    def _build_graph(
+        self,
+    ) -> Tuple[List["Ruleset._Node"], Dict[int, int]]:
+        """Build the dependency DAG.
+
+        Returns the nodes (one per rule, in insertion order) and their
+        in-degrees -- the number of rules that must be applied before each.
+        The in-degree map is keyed by node identity (``id``) rather than rule
+        type so that distinct nodes never share a counter, and is computed as
+        the edges are added (every incoming edge bumps the target's in-degree)
+        rather than re-derived by a second traversal. A node whose in-degree
+        is zero is a root.
+        """
+        rule_to_node: Dict[Type[Rule], "Ruleset._Node"] = {
+            rule: Ruleset._Node(rule) for rule in self._rules
+        }
+        indegree: Dict[int, int] = {id(node): 0 for node in rule_to_node.values()}
+
+        # De-duplicate edges. The same ordering can be declared from both ends
+        # -- rule A lists B in ``dependencies()`` while B lists A in
+        # ``dependents()`` -- which would otherwise add the edge twice,
+        # double-counting the in-degree and duplicating ``dependents`` entries.
+        seen_edges: Set[Tuple[int, int]] = set()
+
+        def add_edge(before: "Ruleset._Node", after: "Ruleset._Node") -> None:
+            """Record that ``before`` must be applied before ``after``."""
+            edge = (id(before), id(after))
+            if edge in seen_edges:
+                return
+            seen_edges.add(edge)
+            before.dependents.append(after)
+            indegree[id(after)] += 1
+
+        for rule in self._rules:
+            node = rule_to_node[rule]
+
+            # Rules that must be applied *before* this rule: dependency -> node.
+            for dependency in rule.dependencies():
+                if dependency in rule_to_node:
+                    add_edge(rule_to_node[dependency], node)
+
+            # Rules that must be applied *after* this rule: node -> dependent.
+            for dependent in rule.dependents():
+                if dependent in rule_to_node:
+                    add_edge(node, rule_to_node[dependent])
+
+        return list(rule_to_node.values()), indegree
+
+    def _contains_cycle(self) -> bool:
+        # Kahn's traversal drops any node stuck in a cycle (its in-degree never
+        # reaches zero), so a shortfall between the ordered nodes and the total
+        # means a cycle exists. This correctly flags a graph that mixes an
+        # acyclic component with a disjoint cycle -- a plain "does any root
+        # exist?" check would be fooled by the acyclic root.
+        order, total = self._topological_order()
+        return len(order) != total