Pretty tree display for Vortex expressions (#4373)

My main motivations here are:
1. Display that is easier to understand quickly, especially in deeply
nested cases.
2. Easier to diff when using `insta`.

I've made it hidden behind a feature, but not sure that's actually
desirable here, it adds one small dependency and it seems unlikely that
the code would be misused in some way that will really bomb perf.

---------

Signed-off-by: Adam Gutglick <[email protected]>

Author: AdamGS

Cargo.lock

@@ -114,6 +114,7 @@ half = { version = "2.6", features = ["std", "num-traits"] }
 hashbrown = "0.15.1"
 humansize = "2.1.3"
 indicatif = "0.18.0"
+insta = "1.43"
 inventory = "0.3.20"
 itertools = "0.14.0"
 jiff = "0.2.0"
@@ -180,6 +181,7 @@ take_mut = "0.2.2"
 tar = "0.4"
 target-lexicon = "0.13"
 tempfile = "3"
+termtree = { version = "0.5" }
 thiserror = "2.0.3"
 tokio = { version = "1.46" }
 tokio-stream = "0.1.17"

Cargo.toml

@@ -39,7 +39,7 @@ impl Dataset for StructListOfInts {
 
     async fn to_vortex_array(&self) -> ArrayRef {
         let names: FieldNames = (0..self.num_columns)
-            .map(|col_idx| (col_idx.to_string()))
+            .map(|col_idx| col_idx.to_string())
             .collect();
         let mut rng = rand::rngs::StdRng::seed_from_u64(0);
 

bench-vortex/src/datasets/struct_list_of_ints.rs

@@ -25,6 +25,7 @@ parking_lot = { workspace = true }
 paste = { workspace = true }
 prost = { workspace = true }
 serde = { workspace = true, optional = true, features = ["derive"] }
+termtree = { workspace = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
@@ -36,6 +37,7 @@ vortex-utils = { workspace = true }
 
 [dev-dependencies]
 anyhow = { workspace = true }
+insta = { workspace = true }
 rstest = { workspace = true }
 vortex-expr = { path = ".", features = ["test-harness"] }
 

vortex-expr/Cargo.toml

@@ -0,0 +1,238 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+pub enum DisplayFormat {
+    Compact,
+    Tree,
+}
+
+/// Configurable display trait for expressions.
+pub trait DisplayAs {
+    fn fmt_as(&self, df: DisplayFormat, f: &mut std::fmt::Formatter) -> std::fmt::Result;
+
+    fn child_names(&self) -> Option<Vec<String>> {
+        None
+    }
+}
+
+pub struct DisplayTreeExpr<'a>(pub &'a dyn crate::VortexExpr);
+
+impl std::fmt::Display for DisplayTreeExpr<'_> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        pub use termtree::Tree;
+        fn make_tree(expr: &dyn crate::VortexExpr) -> Result<Tree<String>, std::fmt::Error> {
+            let node_name = TreeNodeDisplay(expr).to_string();
+
+            // Get child names for display purposes
+            let child_names = DisplayAs::child_names(expr);
+            let children = expr.children();
+
+            let child_trees: Result<Vec<Tree<String>>, _> = if let Some(names) = child_names
+                && names.len() == children.len()
+            {
+                children
+                    .iter()
+                    .zip(names.iter())
+                    .map(|(child, name)| {
+                        let child_tree = make_tree(child.as_ref())?;
+                        Ok(Tree::new(format!("{}: {}", name, child_tree.root))
+                            .with_leaves(child_tree.leaves))
+                    })
+                    .collect()
+            } else {
+                children
+                    .iter()
+                    .map(|child| make_tree(child.as_ref()))
+                    .collect()
+            };
+
+            Ok(Tree::new(node_name).with_leaves(child_trees?))
+        }
+
+        write!(f, "{}", make_tree(self.0)?)
+    }
+}
+
+struct TreeNodeDisplay<'a>(&'a dyn crate::VortexExpr);
+
+impl<'a> std::fmt::Display for TreeNodeDisplay<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt_as(DisplayFormat::Tree, f)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::compute::{BetweenOptions, StrictComparison};
+    use vortex_dtype::{DType, Nullability, PType};
+
+    use crate::{and, between, cast, eq, get_item, gt, lit, not, root, select};
+
+    #[test]
+    fn tree_display_getitem() {
+        let expr = get_item("x", root());
+        println!("{}", expr.display_tree());
+    }
+
+    #[test]
+    fn tree_display_binary() {
+        let expr = gt(get_item("x", root()), lit(5));
+        println!("{}", expr.display_tree());
+    }
+
+    #[test]
+    fn test_child_names_debug() {
+        // Simple test to debug child names display
+        let binary_expr = gt(get_item("x", root()), lit(10));
+        println!("Binary expr tree:\n{}", binary_expr.display_tree());
+
+        let between_expr = between(
+            get_item("score", root()),
+            lit(0),
+            lit(100),
+            BetweenOptions {
+                lower_strict: StrictComparison::NonStrict,
+                upper_strict: StrictComparison::NonStrict,
+            },
+        );
+        println!("Between expr tree:\n{}", between_expr.display_tree());
+    }
+
+    #[test]
+    fn test_display_tree() {
+        use insta::assert_snapshot;
+
+        use crate::{pack, select_exclude};
+        let root_expr = root();
+        assert_snapshot!(root_expr.display_tree().to_string(), @"Root");
+
+        let lit_expr = lit(42);
+        assert_snapshot!(lit_expr.display_tree().to_string(), @"Literal(value: 42i32, dtype: i32)");
+
+        let get_item_expr = get_item("my_field", root());
+        assert_snapshot!(get_item_expr.display_tree().to_string(), @r"
+        GetItem(my_field)
+        └── Root
+        ");
+
+        let binary_expr = gt(get_item("x", root()), lit(10));
+        assert_snapshot!(binary_expr.display_tree().to_string(), @r"
+        Binary(>)
+        ├── lhs: GetItem(x)
+        │   └── Root
+        └── rhs: Literal(value: 10i32, dtype: i32)
+        ");
+
+        let complex_binary = and(
+            eq(get_item("name", root()), lit("alice")),
+            gt(get_item("age", root()), lit(18)),
+        );
+        assert_snapshot!(complex_binary.display_tree().to_string(), @r#"
+        Binary(and)
+        ├── lhs: Binary(=)
+        │   ├── lhs: GetItem(name)
+        │   │   └── Root
+        │   └── rhs: Literal(value: "alice", dtype: utf8)
+        └── rhs: Binary(>)
+            ├── lhs: GetItem(age)
+            │   └── Root
+            └── rhs: Literal(value: 18i32, dtype: i32)
+        "#);
+
+        let select_expr = select(["name", "age"], root());
+        assert_snapshot!(select_expr.display_tree().to_string(), @r#"
+        Select(include): ["name", "age"]
+        └── Root
+        "#);
+
+        let select_exclude_expr = select_exclude(["internal_id", "metadata"], root());
+        assert_snapshot!(select_exclude_expr.display_tree().to_string(), @r#"
+        Select(exclude): ["internal_id", "metadata"]
+        └── Root
+        "#);
+
+        let cast_expr = cast(
+            get_item("value", root()),
+            DType::Primitive(PType::I64, Nullability::NonNullable),
+        );
+        assert_snapshot!(cast_expr.display_tree().to_string(), @r"
+        Cast(target: i64)
+        └── GetItem(value)
+            └── Root
+        ");
+
+        let not_expr = not(eq(get_item("active", root()), lit(true)));
+        assert_snapshot!(not_expr.display_tree().to_string(), @r"
+        Not
+        └── Binary(=)
+            ├── lhs: GetItem(active)
+            │   └── Root
+            └── rhs: Literal(value: true, dtype: bool)
+        ");
+
+        let between_expr = between(
+            get_item("score", root()),
+            lit(0),
+            lit(100),
+            BetweenOptions {
+                lower_strict: StrictComparison::NonStrict,
+                upper_strict: StrictComparison::NonStrict,
+            },
+        );
+        assert_snapshot!(between_expr.display_tree().to_string(), @r"
+        Between
+        ├── array: GetItem(score)
+        │   └── Root
+        ├── lower (NonStrict): Literal(value: 0i32, dtype: i32)
+        └── upper (NonStrict): Literal(value: 100i32, dtype: i32)
+        ");
+
+        // Test nested expression
+        let nested_expr = select(
+            ["result"],
+            cast(
+                between(
+                    get_item("score", root()),
+                    lit(50),
+                    lit(100),
+                    BetweenOptions {
+                        lower_strict: StrictComparison::Strict,
+                        upper_strict: StrictComparison::NonStrict,
+                    },
+                ),
+                DType::Bool(Nullability::NonNullable),
+            ),
+        );
+        assert_snapshot!(nested_expr.display_tree().to_string(), @r#"
+        Select(include): ["result"]
+        └── Cast(target: bool)
+            └── Between
+                ├── array: GetItem(score)
+                │   └── Root
+                ├── lower (Strict): Literal(value: 50i32, dtype: i32)
+                └── upper (NonStrict): Literal(value: 100i32, dtype: i32)
+        "#);
+
+        let select_from_pack_expr = select(
+            ["fizz", "buzz"],
+            pack(
+                [
+                    ("fizz", root()),
+                    ("bar", lit(5)),
+                    ("buzz", eq(lit(42), get_item("answer", root()))),
+                ],
+                Nullability::Nullable,
+            ),
+        );
+        assert_snapshot!(select_from_pack_expr.display_tree().to_string(), @r#"
+        Select(include): ["fizz", "buzz"]
+        └── Pack
+            ├── fizz: Root
+            ├── bar: Literal(value: 5i32, dtype: i32)
+            └── buzz: Binary(=)
+                ├── lhs: Literal(value: 42i32, dtype: i32)
+                └── rhs: GetItem(answer)
+                    └── Root
+        "#);
+    }
+}