Expressions

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

Author: gatesn

vortex-array/src/arrays/scalar_fn/array.rs

@@ -8,7 +8,8 @@ use vortex_error::vortex_ensure;
 use crate::Array;
 use crate::ArrayRef;
 use crate::arrays::ScalarFnVTable;
-use crate::expr::BoundExpression;
+use crate::expr::ReturnDTypeCtx;
+use crate::expr::ScalarFn;
 use crate::stats::ArrayStats;
 use crate::vtable::ArrayVTable;
 use crate::vtable::ArrayVTableExt;
@@ -17,7 +18,7 @@ use crate::vtable::ArrayVTableExt;
 pub struct ScalarFnArray {
     // NOTE(ngates): we should fix vtables so we don't have to hold this
     pub(super) vtable: ArrayVTable,
-    pub(super) bound: BoundExpression,
+    pub(super) bound: ScalarFn,
     pub(super) dtype: DType,
     pub(super) len: usize,
     pub(super) children: Vec<ArrayRef>,
@@ -26,13 +27,8 @@ pub struct ScalarFnArray {
 
 impl ScalarFnArray {
     /// Create a new ScalarFnArray from a scalar function and its children.
-    pub fn try_new(
-        bound: BoundExpression,
-        children: Vec<ArrayRef>,
-        len: usize,
-    ) -> VortexResult<Self> {
-        let arg_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
-        let dtype = bound.return_dtype(&arg_dtypes)?;
+    pub fn try_new(bound: ScalarFn, children: Vec<ArrayRef>, len: usize) -> VortexResult<Self> {
+        let dtype = bound.return_dtype(&ChildArraysReturnDTypeCtx(&children))?;
 
         vortex_ensure!(
             children.iter().all(|c| c.len() == len),
@@ -49,3 +45,14 @@ impl ScalarFnArray {
         })
     }
 }
+
+pub(super) struct ChildArraysReturnDTypeCtx<'a>(pub(super) &'a [ArrayRef]);
+impl ReturnDTypeCtx for ChildArraysReturnDTypeCtx<'_> {
+    fn child_count(&self) -> usize {
+        self.0.len()
+    }
+
+    fn return_dtype(&self, child_idx: usize) -> VortexResult<DType> {
+        Ok(self.0[child_idx].dtype().clone())
+    }
+}

vortex-array/src/arrays/scalar_fn/metadata.rs

@@ -3,10 +3,10 @@
 
 use vortex_dtype::DType;
 
-use crate::expr::BoundExpression;
+use crate::expr::ScalarFn;
 
 #[derive(Clone, Debug)]
 pub struct ScalarFnMetadata {
-    pub(super) bound: BoundExpression,
+    pub(super) bound: ScalarFn,
     pub(super) child_dtypes: Vec<DType>,
 }

vortex-array/src/arrays/scalar_fn/vtable/mod.rs

@@ -25,13 +25,14 @@ use vortex_vector::Vector;
 use crate::Array;
 use crate::ArrayRef;
 use crate::IntoArray;
+use crate::arrays::scalar_fn::array::ChildArraysReturnDTypeCtx;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::metadata::ScalarFnMetadata;
 use crate::execution::ExecutionCtx;
 use crate::expr;
-use crate::expr::BoundExpression;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprVTable;
+use crate::expr::ScalarFn;
 use crate::optimizer::rules::MatchKey;
 use crate::optimizer::rules::Matcher;
 use crate::serde::ArrayChildren;
@@ -115,9 +116,9 @@ impl VTable for ScalarFnVTable {
 
         #[cfg(debug_assertions)]
         {
-            let child_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
+            let ctx = ChildArraysReturnDTypeCtx(&children);
             vortex_error::vortex_ensure!(
-                &metadata.bound.return_dtype(&child_dtypes)? == dtype,
+                &metadata.bound.return_dtype(&ctx)? == dtype,
                 "Return dtype mismatch when building ScalarFnArray"
             );
         }
@@ -162,16 +163,15 @@ pub trait ScalarFnArrayExt: expr::VTable {
         options: Self::Options,
         children: impl Into<Vec<ArrayRef>>,
     ) -> VortexResult<ArrayRef> {
-        let bound = BoundExpression::new_static(self, options);
+        let bound = ScalarFn::new_static(self, options);
 
         let children = children.into();
         vortex_ensure!(
             children.iter().all(|c| c.len() == len),
             "All child arrays must have the same length as the scalar function array"
         );
 
-        let child_dtypes = children.iter().map(|c| c.dtype().clone()).collect_vec();
-        let dtype = bound.return_dtype(&child_dtypes)?;
+        let dtype = bound.return_dtype(&ChildArraysReturnDTypeCtx(&children))?;
 
         let array_vtable: ArrayVTable = ScalarFnVTable {
             vtable: bound.vtable().clone(),

vortex-array/src/expr/display.rs

@@ -5,8 +5,8 @@ use std::fmt::Display;
 use std::fmt::Formatter;
 use std::ops::Deref;
 
-use crate::expr::BoundExpression;
 use crate::expr::Expression;
+use crate::expr::ScalarFn;
 
 pub enum DisplayFormat {
     Compact,
@@ -19,7 +19,7 @@ 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 bound: &BoundExpression = expr.deref();
+            let bound: &ScalarFn = expr.deref();
             let node_name = format!("{}", bound);
 
             // Get child names for display purposes

vortex-array/src/expr/expression.rs

@@ -9,18 +9,18 @@ use std::hash::Hash;
 use std::ops::Deref;
 use std::sync::Arc;
 
-use itertools::Itertools;
 use vortex_dtype::DType;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 
 use crate::ArrayRef;
+use crate::expr::ReturnDTypeCtx;
 use crate::expr::Root;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
-use crate::expr::bound::BoundExpression;
 use crate::expr::display::DisplayTreeExpr;
+use crate::expr::scalar_fn::ScalarFn;
 use crate::expr::stats::Stat;
 
 /// A node in a Vortex expression tree.
@@ -29,36 +29,39 @@ use crate::expr::stats::Stat;
 /// expression consists of an encoding (vtable), heap-allocated metadata, and child expressions.
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct Expression {
-    /// The bound expression for this node.
-    bound: BoundExpression,
+    /// The scalar function for this node.
+    scalar_fn: ScalarFn,
     /// Any children of this expression.
     children: Arc<[Expression]>,
 }
 
 impl Deref for Expression {
-    type Target = BoundExpression;
+    type Target = ScalarFn;
 
     fn deref(&self) -> &Self::Target {
-        &self.bound
+        &self.scalar_fn
     }
 }
 
 impl Expression {
-    /// Create a new expression node from a bound expression and its children.
+    /// Create a new expression node from a scalar function and children.
     pub fn try_new(
-        bound: BoundExpression,
+        scalar_fn: ScalarFn,
         children: impl Into<Arc<[Expression]>>,
     ) -> VortexResult<Self> {
         let children: Arc<[Expression]> = children.into();
 
         vortex_ensure!(
-            bound.signature().arity().matches(children.len()),
+            scalar_fn.signature().arity().matches(children.len()),
             "Expression arity mismatch: expected {} children but got {}",
-            bound.signature().arity(),
+            scalar_fn.signature().arity(),
             children.len()
         );
 
-        Ok(Self { bound, children })
+        Ok(Self {
+            scalar_fn,
+            children,
+        })
     }
 
     /// Returns true if this expression is of the given vtable type.
@@ -102,24 +105,18 @@ 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()
-            .map(|c| c.return_dtype(scope))
-            .try_collect()?;
-        self.bound.return_dtype(&dtypes)
+        self.scalar_fn.return_dtype(&ExpressionReturnDTypeCtx {
+            expr: self,
+            scope_dtype: scope,
+        })
     }
 
     /// Evaluates the expression in the given scope, returning an array.
     pub fn evaluate(&self, scope: &ArrayRef) -> VortexResult<ArrayRef> {
         if self.is::<Root>() {
             return Ok(scope.clone());
         }
-        self.bound.evaluate(self, scope)
+        self.scalar_fn.evaluate(self, scope)
     }
 
     /// An expression over zone-statistics which implies all records in the zone evaluate to false.
@@ -237,3 +234,29 @@ impl Display for Expression {
         self.fmt_sql(f)
     }
 }
+
+pub(super) struct ExpressionReturnDTypeCtx<'a> {
+    pub(super) expr: &'a Expression,
+    pub(super) scope_dtype: &'a DType,
+}
+
+impl ReturnDTypeCtx for ExpressionReturnDTypeCtx<'_> {
+    fn child_count(&self) -> usize {
+        self.expr.children().len()
+    }
+
+    fn return_dtype(&self, child_idx: usize) -> VortexResult<DType> {
+        let child = &self.expr.children()[child_idx];
+
+        if child.is::<Root>() {
+            return Ok(self.scope_dtype.clone());
+        }
+
+        let ctx = ExpressionReturnDTypeCtx {
+            expr: child,
+            scope_dtype: self.scope_dtype,
+        };
+        let child_fn: &ScalarFn = child.deref();
+        child_fn.return_dtype(&ctx)
+    }
+}

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

@@ -19,6 +19,7 @@ use crate::expr::Arity;
 use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
+use crate::expr::ReturnDTypeCtx;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
@@ -112,19 +113,23 @@ impl VTable for Between {
         )
     }
 
-    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
-        let arr_dt = &arg_dtypes[0];
-        let lower_dt = &arg_dtypes[1];
-        let upper_dt = &arg_dtypes[2];
-
-        if !arr_dt.eq_ignore_nullability(lower_dt) {
+    fn return_dtype(
+        &self,
+        _options: &Self::Options,
+        ctx: &dyn ReturnDTypeCtx,
+    ) -> VortexResult<DType> {
+        let arr_dt = ctx.return_dtype(0)?;
+        let lower_dt = ctx.return_dtype(1)?;
+        let upper_dt = ctx.return_dtype(2)?;
+
+        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

@@ -24,6 +24,7 @@ use crate::expr::Arity;
 use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
+use crate::expr::ReturnDTypeCtx;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
@@ -80,12 +81,12 @@ impl VTable for Binary {
         write!(f, ")")
     }
 
-    fn return_dtype(&self, operator: &Operator, arg_dtypes: &[DType]) -> VortexResult<DType> {
-        let lhs = &arg_dtypes[0];
-        let rhs = &arg_dtypes[1];
+    fn return_dtype(&self, operator: &Operator, ctx: &dyn ReturnDTypeCtx) -> VortexResult<DType> {
+        let lhs = ctx.return_dtype(0)?;
+        let rhs = ctx.return_dtype(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!(

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

@@ -18,6 +18,7 @@ use crate::expr::Arity;
 use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
+use crate::expr::ReturnDTypeCtx;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
@@ -69,7 +70,7 @@ impl VTable for Cast {
         write!(f, ")")
     }
 
-    fn return_dtype(&self, dtype: &DType, _arg_dtypes: &[DType]) -> VortexResult<DType> {
+    fn return_dtype(&self, dtype: &DType, _ctx: &dyn ReturnDTypeCtx) -> VortexResult<DType> {
         Ok(dtype.clone())
     }
 

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

@@ -28,6 +28,7 @@ use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::Expression;
+use crate::expr::ReturnDTypeCtx;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
@@ -76,10 +77,10 @@ impl VTable for DynamicComparison {
     fn return_dtype(
         &self,
         dynamic: &DynamicComparisonExpr,
-        arg_dtypes: &[DType],
+        ctx: &dyn ReturnDTypeCtx,
     ) -> VortexResult<DType> {
-        let lhs = &arg_dtypes[0];
-        if !dynamic.rhs.dtype.eq_ignore_nullability(lhs) {
+        let lhs = ctx.return_dtype(0)?;
+        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/get_item.rs

@@ -26,6 +26,7 @@ use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::Expression;
 use crate::expr::Pack;
+use crate::expr::ReturnDTypeCtx;
 use crate::expr::SimplifyCtx;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
@@ -77,8 +78,12 @@ impl VTable for GetItem {
         write!(f, ".{}", field_name)
     }
 
-    fn return_dtype(&self, field_name: &FieldName, arg_dtypes: &[DType]) -> VortexResult<DType> {
-        let struct_dtype = &arg_dtypes[0];
+    fn return_dtype(
+        &self,
+        field_name: &FieldName,
+        ctx: &dyn ReturnDTypeCtx,
+    ) -> VortexResult<DType> {
+        let struct_dtype = ctx.return_dtype(0)?;
         let field_dtype = struct_dtype
             .as_struct_fields_opt()
             .and_then(|st| st.field(field_name))