benchmark

datafusion/common/Cargo.toml

@@ -54,6 +54,7 @@ arrow-array = { workspace = true }
 arrow-buffer = { workspace = true }
 arrow-schema = { workspace = true }
 chrono = { workspace = true }
+criterion = { version = "0.5", features = ["async_tokio"] }
 half = { workspace = true }
 hashbrown = { workspace = true }
 indexmap = { workspace = true }
@@ -72,3 +73,7 @@ instant = { version = "0.1", features = ["wasm-bindgen"] }
 
 [dev-dependencies]
 rand = { workspace = true }
+
+[[bench]]
+harness = false
+name = "tree_node"

datafusion/common/benches/tree_node.rs

@@ -0,0 +1,47 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+extern crate arrow;
+#[macro_use]
+extern crate criterion;
+
+use crate::criterion::Criterion;
+use datafusion_common::tree_node::tests::test_dyn_tree_node::{
+    tall_tree, visit_tree_new, visit_tree_old, wide_tree,
+};
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let tall_tree = tall_tree();
+    let wide_tree = wide_tree();
+
+    c.bench_function("visit old tall tree", |b| {
+        b.iter(|| visit_tree_old(tall_tree.clone()))
+    });
+    c.bench_function("visit new tall tree", |b| {
+        b.iter(|| visit_tree_new(tall_tree.clone()))
+    });
+
+    c.bench_function("visit old wide tree", |b| {
+        b.iter(|| visit_tree_old(wide_tree.clone()))
+    });
+    c.bench_function("visit new wide tree", |b| {
+        b.iter(|| visit_tree_new(wide_tree.clone()))
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

datafusion/common/src/tree_node.rs

@@ -1337,8 +1337,8 @@ impl<T: ConcreteTreeNode> TreeNode for T {
     }
 }
 
-#[cfg(test)]
-pub(crate) mod tests {
+// #[cfg(test)]
+pub mod tests {
     use crate::tree_node::{
         DynTreeNode, Transformed, TreeNode, TreeNodeIterator, TreeNodeRecursion,
         TreeNodeRewriter, TreeNodeVisitor,
@@ -2633,18 +2633,62 @@ pub(crate) mod tests {
 
         node_tests!(ArcTestNode);
 
-        #[test]
-        fn test_large_tree() {
+        pub fn tall_tree() -> Arc<DynTestNode<String>> {
             let mut item = ArcTestNode::new_leaf("initial".to_string());
             for i in 0..3000 {
                 item =
                     ArcTestNode::new_with_children(vec![item], format!("parent-{}", i));
             }
+            item
+        }
+
+        pub fn wide_tree() -> Arc<DynTestNode<String>> {
+            let mut item = ArcTestNode::new_leaf("initial".to_string());
+            for i in 0..22 {
+                item = ArcTestNode::new_with_children(
+                    vec![item.clone(), item],
+                    format!("parent-{}", i),
+                );
+            }
+            item
+        }
+
+        pub fn visit_tree_new(item: Arc<DynTestNode<String>>) {
+            let mut visitor =
+                TestVisitor::new(Box::new(visit_continue), Box::new(visit_continue));
+
+            item.visit_new(&mut visitor).unwrap();
+        }
 
+        pub fn visit_tree_old(item: Arc<DynTestNode<String>>) {
             let mut visitor =
                 TestVisitor::new(Box::new(visit_continue), Box::new(visit_continue));
 
             item.visit(&mut visitor).unwrap();
         }
+
+        #[test]
+        fn test_visit_tall_tree_new() {
+            let item = tall_tree();
+            visit_tree_new(item);
+        }
+
+        #[test]
+        fn test_visit_tall_tree_old() {
+            let item = tall_tree();
+            visit_tree_old(item);
+        }
+
+        #[test]
+        fn test_visit_wide_tree_new() {
+            let item = wide_tree();
+            visit_tree_new(item);
+        }
+
+        #[test]
+        fn test_visit_wide_tree_old() {
+            let item = wide_tree();
+            visit_tree_old(item);
+        }
     }
 }