Execute ScalarFn arrays (#5611)

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

Author: gatesn

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

@@ -2,11 +2,16 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_dtype::DType;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 
+use crate::Array;
 use crate::ArrayRef;
+use crate::arrays::ScalarFnVTable;
 use crate::expr::functions::scalar::ScalarFn;
 use crate::stats::ArrayStats;
 use crate::vtable::ArrayVTable;
+use crate::vtable::ArrayVTableExt;
 
 #[derive(Clone, Debug)]
 pub struct ScalarFnArray {
@@ -18,3 +23,25 @@ pub struct ScalarFnArray {
     pub(super) children: Vec<ArrayRef>,
     pub(super) stats: ArrayStats,
 }
+
+impl ScalarFnArray {
+    /// Create a new ScalarFnArray from a scalar function and its children.
+    pub fn try_new(scalar_fn: ScalarFn, children: Vec<ArrayRef>, len: usize) -> VortexResult<Self> {
+        let arg_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
+        let dtype = scalar_fn.return_dtype(&arg_dtypes)?;
+
+        vortex_ensure!(
+            children.iter().all(|c| c.len() == len),
+            "ScalarFnArray must have children equal to the array length"
+        );
+
+        Ok(Self {
+            vtable: ScalarFnVTable::new(scalar_fn.vtable().clone()).into_vtable(),
+            scalar_fn,
+            dtype,
+            len,
+            children,
+            stats: Default::default(),
+        })
+    }
+}

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

@@ -5,4 +5,5 @@ mod array;
 mod metadata;
 mod vtable;
 
+pub use array::*;
 pub use vtable::*;

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

@@ -53,6 +53,12 @@ pub struct ScalarFnVTable {
     vtable: functions::ScalarFnVTable,
 }
 
+impl ScalarFnVTable {
+    pub fn new(vtable: functions::ScalarFnVTable) -> Self {
+        Self { vtable }
+    }
+}
+
 impl VTable for ScalarFnVTable {
     type Array = ScalarFnArray;
     type Metadata = ScalarFnMetadata;

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

@@ -9,12 +9,16 @@ use std::sync::Arc;
 use itertools::Itertools;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_session::SessionVar;
+use vortex_vector::Datum;
+use vortex_vector::ScalarOps;
 use vortex_vector::Vector;
+use vortex_vector::VectorMutOps;
 
 use crate::ArrayRef;
+use crate::IntoArray;
+use crate::arrays::ScalarFnArray;
 use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
@@ -92,12 +96,27 @@ impl VTable for ScalarFnExpr {
         expr.data().return_dtype(&arg_dtypes)
     }
 
-    fn evaluate(&self, _expr: &ExpressionView<Self>, _scope: &ArrayRef) -> VortexResult<ArrayRef> {
-        vortex_bail!("Scalar function evaluation not yet implemented")
+    fn evaluate(&self, expr: &ExpressionView<Self>, scope: &ArrayRef) -> VortexResult<ArrayRef> {
+        let children: Vec<_> = expr
+            .children()
+            .iter()
+            .map(|child| child.evaluate(scope))
+            .try_collect()?;
+        Ok(ScalarFnArray::try_new(expr.data().clone(), children, scope.len())?.into_array())
     }
 
-    fn execute(&self, _data: &Self::Instance, _args: ExecutionArgs) -> VortexResult<Vector> {
-        vortex_bail!("Scalar function execution not yet implemented")
+    fn execute(&self, func: &ScalarFn, args: ExecutionArgs) -> VortexResult<Vector> {
+        let expr_args = functions::ExecutionArgs::new(
+            args.row_count,
+            args.return_dtype,
+            args.dtypes,
+            args.vectors.into_iter().map(Datum::Vector).collect(),
+        );
+        let result = func.execute(&expr_args)?;
+        Ok(match result {
+            Datum::Scalar(s) => s.repeat(args.row_count).freeze(),
+            Datum::Vector(v) => v,
+        })
     }
 
     fn stat_falsification(