refactor(common): unify the node.find() and node.match() methods to transparently support types and patterns

Author: kszucs

ibis/common/graph.py

@@ -162,38 +162,22 @@ def map(self, fn: Callable, filter: Optional[Any] = None) -> dict[Node, Any]:
             results[node] = fn(node, results, **kwargs)
         return results
 
-    def find(self, type: type | tuple[type], filter: Optional[Any] = None) -> set[Node]:
-        """Find all nodes of a given type in the graph.
-
-        Parameters
-        ----------
-        type
-            Type or tuple of types to find.
-        filter
-            Pattern-like object to filter out nodes from the traversal. The traversal
-            will only visit nodes that match the given pattern and stop otherwise.
-
-        Returns
-        -------
-        The set of nodes matching the given type.
-        """
-        nodes = Graph.from_bfs(self, filter=filter).nodes()
-        return {node for node in nodes if isinstance(node, type)}
-
-    @experimental
-    def match(
-        self, pat: Any, filter: Optional[Any] = None, context: Optional[dict] = None
-    ) -> set[Node]:
-        """Find all nodes matching a given pattern in the graph.
-
-        A more advanced version of find, this method allows to match nodes based on
-        the more flexible pattern matching system implemented in the pattern module.
+    def find(
+        self,
+        pat: type | tuple[type],
+        filter: Optional[Any] = None,
+        context: Optional[dict] = None,
+    ) -> list[Node]:
+        """Find all nodes matching a given pattern or type in the graph.
+
+        Allow to match nodes based on the flexible pattern matching system implemented
+        in the pattern module, but also provide a fast path for matching based on the
+        type of the node.
 
         Parameters
         ----------
         pat
-            Pattern to match. `ibis.common.pattern()` function is used to coerce the
-            input value into a pattern. See the pattern module for more details.
+            Python type or `Pattern` to match.
         filter
             Pattern-like object to filter out nodes from the traversal. The traversal
             will only visit nodes that match the given pattern and stop otherwise.
@@ -202,12 +186,16 @@ def match(
 
         Returns
         -------
-        The set of nodes matching the given pattern.
+        The list of nodes matching the given pattern. The order of the nodes is
+        determined by a breadth-first search.
         """
-        pat = pattern(pat)
-        ctx = context or {}
         nodes = Graph.from_bfs(self, filter=filter).nodes()
-        return {node for node in nodes if pat.match(node, ctx) is not NoMatch}
+        if isinstance(pat, (tuple, type)):
+            return [node for node in nodes if isinstance(node, pat)]
+        else:
+            pat = pattern(pat)
+            ctx = context or {}
+            return [node for node in nodes if pat.match(node, ctx) is not NoMatch]
 
     @experimental
     def replace(

ibis/common/tests/test_graph.py

@@ -309,12 +309,38 @@ def test_recursive_lookup():
     )
 
 
-def test_node_match():
-    result = A.match(If(_.name == "C"))
-    assert result == {C}
+def test_node_find_using_type():
+    class FooNode(MyNode):
+        pass
+
+    class BarNode(MyNode):
+        pass
+
+    C = BarNode(name="C", children=[])
+    D = FooNode(name="D", children=[])
+    E = BarNode(name="E", children=[])
+    B = FooNode(name="B", children=[D, E])
+    A = MyNode(name="A", children=[B, C])
+
+    result = A.find(MyNode)
+    assert result == [A, B, C, D, E]
+
+    result = A.find(FooNode)
+    assert result == [B, D]
+
+    result = A.find(BarNode)
+    assert result == [C, E]
+
+    result = A.find((FooNode, BarNode))
+    assert result == [B, C, D, E]
+
+
+def test_node_find_using_pattern():
+    result = A.find(If(_.name == "C"))
+    assert result == [C]
 
-    result = A.match(Object(MyNode, name=Eq("D")))
-    assert result == {D}
+    result = A.find(Object(MyNode, name=Eq("D")))
+    assert result == [D]
 
-    result = A.match(If(_.children))
-    assert result == {A, B}
+    result = A.find(If(_.children))
+    assert result == [A, B]

ibis/expr/analysis.py

@@ -151,7 +151,7 @@ def pushdown_selection_filters(parent, predicates):
 
     simplified = []
     for pred in predicates:
-        if pred.match(conflicting_projection, filter=p.Value):
+        if pred.find(conflicting_projection, filter=p.Value):
             return default
         try:
             simplified.append(pred.replace(pushdown_pattern))