Flat Layout Execution

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

Author: gatesn

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

@@ -0,0 +1,70 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use itertools::Itertools;
+use vortex_dtype::DType;
+use vortex_error::VortexResult;
+use vortex_vector::Datum;
+use vortex_vector::Scalar;
+use vortex_vector::Vector;
+
+use crate::expr::ExecutionArgs;
+use crate::expr::ScalarFn;
+use crate::kernel::Kernel;
+use crate::kernel::KernelRef;
+
+/// A CPU kernel for executing scalar functions.
+#[derive(Debug)]
+pub struct ScalarFnKernel {
+    /// The scalar function to apply.
+    scalar_fn: ScalarFn,
+
+    /// Inputs to the kernel, either constants or other kernels.
+    inputs: Vec<KernelInput>,
+    /// The input data types
+    input_dtypes: Vec<DType>,
+    /// The row count for vector inputs
+    row_count: usize,
+    /// The return data type
+    return_dtype: DType,
+}
+
+#[derive(Debug)]
+enum KernelInput {
+    Scalar(Scalar),
+    Vector(KernelRef),
+}
+
+impl Kernel for ScalarFnKernel {
+    fn execute(self: Box<Self>) -> VortexResult<Vector> {
+        let datums: Vec<_> = self
+            .inputs
+            .into_iter()
+            .map(|input| {
+                Ok(match input {
+                    KernelInput::Scalar(s) => Datum::Scalar(s),
+                    KernelInput::Vector(k) => Datum::Vector(k.execute()?),
+                })
+            })
+            .try_collect()?;
+
+        let args = ExecutionArgs {
+            datums,
+            dtypes: self.input_dtypes,
+            row_count: self.row_count,
+            return_dtype: self.return_dtype,
+        };
+
+        Ok(self.scalar_fn.execute(args)?.ensure_vector(self.row_count))
+    }
+
+    fn children(&self) -> Vec<&KernelRef> {
+        self.inputs
+            .iter()
+            .filter_map(|input| match input {
+                KernelInput::Vector(k) => Some(k),
+                KernelInput::Scalar(_) => None,
+            })
+            .collect()
+    }
+}

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

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 mod array;
+mod kernel;
 mod metadata;
 mod vtable;
 

vortex-array/src/kernel/mod.rs

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Kernels represent the CPU physical plan for array execution.
+
+use std::fmt::Debug;
+use std::fmt::Display;
+
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::Vector;
+
+/// A boxed reference to a kernel.
+pub type KernelRef = Box<dyn Kernel>;
+
+/// A trait representing the physical CPU execution of an array tree.
+pub trait Kernel: 'static + Send + Debug {
+    /// Execute the kernel and produce a vector result.
+    fn execute(self: Box<Self>) -> VortexResult<Vector>;
+
+    /// Return the child kernels of this kernel.
+    fn children(&self) -> Vec<&KernelRef>;
+
+    /// Report an estimated cost for computing this kernel over the given filter mask.
+    ///
+    /// This is obviously a very rough estimate, but is used to decide when a filter should be
+    /// pushed through a kernel using [`Kernel::push_down_filter`].
+    ///
+    /// Return [`f64::INFINITY`] if the kernel has unknown cost, meaning filters will _always_
+    /// be pushed through the kernel if possible.
+    fn cost_estimate(&self, selection: &Mask) -> f64 {
+        _ = selection;
+        f64::INFINITY
+    }
+
+    /// Push a selection mask through this kernel.
+    ///
+    /// Return `Ok(None)` if the filter cannot be pushed down.
+    fn push_down_filter(&self, selection: &Mask) -> VortexResult<Option<KernelRef>> {
+        _ = selection;
+        Ok(None)
+    }
+}

vortex-array/src/lib.rs

@@ -35,6 +35,7 @@ pub mod expr;
 mod expression;
 mod hash;
 pub mod iter;
+mod kernel;
 mod mask;
 mod mask_future;
 mod metadata;