Expressions

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/expr/bound.rs

@@ -7,15 +7,18 @@ use std::fmt::Display;
 use std::fmt::Formatter;
 use std::hash::Hash;
 use std::hash::Hasher;
+use std::ops::Deref;
 
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
 use vortex_utils::debug_with::DebugWith;
 use vortex_vector::Datum;
 
+use crate::ArrayRef;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::ExprVTable;
+use crate::expr::Expression;
 use crate::expr::VTable;
 use crate::expr::options::ExpressionOptions;
 use crate::expr::signature::ExpressionSignature;
@@ -67,36 +70,44 @@ impl BoundExpression {
     pub fn options(&self) -> ExpressionOptions<'_> {
         ExpressionOptions {
             vtable: &self.vtable,
-            options: self.options.as_ref(),
+            options: self.options.deref(),
         }
     }
 
     /// Signature information for this expression.
     pub fn signature(&self) -> ExpressionSignature<'_> {
         ExpressionSignature {
             vtable: &self.vtable,
-            options: self.options.as_ref(),
+            options: self.options.deref(),
         }
     }
 
     /// Compute the return [`DType`] of this expression given the input argument types.
     pub fn return_dtype(&self, arg_types: &[DType]) -> VortexResult<DType> {
         self.vtable
             .as_dyn()
-            .return_dtype(self.options.as_ref(), arg_types)
+            .return_dtype(self.options.deref(), arg_types)
+    }
+
+    /// Evaluate the expression, returning an ArrayRef.
+    ///
+    /// NOTE: this function will soon be deprecated as all expressions will evaluate trivially
+    ///  into an ExprArray.
+    pub fn evaluate(&self, expr: &Expression, scope: &ArrayRef) -> VortexResult<ArrayRef> {
+        self.vtable.as_dyn().evaluate(expr, scope)
     }
 
     /// Execute the expression given the input arguments.
     pub fn execute(&self, ctx: ExecutionArgs) -> VortexResult<Datum> {
-        self.vtable.as_dyn().execute(self.options.as_ref(), ctx)
+        self.vtable.as_dyn().execute(self.options.deref(), ctx)
     }
 }
 
 impl Clone for BoundExpression {
     fn clone(&self) -> Self {
         BoundExpression {
             vtable: self.vtable.clone(),
-            options: self.vtable.as_dyn().options_clone(self.options.as_ref()),
+            options: self.vtable.as_dyn().options_clone(self.options.deref()),
         }
     }
 }
@@ -110,7 +121,7 @@ impl Debug for BoundExpression {
                 &DebugWith(|fmt| {
                     self.vtable
                         .as_dyn()
-                        .options_debug(self.options.as_ref(), fmt)
+                        .options_debug(self.options.deref(), fmt)
                 }),
             )
             .finish()
@@ -122,7 +133,7 @@ impl Display for BoundExpression {
         write!(f, "{}(", self.vtable.id())?;
         self.vtable
             .as_dyn()
-            .options_display(self.options.as_ref(), f)?;
+            .options_display(self.options.deref(), f)?;
         write!(f, ")")
     }
 }
@@ -133,7 +144,7 @@ impl PartialEq for BoundExpression {
             && self
                 .vtable
                 .as_dyn()
-                .options_eq(self.options.as_ref(), other.options.as_ref())
+                .options_eq(self.options.deref(), other.options.deref())
     }
 }
 impl Eq for BoundExpression {}
@@ -143,6 +154,6 @@ impl Hash for BoundExpression {
         self.vtable.hash(state);
         self.vtable
             .as_dyn()
-            .options_hash(self.options.as_ref(), state);
+            .options_hash(self.options.deref(), state);
     }
 }

vortex-array/src/expr/display.rs

@@ -3,7 +3,9 @@
 
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::ops::Deref;
 
+use crate::expr::BoundExpression;
 use crate::expr::Expression;
 
 pub enum DisplayFormat {
@@ -17,7 +19,8 @@ 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 node_name = format!("{}", expr);
+            let bound: &BoundExpression = expr.deref();
+            let node_name = format!("{}", bound);
 
             // Get child names for display purposes
             let child_names = (0..expr.children().len()).map(|i| expr.signature().child_name(i));
@@ -96,10 +99,10 @@ mod tests {
         use insta::assert_snapshot;
 
         let root_expr = root();
-        assert_snapshot!(root_expr.display_tree().to_string(), @"vortex.root");
+        assert_snapshot!(root_expr.display_tree().to_string(), @"vortex.root()");
 
         let lit_expr = lit(42);
-        assert_snapshot!(lit_expr.display_tree().to_string(), @"vortex.literal 42i32");
+        assert_snapshot!(lit_expr.display_tree().to_string(), @"vortex.literal(42i32)");
 
         let get_item_expr = get_item("my_field", root());
         assert_snapshot!(get_item_expr.display_tree().to_string(), @r#"
@@ -123,24 +126,24 @@ mod tests {
         vortex.binary and
         ├── lhs: vortex.binary =
         │   ├── lhs: vortex.get_item "name"
-        │   │   └── input: vortex.root
+        │   │   └── input: $
         │   └── rhs: vortex.literal "alice"
         └── rhs: vortex.binary >
             ├── lhs: vortex.get_item "age"
-            │   └── input: vortex.root
+            │   └── input: $
             └── rhs: vortex.literal 18i32
         "#);
 
         let select_expr = select(["name", "age"], root());
         assert_snapshot!(select_expr.display_tree().to_string(), @r"
         vortex.select include={name, age}
-        └── child: vortex.root
+        └── child: $
         ");
 
         let select_exclude_expr = select_exclude(["internal_id", "metadata"], root());
         assert_snapshot!(select_exclude_expr.display_tree().to_string(), @r"
         vortex.select exclude={internal_id, metadata}
-        └── child: vortex.root
+        └── child: $
         ");
 
         let cast_expr = cast(
@@ -150,7 +153,7 @@ mod tests {
         assert_snapshot!(cast_expr.display_tree().to_string(), @r#"
         vortex.cast i64
         └── input: vortex.get_item "value"
-            └── input: vortex.root
+            └── input: $
         "#);
 
         let not_expr = not(eq(get_item("active", root()), lit(true)));

vortex-array/src/expr/expression.rs

@@ -16,6 +16,7 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 
 use crate::ArrayRef;
+use crate::expr::Root;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::bound::BoundExpression;
@@ -101,6 +102,10 @@ impl Expression {
 
     /// Computes the return dtype of this expression given the input dtype.
     pub fn return_dtype(&self, scope: &DType) -> VortexResult<DType> {
+        if self.is::<Root>() {
+            return Ok(scope.clone());
+        }
+
         let dtypes: Vec<_> = self
             .children
             .iter()
@@ -110,8 +115,11 @@ impl Expression {
     }
 
     /// Evaluates the expression in the given scope, returning an array.
-    pub fn evaluate(&self, _scope: &ArrayRef) -> VortexResult<ArrayRef> {
-        todo!("Return an ExprArray")
+    pub fn evaluate(&self, scope: &ArrayRef) -> VortexResult<ArrayRef> {
+        if self.is::<Root>() {
+            return Ok(scope.clone());
+        }
+        self.bound.evaluate(self, scope)
     }
 
     /// An expression over zone-statistics which implies all records in the zone evaluate to false.

vortex-array/src/expr/exprs/between.rs

@@ -117,14 +117,14 @@ impl VTable for Between {
         let lower_dt = &arg_dtypes[1];
         let upper_dt = &arg_dtypes[2];
 
-        if !arr_dt.eq_ignore_nullability(&lower_dt) {
+        if !arr_dt.eq_ignore_nullability(lower_dt) {
             vortex_bail!(
                 "Array dtype {} does not match lower dtype {}",
                 arr_dt,
                 lower_dt
             );
         }
-        if !arr_dt.eq_ignore_nullability(&upper_dt) {
+        if !arr_dt.eq_ignore_nullability(upper_dt) {
             vortex_bail!(
                 "Array dtype {} does not match upper dtype {}",
                 arr_dt,

vortex-array/src/expr/exprs/binary.rs

@@ -52,7 +52,7 @@ impl VTable for Binary {
 
     fn deserialize(&self, metadata: &[u8]) -> VortexResult<Self::Options> {
         let opts = pb::BinaryOpts::decode(metadata)?;
-        Ok(Operator::try_from(opts.op)?)
+        Operator::try_from(opts.op)
     }
 
     fn arity(&self, _options: &Self::Options) -> Arity {
@@ -85,7 +85,7 @@ impl VTable for Binary {
         let rhs = &arg_dtypes[1];
 
         if operator.is_arithmetic() {
-            if lhs.is_primitive() && lhs.eq_ignore_nullability(&rhs) {
+            if lhs.is_primitive() && lhs.eq_ignore_nullability(rhs) {
                 return Ok(lhs.with_nullability(lhs.nullability() | rhs.nullability()));
             }
             vortex_bail!(
@@ -601,15 +601,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn test_debug_print() {
-        let expr = gt(lit(1), lit(2));
-        assert_eq!(
-            format!("{expr:?}"),
-            "Expression { vtable: vortex.binary, data: >, children: [Expression { vtable: vortex.literal, data: 1i32, children: [] }, Expression { vtable: vortex.literal, data: 2i32, children: [] }] }"
-        );
-    }
-
     #[test]
     fn test_display_print() {
         let expr = gt(lit(1), lit(2));

vortex-array/src/expr/exprs/dynamic.rs

@@ -79,7 +79,7 @@ impl VTable for DynamicComparison {
         arg_dtypes: &[DType],
     ) -> VortexResult<DType> {
         let lhs = &arg_dtypes[0];
-        if !dynamic.rhs.dtype.eq_ignore_nullability(&lhs) {
+        if !dynamic.rhs.dtype.eq_ignore_nullability(lhs) {
             vortex_bail!(
                 "Incompatible dtypes for dynamic comparison: expected {} (ignore nullability) but got {}",
                 &dynamic.rhs.dtype,

vortex-array/src/expr/exprs/literal.rs

@@ -47,11 +47,10 @@ impl VTable for Literal {
 
     fn deserialize(&self, metadata: &[u8]) -> VortexResult<Self::Options> {
         let ops = pb::LiteralOpts::decode(metadata)?;
-        Ok(ops
-            .value
+        ops.value
             .as_ref()
             .ok_or_else(|| vortex_err!("Literal metadata missing value"))?
-            .try_into()?)
+            .try_into()
     }
 
     fn arity(&self, _options: &Self::Options) -> Arity {
@@ -154,8 +153,8 @@ impl VTable for Literal {
 ///
 /// let number = lit(34i32);
 ///
-/// let literal = number.as_::<Literal>();
-/// assert_eq!(literal.data(), &Scalar::primitive(34i32, Nullability::NonNullable));
+/// let scalar = number.as_::<Literal>();
+/// assert_eq!(scalar, &Scalar::primitive(34i32, Nullability::NonNullable));
 /// ```
 pub fn lit(value: impl Into<Scalar>) -> Expression {
     Literal.new_expr(value.into(), [])

vortex-array/src/expr/mod.rs

@@ -92,7 +92,7 @@ pub struct ExactExpr(pub Expression);
 impl PartialEq for ExactExpr {
     fn eq(&self, other: &Self) -> bool {
         self.0.id() == other.0.id()
-            && ptr::addr_eq(&*self.0.options().as_any(), &*other.0.options().as_any())
+            && ptr::addr_eq(self.0.options().as_any(), other.0.options().as_any())
     }
 }
 impl Eq for ExactExpr {}

vortex-array/src/expr/simplify.rs

@@ -22,7 +22,7 @@ impl Expression {
             let children: Vec<_> = expr
                 .children()
                 .iter()
-                .map(|c| inner(c, &cache))
+                .map(|c| inner(c, cache))
                 .try_collect()?;
 
             if children.iter().any(|c| c.is_some()) {
@@ -35,11 +35,18 @@ impl Expression {
 
                 let new_expr = expr.clone().with_children(new_children)?;
 
-                // Then we simplify the new expression
-                new_expr.vtable().as_dyn().simplify(&new_expr, cache)
+                // Then we simplify the new expression, and since we rewrote the expression we must
+                // always return a new expression (even if simplification returns None)
+                Ok(Some(
+                    new_expr
+                        .vtable()
+                        .as_dyn()
+                        .simplify(&new_expr, cache)?
+                        .unwrap_or(new_expr),
+                ))
             } else {
                 // Otherwise, we attempt to simplify the current expression
-                expr.vtable().as_dyn().simplify(&expr, cache)
+                expr.vtable().as_dyn().simplify(expr, cache)
             }
         }
 

vortex-array/src/expr/vtable.rs

@@ -338,7 +338,7 @@ impl<V: VTable> DynExprVTable for VTableAdapter<V> {
     }
 
     fn options_deserialize(&self, bytes: &[u8]) -> VortexResult<Box<dyn Any + Send + Sync>> {
-        Ok(Box::new(V::deserialize(&self.0, bytes)))
+        Ok(Box::new(V::deserialize(&self.0, bytes)?))
     }
 
     fn options_clone(&self, options: &dyn Any) -> Box<dyn Any + Send + Sync> {
@@ -357,11 +357,11 @@ impl<V: VTable> DynExprVTable for VTableAdapter<V> {
     }
 
     fn options_display(&self, options: &dyn Any, fmt: &mut Formatter<'_>) -> fmt::Result {
-        write!(fmt, "{}", downcast::<V>(options))
+        Display::fmt(downcast::<V>(options), fmt)
     }
 
     fn options_debug(&self, options: &dyn Any, fmt: &mut Formatter<'_>) -> fmt::Result {
-        write!(fmt, "{:?}", downcast::<V>(options))
+        Debug::fmt(downcast::<V>(options), fmt)
     }
 
     fn return_dtype(&self, options: &dyn Any, arg_dtypes: &[DType]) -> VortexResult<DType> {