Expression execution to be passed input vectors (#5507)

This execution model trades off short-circuiting in favor of CSE.

It also allows expressions to be executed independently of the
expression tree, which may be useful depending on how we ultimately end
up wrapping up scalar expressions into arrays.

@joseph-isaacs the alternative would be to pass a `trait ExecutionArgs`
and then calling `child(usize)` could lazily compute the child, bringing
back the flexibility of short-circuiting depending on how the expression
is ultimately executed.

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

Author: gatesn

vortex-array/src/expr/expression.rs

@@ -7,13 +7,16 @@ use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::sync::Arc;
 
+use itertools::Itertools;
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult};
-use vortex_vector::Vector;
+use vortex_vector::{Vector, VectorOps};
 
 use crate::ArrayRef;
 use crate::expr::display::DisplayTreeExpr;
-use crate::expr::{ChildName, ExprId, ExprVTable, ExpressionView, StatsCatalog, VTable};
+use crate::expr::{
+    ChildName, ExecutionArgs, ExprId, ExprVTable, ExpressionView, Root, StatsCatalog, VTable,
+};
 use crate::stats::Stat;
 
 /// A node in a Vortex expression tree.
@@ -144,7 +147,31 @@ impl Expression {
 
     /// Executes the expression over the given vector input scope.
     pub fn execute(&self, vector: &Vector, dtype: &DType) -> VortexResult<Vector> {
-        self.vtable.as_dyn().execute(self, vector, dtype)
+        // We special-case the "root" expression that must extract that scope vector directly.
+        if self.is::<Root>() {
+            return Ok(vector.clone());
+        }
+
+        let return_dtype = self.return_dtype(dtype)?;
+        let child_dtypes: Vec<_> = self
+            .children
+            .iter()
+            .map(|child| child.return_dtype(dtype))
+            .try_collect()?;
+        let child_vectors: Vec<_> = self
+            .children
+            .iter()
+            .map(|child| child.execute(vector, dtype))
+            .try_collect()?;
+
+        let args = ExecutionArgs {
+            vectors: child_vectors,
+            dtypes: child_dtypes,
+            row_count: vector.len(),
+            return_dtype,
+        };
+
+        self.vtable.as_dyn().execute(&self.data, args)
     }
 
     /// An expression over zone-statistics which implies all records in the zone evaluate to false.

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

@@ -13,7 +13,9 @@ use vortex_vector::Vector;
 use crate::ArrayRef;
 use crate::compute::cast as compute_cast;
 use crate::expr::expression::Expression;
-use crate::expr::{ChildName, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::expr::{
+    ChildName, ExecutionArgs, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt,
+};
 use crate::stats::Stat;
 
 /// A cast expression that converts values to a target data type.
@@ -88,16 +90,6 @@ impl VTable for Cast {
         })
     }
 
-    fn execute(
-        &self,
-        expr: &ExpressionView<Self>,
-        vector: &Vector,
-        dtype: &DType,
-    ) -> VortexResult<Vector> {
-        let input = expr.child(0).execute(vector, dtype)?;
-        vortex_compute::cast::Cast::cast(&input, dtype)
-    }
-
     fn stat_expression(
         &self,
         expr: &ExpressionView<Self>,
@@ -129,6 +121,14 @@ impl VTable for Cast {
             }
         }
     }
+
+    fn execute(&self, target_dtype: &DType, mut args: ExecutionArgs) -> VortexResult<Vector> {
+        let input = args
+            .vectors
+            .pop()
+            .vortex_expect("missing input for Cast expression");
+        vortex_compute::cast::Cast::cast(&input, target_dtype)
+    }
 }
 
 /// Creates an expression that casts values to a target data type.

vortex-array/src/expr/exprs/get_item/mod.rs

@@ -9,13 +9,15 @@ use std::ops::Not;
 use prost::Message;
 use vortex_compute::mask::MaskValidity;
 use vortex_dtype::{DType, FieldName, FieldPath, Nullability};
-use vortex_error::{VortexResult, vortex_bail, vortex_err};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_proto::expr as pb;
 use vortex_vector::{Vector, VectorOps};
 
 use crate::compute::mask;
 use crate::expr::exprs::root::root;
-use crate::expr::{ChildName, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::expr::{
+    ChildName, ExecutionArgs, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt,
+};
 use crate::stats::Stat;
 use crate::{ArrayRef, ToCanonical};
 
@@ -98,29 +100,6 @@ impl VTable for GetItem {
         }
     }
 
-    fn execute(
-        &self,
-        expr: &ExpressionView<Self>,
-        vector: &Vector,
-        dtype: &DType,
-    ) -> VortexResult<Vector> {
-        let child_dtype = expr.child(0).return_dtype(dtype)?;
-        let struct_dtype = child_dtype
-            .as_struct_fields_opt()
-            .ok_or_else(|| vortex_err!("Expected struct dtype for child of GetItem expression"))?;
-        let field_idx = struct_dtype
-            .find(expr.data())
-            .ok_or_else(|| vortex_err!("Field {} not found in struct dtype", expr.data()))?;
-
-        let struct_vector = expr.child(0).execute(vector, dtype)?.into_struct();
-
-        // We must intersect the validity with that of the parent struct
-        let field = struct_vector.fields()[field_idx].clone();
-        let field = MaskValidity::mask_validity(field, struct_vector.validity());
-
-        Ok(field)
-    }
-
     fn stat_expression(
         &self,
         expr: &ExpressionView<Self>,
@@ -137,6 +116,27 @@ impl VTable for GetItem {
         //  falsify expressions, but for now I'm preserving the existing buggy behavior.
         catalog.stats_ref(&FieldPath::from_name(expr.data().clone()), stat)
     }
+
+    fn execute(&self, field_name: &FieldName, mut args: ExecutionArgs) -> VortexResult<Vector> {
+        let struct_dtype = args.dtypes[0]
+            .as_struct_fields_opt()
+            .ok_or_else(|| vortex_err!("Expected struct dtype for child of GetItem expression"))?;
+        let field_idx = struct_dtype
+            .find(field_name)
+            .ok_or_else(|| vortex_err!("Field {} not found in struct dtype", field_name))?;
+
+        let struct_vector = args
+            .vectors
+            .pop()
+            .vortex_expect("missing input")
+            .into_struct();
+
+        // We must intersect the validity with that of the parent struct
+        let field = struct_vector.fields()[field_idx].clone();
+        let field = MaskValidity::mask_validity(field, struct_vector.validity());
+
+        Ok(field)
+    }
 }
 
 /// Creates an expression that accesses a field from the root array.

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

@@ -5,15 +5,17 @@ use std::fmt::Formatter;
 use std::ops::Not;
 
 use vortex_dtype::{DType, Nullability};
-use vortex_error::{VortexResult, vortex_bail};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_mask::Mask;
 use vortex_vector::bool::BoolVector;
 use vortex_vector::{Vector, VectorOps};
 
 use crate::arrays::{BoolArray, ConstantArray};
 use crate::expr::exprs::binary::eq;
 use crate::expr::exprs::literal::lit;
-use crate::expr::{ChildName, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::expr::{
+    ChildName, ExecutionArgs, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt,
+};
 use crate::stats::Stat;
 use crate::{Array, ArrayRef, IntoArray};
 
@@ -71,20 +73,6 @@ impl VTable for IsNull {
         }
     }
 
-    fn execute(
-        &self,
-        expr: &ExpressionView<Self>,
-        vector: &Vector,
-        dtype: &DType,
-    ) -> VortexResult<Vector> {
-        let child = expr.child(0).execute(vector, dtype)?;
-        Ok(BoolVector::new(
-            child.validity().to_bit_buffer().not(),
-            Mask::new_true(child.len()),
-        )
-        .into())
-    }
-
     fn stat_falsification(
         &self,
         expr: &ExpressionView<Self>,
@@ -93,6 +81,15 @@ impl VTable for IsNull {
         let null_count_expr = expr.child(0).stat_expression(Stat::NullCount, catalog)?;
         Some(eq(null_count_expr, lit(0u64)))
     }
+
+    fn execute(&self, _data: &Self::Instance, mut args: ExecutionArgs) -> VortexResult<Vector> {
+        let child = args.vectors.pop().vortex_expect("Missing input child");
+        Ok(BoolVector::new(
+            child.validity().to_bit_buffer().not(),
+            Mask::new_true(child.len()),
+        )
+        .into())
+    }
 }
 
 /// Creates an expression that checks for null values.

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

@@ -5,12 +5,14 @@ use std::fmt::Formatter;
 
 use vortex_compute::logical::LogicalNot;
 use vortex_dtype::DType;
-use vortex_error::{VortexResult, vortex_bail};
+use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_vector::Vector;
 
 use crate::ArrayRef;
 use crate::compute::invert;
-use crate::expr::{ChildName, ExprId, Expression, ExpressionView, VTable, VTableExt};
+use crate::expr::{
+    ChildName, ExecutionArgs, ExprId, Expression, ExpressionView, VTable, VTableExt,
+};
 
 /// Expression that logically inverts boolean values.
 pub struct Not;
@@ -69,13 +71,8 @@ impl VTable for Not {
         invert(&child_result)
     }
 
-    fn execute(
-        &self,
-        expr: &ExpressionView<Self>,
-        vector: &Vector,
-        dtype: &DType,
-    ) -> VortexResult<Vector> {
-        let child = expr.child(0).execute(vector, dtype)?;
+    fn execute(&self, _data: &Self::Instance, mut args: ExecutionArgs) -> VortexResult<Vector> {
+        let child = args.vectors.pop().vortex_expect("Missing input child");
         Ok(child.into_bool().not().into())
     }
 }

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

@@ -9,7 +9,9 @@ use vortex_vector::Vector;
 
 use crate::ArrayRef;
 use crate::expr::expression::Expression;
-use crate::expr::{ChildName, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt};
+use crate::expr::{
+    ChildName, ExecutionArgs, ExprId, ExpressionView, StatsCatalog, VTable, VTableExt,
+};
 use crate::stats::Stat;
 
 /// An expression that returns the full scope of the expression evaluation.
@@ -60,13 +62,8 @@ impl VTable for Root {
         Ok(scope.clone())
     }
 
-    fn execute(
-        &self,
-        _expr: &ExpressionView<Self>,
-        vector: &Vector,
-        _dtype: &DType,
-    ) -> VortexResult<Vector> {
-        Ok(vector.clone())
+    fn execute(&self, _data: &Self::Instance, _args: ExecutionArgs) -> VortexResult<Vector> {
+        vortex_bail!("Root expression is not executable")
     }
 
     fn stat_expression(

vortex-array/src/expr/exprs/select/mod.rs

@@ -4,6 +4,7 @@
 pub mod transform;
 
 use std::fmt::{Display, Formatter};
+use std::sync::Arc;
 
 use itertools::Itertools;
 use prost::Message;
@@ -12,10 +13,11 @@ use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_proto::expr::select_opts::Opts;
 use vortex_proto::expr::{FieldNames as ProtoFieldNames, SelectOpts};
 use vortex_vector::Vector;
+use vortex_vector::struct_::StructVector;
 
 use crate::expr::expression::Expression;
 use crate::expr::field::DisplayFieldNames;
-use crate::expr::{ChildName, ExprId, ExpressionView, VTable, VTableExt};
+use crate::expr::{ChildName, ExecutionArgs, ExprId, ExpressionView, VTable, VTableExt};
 use crate::{ArrayRef, IntoArray, ToCanonical};
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -146,15 +148,45 @@ impl VTable for Select {
         .into_array())
     }
 
-    fn execute(
-        &self,
-        _expr: &ExpressionView<Self>,
-        _vector: &Vector,
-        _dtype: &DType,
-    ) -> VortexResult<Vector> {
-        vortex_bail!(
-            "Select expressions cannot be executed. They must be removed during optimization."
-        )
+    fn execute(&self, selection: &FieldSelection, mut args: ExecutionArgs) -> VortexResult<Vector> {
+        let child = args
+            .vectors
+            .pop()
+            .vortex_expect("Missing input child")
+            .into_struct();
+        let child_fields = args
+            .dtypes
+            .pop()
+            .vortex_expect("Missing input dtype")
+            .into_struct_fields();
+
+        let field_indices: Vec<usize> = match selection {
+            FieldSelection::Include(f) => f
+                .iter()
+                .map(|name| {
+                    child_fields
+                        .find(name)
+                        .ok_or_else(|| vortex_err!("Field {} not found in struct dtype", name))
+                })
+                .try_collect(),
+            FieldSelection::Exclude(names) => child_fields
+                .names()
+                .iter()
+                .filter(|&f| !names.as_ref().contains(f))
+                .map(|name| {
+                    child_fields
+                        .find(name)
+                        .ok_or_else(|| vortex_err!("Field {} not found in struct dtype", name))
+                })
+                .try_collect(),
+        }?;
+
+        let (fields, mask) = child.into_parts();
+        let new_fields = field_indices
+            .iter()
+            .map(|&idx| fields[idx].clone())
+            .collect();
+        Ok(unsafe { StructVector::new_unchecked(Arc::new(new_fields), mask) }.into())
     }
 }