perf(common): improve the performance of replacing nodes using mappin…

…gs (#8638)
ibis-project · Mar 13, 2024 · a2e733a · a2e733a
1 parent 1a94644
commit a2e733a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/ibis/common/graph.py b/ibis/common/graph.py
@@ -193,7 +193,7 @@ def fn(node, _, **kwargs):
             try:
                 return obj[node]
             except KeyError:
-                return node.__class__(**kwargs)
+                return node.__recreate__(kwargs)
     elif callable(obj):
         fn = obj
     else:

diff --git a/ibis/common/tests/test_graph_benchmarks.py b/ibis/common/tests/test_graph_benchmarks.py
@@ -52,7 +52,13 @@ def test_dfs(benchmark):
     benchmark(Graph.from_dfs, node)
 
 
-def test_replace(benchmark):
+def test_replace_pattern(benchmark):
     node = generate_node(500)
     pattern = Object(MyNode, a=Between(lower=100)) >> _.copy(a=_.a + 1)
     benchmark(node.replace, pattern)
+
+
+def test_replace_mapping(benchmark):
+    node = generate_node(500)
+    subs = {generate_node(1): generate_node(0)}
+    benchmark(node.replace, subs)