use std::fmt::Display;
use std::fmt::Formatter;
use std::ops::Deref;
use crate::expr::Expression;
use crate::scalar_fn::ScalarFnRef;
pub enum DisplayFormat {
Compact,
Tree,
}
pub struct DisplayTreeExpr<'a>(pub &'a Expression);
impl Display for DisplayTreeExpr<'_> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
pub use termtree::Tree;
fn make_tree(expr: &Expression) -> Result<Tree<String>, std::fmt::Error> {
let scalar_fn: &ScalarFnRef = expr.deref();
let node_name = format!("{}", scalar_fn);
let child_names = (0..expr.children().len()).map(|i| expr.signature().child_name(i));
let children = expr.children();
let child_trees: Result<Vec<Tree<String>>, std::fmt::Error> = children
.iter()
.zip(child_names)
.map(|(child, name)| {
let child_tree = make_tree(child)?;
Ok::<Tree<String>, std::fmt::Error>(
Tree::new(format!("{}: {}", name, child_tree.root))
.with_leaves(child_tree.leaves),
)
})
.collect();
Ok(Tree::new(node_name).with_leaves(child_trees?))
}
write!(f, "{}", make_tree(self.0)?)
}
}
#[cfg(test)]
mod tests {
use crate::dtype::DType;
use crate::dtype::Nullability;
use crate::dtype::PType;
use crate::expr::and;
use crate::expr::between;
use crate::expr::cast;
use crate::expr::eq;
use crate::expr::get_item;
use crate::expr::gt;
use crate::expr::lit;
use crate::expr::not;
use crate::expr::pack;
use crate::expr::root;
use crate::expr::select;
use crate::expr::select_exclude;
use crate::scalar_fn::fns::between::BetweenOptions;
use crate::scalar_fn::fns::between::StrictComparison;
#[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() {
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_root() {
use insta::assert_snapshot;
let root_expr = root();
assert_snapshot!(root_expr.display_tree().to_string(), @"vortex.root()");
}
#[test]
fn test_display_tree_literal() {
use insta::assert_snapshot;
let lit_expr = lit(42);
assert_snapshot!(lit_expr.display_tree().to_string(), @"vortex.literal(42i32)");
}
#[test]
fn test_display_tree_get_item() {
use insta::assert_snapshot;
let get_item_expr = get_item("my_field", root());
assert_snapshot!(get_item_expr.display_tree().to_string(), @r"
vortex.get_item(my_field)
└── input: vortex.root()
");
}
#[test]
fn test_display_tree_binary() {
use insta::assert_snapshot;
let binary_expr = gt(get_item("x", root()), lit(10));
assert_snapshot!(binary_expr.display_tree().to_string(), @r"
vortex.binary(>)
├── lhs: vortex.get_item(x)
│ └── input: vortex.root()
└── rhs: vortex.literal(10i32)
");
}
#[test]
fn test_display_tree_complex_binary() {
use insta::assert_snapshot;
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#"
vortex.binary(and)
├── lhs: vortex.binary(=)
│ ├── lhs: vortex.get_item(name)
│ │ └── input: vortex.root()
│ └── rhs: vortex.literal("alice")
└── rhs: vortex.binary(>)
├── lhs: vortex.get_item(age)
│ └── input: vortex.root()
└── rhs: vortex.literal(18i32)
"#);
}
#[test]
fn test_display_tree_select() {
use insta::assert_snapshot;
let select_expr = select(["name", "age"], root());
assert_snapshot!(select_expr.display_tree().to_string(), @r"
vortex.select({name, age})
└── child: vortex.root()
");
}
#[test]
fn test_display_tree_select_exclude() {
use insta::assert_snapshot;
let select_exclude_expr = select_exclude(["internal_id", "metadata"], root());
assert_snapshot!(select_exclude_expr.display_tree().to_string(), @r"
vortex.select(~{internal_id, metadata})
└── child: vortex.root()
");
}
#[test]
fn test_display_tree_cast() {
use insta::assert_snapshot;
let cast_expr = cast(
get_item("value", root()),
DType::Primitive(PType::I64, Nullability::NonNullable),
);
assert_snapshot!(cast_expr.display_tree().to_string(), @r"
vortex.cast(i64)
└── input: vortex.get_item(value)
└── input: vortex.root()
");
}
#[test]
fn test_display_tree_not() {
use insta::assert_snapshot;
let not_expr = not(eq(get_item("active", root()), lit(true)));
assert_snapshot!(not_expr.display_tree().to_string(), @r"
vortex.not()
└── input: vortex.binary(=)
├── lhs: vortex.get_item(active)
│ └── input: vortex.root()
└── rhs: vortex.literal(true)
");
}
#[test]
fn test_display_tree_between() {
use insta::assert_snapshot;
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"
vortex.between(lower_strict: <=, upper_strict: <=)
├── array: vortex.get_item(score)
│ └── input: vortex.root()
├── lower: vortex.literal(0i32)
└── upper: vortex.literal(100i32)
");
}
#[test]
fn test_display_tree_nested() {
use insta::assert_snapshot;
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"
vortex.select({result})
└── child: vortex.cast(bool)
└── input: vortex.between(lower_strict: <, upper_strict: <=)
├── array: vortex.get_item(score)
│ └── input: vortex.root()
├── lower: vortex.literal(50i32)
└── upper: vortex.literal(100i32)
");
}
#[test]
fn test_display_tree_pack() {
use insta::assert_snapshot;
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"
vortex.select({fizz, buzz})
└── child: vortex.pack(names: [fizz, bar, buzz], nullability: Nullable)
├── fizz: vortex.root()
├── bar: vortex.literal(5i32)
└── buzz: vortex.binary(=)
├── lhs: vortex.literal(42i32)
└── rhs: vortex.get_item(answer)
└── input: vortex.root()
");
}
}