Execute fn

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())
     }
 }