diff --git a/ibis/common/graph.py b/ibis/common/graph.py index 31dc7829d5a8e..9921898ddfcb4 100644 --- a/ibis/common/graph.py +++ b/ibis/common/graph.py @@ -147,8 +147,9 @@ def fn(node, _, **kwargs): # need to first reconstruct the node from the possible rewritten # children, so we can match on the new node containing the rewritten # child arguments, this way we can propagate the rewritten nodes - # upward in the hierarchy - recreated = node.__class__(**kwargs) + # upward in the hierarchy, using a specialized __recreate__ method + # improves the performance by 17% compared node.__class__(**kwargs) + recreated = node.__recreate__(kwargs) if (result := obj.match(recreated, ctx)) is NoMatch: return recreated else: @@ -172,6 +173,11 @@ def fn(node, _, **kwargs): class Node(Hashable): __slots__ = () + @classmethod + def __recreate__(cls, kwargs: Any) -> Self: + """Reconstruct the node from the given arguments.""" + return cls(**kwargs) + @property @abstractmethod def __args__(self) -> tuple[Any, ...]: diff --git a/ibis/common/tests/test_graph_benchmarks.py b/ibis/common/tests/test_graph_benchmarks.py index cf9f9df52b7b3..45f63211a1e88 100644 --- a/ibis/common/tests/test_graph_benchmarks.py +++ b/ibis/common/tests/test_graph_benchmarks.py @@ -6,8 +6,10 @@ from typing_extensions import Self # noqa: TCH002 from ibis.common.collections import frozendict +from ibis.common.deferred import _ from ibis.common.graph import Graph, Node from ibis.common.grounds import Concrete +from ibis.common.patterns import Between, Object class MyNode(Concrete, Node): @@ -24,7 +26,7 @@ def generate_node(depth): if depth == 0: return MyNode(10, "20", c=(30, 40), d=frozendict(e=50, f=60)) return MyNode( - 1, + depth, "2", c=(3, 4), d=frozendict(e=5, f=6), @@ -48,3 +50,9 @@ def test_bfs(benchmark): def test_dfs(benchmark): node = generate_node(500) benchmark(Graph.from_dfs, node) + + +def test_replace(benchmark): + node = generate_node(500) + pattern = Object(MyNode, a=Between(lower=100)) >> _.copy(a=_.a + 1) + benchmark(node.replace, pattern)