feat: teach VortexExpr to Display (#1293)

This was helpful when debugging the pruner. Seems generally good to have. The Debug output of these expressions is way too big to quickly digest.
spiraldb · Nov 13, 2024 · a24819a · a24819a
1 parent e22beed
commit a24819a
Show file tree

Hide file tree

Showing 9 changed files with 253 additions and 6 deletions.
diff --git a/vortex-expr/src/binary.rs b/vortex-expr/src/binary.rs
@@ -1,4 +1,5 @@
 use std::any::Any;
+use std::fmt::Display;
 use std::sync::Arc;
 
 use vortex_array::aliases::hash_set::HashSet;
@@ -34,6 +35,12 @@ impl BinaryExpr {
     }
 }
 
+impl Display for BinaryExpr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "({} {} {})", self.lhs, self.operator, self.rhs)
+    }
+}
+
 impl VortexExpr for BinaryExpr {
     fn as_any(&self) -> &dyn Any {
         self

diff --git a/vortex-expr/src/column.rs b/vortex-expr/src/column.rs
@@ -1,4 +1,5 @@
 use std::any::Any;
+use std::fmt::Display;
 
 use vortex_array::aliases::hash_set::HashSet;
 use vortex_array::array::StructArray;
@@ -36,6 +37,12 @@ impl From<usize> for Column {
     }
 }
 
+impl Display for Column {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.field)
+    }
+}
+
 impl VortexExpr for Column {
     fn as_any(&self) -> &dyn Any {
         self

diff --git a/vortex-expr/src/identity.rs b/vortex-expr/src/identity.rs
@@ -1,4 +1,5 @@
 use std::any::Any;
+use std::fmt::Display;
 
 use vortex_array::Array;
 use vortex_error::VortexResult;
@@ -8,6 +9,12 @@ use crate::{unbox_any, VortexExpr};
 #[derive(Debug, Eq, PartialEq)]
 pub struct Identity;
 
+impl Display for Identity {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "[]")
+    }
+}
+
 impl VortexExpr for Identity {
     fn as_any(&self) -> &dyn Any {
         self

diff --git a/vortex-expr/src/lib.rs b/vortex-expr/src/lib.rs
@@ -1,5 +1,5 @@
 use std::any::Any;
-use std::fmt::Debug;
+use std::fmt::{Debug, Display};
 use std::sync::Arc;
 
 use vortex_array::aliases::hash_set::HashSet;
@@ -25,7 +25,7 @@ use vortex_dtype::field::Field;
 use vortex_error::{VortexExpect, VortexResult};
 
 /// Represents logical operation on [`Array`]s
-pub trait VortexExpr: Debug + Send + Sync + PartialEq<dyn Any> {
+pub trait VortexExpr: Debug + Send + Sync + PartialEq<dyn Any> + Display {
     /// Convert expression reference to reference of [`Any`] type
     fn as_any(&self) -> &dyn Any;
 
@@ -77,9 +77,56 @@ pub fn unbox_any(any: &dyn Any) -> &dyn Any {
     }
 }
 
+pub fn join_write<I, T: Display>(
+    f: &mut std::fmt::Formatter<'_>,
+    start: &str,
+    values: I,
+    delimiter: &str,
+    end: &str,
+) -> std::fmt::Result
+where
+    I: IntoIterator<Item = T>,
+{
+    write!(f, "{}", start)?;
+    let mut first = true;
+    for value in values {
+        if !first {
+            write!(f, "{}", delimiter)?;
+        }
+        first = false;
+        write!(f, "{}", value)?;
+    }
+    write!(f, "{}", end)
+}
+
+pub fn join_write_fun<F, I>(
+    f: &mut std::fmt::Formatter<'_>,
+    start: &str,
+    delimiter: &str,
+    end: &str,
+    formatters: I,
+) -> std::fmt::Result
+where
+    F: FnOnce(&mut std::fmt::Formatter<'_>) -> std::fmt::Result,
+    I: IntoIterator<Item = F>,
+{
+    write!(f, "{}", start)?;
+    let mut first = true;
+    for formatter in formatters {
+        if !first {
+            write!(f, "{}", delimiter)?;
+        }
+        first = false;
+        formatter(f)?;
+    }
+    write!(f, "{}", end)
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_dtype::field::Field;
+    use vortex_dtype::{DType, Nullability, PType, StructDType};
+    use vortex_scalar::{Scalar, ScalarValue};
 
     use super::*;
 
@@ -100,4 +147,114 @@ mod tests {
         let conjunction = split_conjunction(&expr);
         assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}");
     }
+
+    #[test]
+    fn expr_display() {
+        assert_eq!(Column::new(Field::Name("a".to_string())).to_string(), "$a");
+        assert_eq!(Column::new(Field::Index(1)).to_string(), "[1]");
+        assert_eq!(Identity.to_string(), "[]");
+        assert_eq!(Identity.to_string(), "[]");
+
+        let col1: Arc<dyn VortexExpr> = Arc::new(Column::new(Field::Name("col1".to_string())));
+        let col2: Arc<dyn VortexExpr> = Arc::new(Column::new(Field::Name("col2".to_string())));
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::And, col2.clone()).to_string(),
+            "($col1 and $col2)"
+        );
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::Or, col2.clone()).to_string(),
+            "($col1 or $col2)"
+        );
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::Eq, col2.clone()).to_string(),
+            "($col1 = $col2)"
+        );
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::NotEq, col2.clone()).to_string(),
+            "($col1 != $col2)"
+        );
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::Gt, col2.clone()).to_string(),
+            "($col1 > $col2)"
+        );
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::Gte, col2.clone()).to_string(),
+            "($col1 >= $col2)"
+        );
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::Lt, col2.clone()).to_string(),
+            "($col1 < $col2)"
+        );
+        assert_eq!(
+            BinaryExpr::new(col1.clone(), Operator::Lte, col2.clone()).to_string(),
+            "($col1 <= $col2)"
+        );
+
+        assert_eq!(
+            BinaryExpr::new(
+                Arc::new(BinaryExpr::new(col1.clone(), Operator::Lt, col2.clone())),
+                Operator::Or,
+                Arc::new(BinaryExpr::new(col1.clone(), Operator::NotEq, col2.clone()))
+            )
+            .to_string(),
+            "(($col1 < $col2) or ($col1 != $col2))"
+        );
+
+        assert_eq!(Not::new(col1.clone()).to_string(), "!$col1");
+
+        assert_eq!(
+            Select::include(vec![Field::Name("col1".to_string())]).to_string(),
+            "Include($col1)"
+        );
+        assert_eq!(
+            Select::include(vec![
+                Field::Name("col1".to_string()),
+                Field::Name("col2".to_string())
+            ])
+            .to_string(),
+            "Include($col1,$col2)"
+        );
+        assert_eq!(
+            Select::exclude(vec![
+                Field::Name("col1".to_string()),
+                Field::Name("col2".to_string()),
+                Field::Index(1),
+            ])
+            .to_string(),
+            "Exclude($col1,$col2,[1])"
+        );
+
+        assert_eq!(Literal::new(Scalar::from(0_u8)).to_string(), "0_u8");
+        assert_eq!(Literal::new(Scalar::from(0.0_f32)).to_string(), "0_f32");
+        assert_eq!(
+            Literal::new(Scalar::from(i64::MAX)).to_string(),
+            "9223372036854775807_i64"
+        );
+        assert_eq!(Literal::new(Scalar::from(true)).to_string(), "true");
+        assert_eq!(
+            Literal::new(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
+            "null"
+        );
+
+        assert_eq!(
+            Literal::new(Scalar::new(
+                DType::Struct(
+                    StructDType::new(
+                        Arc::from([Arc::from("dog"), Arc::from("cat")]),
+                        vec![
+                            DType::Primitive(PType::U32, Nullability::NonNullable),
+                            DType::Utf8(Nullability::NonNullable)
+                        ],
+                    ),
+                    Nullability::NonNullable
+                ),
+                ScalarValue::List(Arc::from([
+                    ScalarValue::from(32_u32),
+                    ScalarValue::from("rufus".to_string())
+                ]))
+            ))
+            .to_string(),
+            "{dog:32_u32,cat:rufus}"
+        );
+    }
 }
diff --git a/vortex-expr/src/literal.rs b/vortex-expr/src/literal.rs
@@ -1,4 +1,5 @@
 use std::any::Any;
+use std::fmt::Display;
 
 use vortex_array::array::ConstantArray;
 use vortex_array::{Array, IntoArray};
@@ -22,6 +23,12 @@ impl Literal {
     }
 }
 
+impl Display for Literal {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.value)
+    }
+}
+
 impl VortexExpr for Literal {
     fn as_any(&self) -> &dyn Any {
         self

diff --git a/vortex-expr/src/not.rs b/vortex-expr/src/not.rs
@@ -1,4 +1,5 @@
 use std::any::Any;
+use std::fmt::Display;
 use std::sync::Arc;
 
 use vortex_array::aliases::hash_set::HashSet;
@@ -23,6 +24,13 @@ impl Not {
     }
 }
 
+impl Display for Not {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "!")?;
+        self.child.fmt(f)
+    }
+}
+
 impl VortexExpr for Not {
     fn as_any(&self) -> &dyn Any {
         self

diff --git a/vortex-expr/src/select.rs b/vortex-expr/src/select.rs
@@ -1,11 +1,12 @@
 use std::any::Any;
+use std::fmt::Display;
 
 use vortex_array::aliases::hash_set::HashSet;
 use vortex_array::Array;
 use vortex_dtype::field::Field;
 use vortex_error::{vortex_err, VortexResult};
 
-use crate::{unbox_any, VortexExpr};
+use crate::{join_write, unbox_any, VortexExpr};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Select {
@@ -23,6 +24,15 @@ impl Select {
     }
 }
 
+impl Display for Select {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Select::Include(fields) => join_write(f, "Include(", fields, ",", ")"),
+            Select::Exclude(fields) => join_write(f, "Exclude(", fields, ",", ")"),
+        }
+    }
+}
+
 impl VortexExpr for Select {
     fn as_any(&self) -> &dyn Any {
         self

diff --git a/vortex-serde/src/file/pruning.rs b/vortex-serde/src/file/pruning.rs
@@ -1,14 +1,17 @@
 // This code doesn't have usage outside of tests yet, remove once usage is added
 #![allow(dead_code)]
 
+use std::fmt::Display;
 use std::sync::Arc;
 
 use vortex_array::aliases::hash_map::HashMap;
 use vortex_array::aliases::hash_set::HashSet;
 use vortex_array::stats::Stat;
 use vortex_dtype::field::Field;
 use vortex_dtype::Nullability;
-use vortex_expr::{BinaryExpr, Column, Literal, Not, Operator, VortexExpr};
+use vortex_expr::{
+    join_write, join_write_fun, BinaryExpr, Column, Literal, Not, Operator, VortexExpr,
+};
 use vortex_scalar::Scalar;
 
 #[derive(Debug, Clone)]
@@ -17,6 +20,25 @@ pub struct PruningPredicate {
     required_stats: HashMap<Field, HashSet<Stat>>,
 }
 
+impl Display for PruningPredicate {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "PruningPredicate({}, ", self.expr)?;
+        join_write_fun(
+            f,
+            "{",
+            ",",
+            "}",
+            self.required_stats.iter().map(|(k, v)| {
+                move |f: &mut std::fmt::Formatter<'_>| {
+                    write!(f, "{}: ", k)?;
+                    join_write(f, "{", v, ",", "}")
+                }
+            }),
+        )?;
+        write!(f, ")")
+    }
+}
+
 impl PruningPredicate {
     pub fn try_new(original_expr: &Arc<dyn VortexExpr>) -> Option<Self> {
         let (expr, required_stats) = convert_to_pruning_expression(original_expr);
@@ -526,4 +548,20 @@ mod tests {
         )))) as _;
         assert!(PruningPredicate::try_new(&or_expr).is_none());
     }
+
+    #[test]
+    fn display_pruning_predicate() {
+        let column = Field::from("a");
+        let other_col = Arc::new(Literal::new(42.into()));
+        let not_eq_expr = Arc::new(BinaryExpr::new(
+            Arc::new(Column::new(column.clone())),
+            Operator::Lt,
+            other_col.clone(),
+        )) as _;
+
+        assert_eq!(
+            PruningPredicate::try_new(&not_eq_expr).unwrap().to_string(),
+            "PruningPredicate(($a_min >= 42_i32), {$a: {min}})"
+        );
+    }
 }