Scalar functions

vortex-array/src/arrays/mod.rs

@@ -29,13 +29,13 @@ mod listview;
 mod masked;
 mod null;
 mod primitive;
+mod scalar_fn;
 mod struct_;
 mod varbin;
 mod varbinview;
 
 #[cfg(feature = "arbitrary")]
 pub mod arbitrary;
-
 // TODO(connor): Export exact types, not glob.
 
 pub use bool::*;

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

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::DType;
+
+use crate::ArrayRef;
+use crate::expr::functions::scalar::ScalarFn;
+use crate::stats::ArrayStats;
+use crate::vtable::ArrayVTable;
+
+#[derive(Clone, Debug)]
+pub struct ScalarFnArray {
+    // NOTE(ngates): we should fix vtables so we don't have to hold this
+    pub(super) vtable: ArrayVTable,
+    pub(super) scalar_fn: ScalarFn,
+    pub(super) dtype: DType,
+    pub(super) len: usize,
+    pub(super) children: Vec<ArrayRef>,
+    pub(super) stats: ArrayStats,
+}

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

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::DType;
+
+use crate::expr::functions::scalar::ScalarFn;
+
+#[derive(Clone, Debug)]
+pub struct ScalarFnMetadata {
+    pub(super) scalar_fn: ScalarFn,
+    pub(super) child_dtypes: Vec<DType>,
+}

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

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod array;
+mod metadata;
+mod vtable;

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

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::hash::Hash;
+use std::hash::Hasher;
+
+use vortex_dtype::DType;
+
+use crate::ArrayEq;
+use crate::ArrayHash;
+use crate::Precision;
+use crate::arrays::scalar_fn::array::ScalarFnArray;
+use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
+use crate::stats::StatsSetRef;
+use crate::vtable::BaseArrayVTable;
+
+impl BaseArrayVTable<ScalarFnVTable> for ScalarFnVTable {
+    fn len(array: &ScalarFnArray) -> usize {
+        array.len
+    }
+
+    fn dtype(array: &ScalarFnArray) -> &DType {
+        &array.dtype
+    }
+
+    fn stats(array: &ScalarFnArray) -> StatsSetRef<'_> {
+        array.stats.to_ref(array.as_ref())
+    }
+
+    fn array_hash<H: Hasher>(array: &ScalarFnArray, state: &mut H, precision: Precision) {
+        array.len.hash(state);
+        array.dtype.hash(state);
+        array.scalar_fn.hash(state);
+        for child in &array.children {
+            child.array_hash(state, precision);
+        }
+    }
+
+    fn array_eq(array: &ScalarFnArray, other: &ScalarFnArray, precision: Precision) -> bool {
+        if array.len != other.len {
+            return false;
+        }
+        if array.dtype != other.dtype {
+            return false;
+        }
+        if array.scalar_fn != other.scalar_fn {
+            return false;
+        }
+        for (child, other_child) in array.children.iter().zip(other.children.iter()) {
+            if !child.array_eq(other_child, precision) {
+                return false;
+            }
+        }
+        true
+    }
+}

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

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use itertools::Itertools;
+use vortex_error::VortexExpect;
+use vortex_vector::Datum;
+
+use crate::Array;
+use crate::Canonical;
+use crate::arrays::scalar_fn::array::ScalarFnArray;
+use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
+use crate::expr::functions::ExecutionCtx;
+use crate::vectors::VectorIntoArray;
+use crate::vtable::CanonicalVTable;
+
+impl CanonicalVTable<ScalarFnVTable> for ScalarFnVTable {
+    fn canonicalize(array: &ScalarFnArray) -> Canonical {
+        let child_dtypes: Vec<_> = array.children.iter().map(|c| c.dtype().clone()).collect();
+        let child_datums: Vec<_> = array
+            .children()
+            .iter()
+            // TODO(ngates): we could make all execution operate over datums
+            .map(|child| child.execute().map(Datum::Vector))
+            .try_collect()
+            // FIXME(ngates): canonicalizing really ought to be fallible
+            .vortex_expect(
+                "Failed to execute child array during canonicalization of ScalarFnArray",
+            );
+
+        let ctx = ExecutionCtx::new(array.len, array.dtype.clone(), child_dtypes, child_datums);
+
+        let result_vector = array
+            .scalar_fn
+            .execute(&ctx)
+            .vortex_expect("Canonicalize should be fallible")
+            .into_vector()
+            .vortex_expect("Canonicalize should return a vector");
+
+        result_vector.into_array(&array.dtype).to_canonical()
+    }
+}

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

@@ -0,0 +1,129 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod array;
+mod canonical;
+mod operations;
+mod validity;
+mod visitor;
+
+use itertools::Itertools;
+use vortex_buffer::BufferHandle;
+use vortex_dtype::DType;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_vector::Vector;
+
+use crate::Array;
+use crate::arrays::scalar_fn::array::ScalarFnArray;
+use crate::arrays::scalar_fn::metadata::ScalarFnMetadata;
+use crate::execution::ExecutionCtx;
+use crate::expr::functions;
+use crate::serde::ArrayChildren;
+use crate::vtable;
+use crate::vtable::ArrayId;
+use crate::vtable::ArrayVTable;
+use crate::vtable::ArrayVTableExt;
+use crate::vtable::NotSupported;
+use crate::vtable::VTable;
+
+vtable!(ScalarFn);
+
+#[derive(Clone, Debug)]
+pub struct ScalarFnVTable {
+    vtable: functions::ScalarFnVTable,
+}
+
+impl VTable for ScalarFnVTable {
+    type Array = ScalarFnArray;
+    type Metadata = ScalarFnMetadata;
+    type ArrayVTable = Self;
+    type CanonicalVTable = Self;
+    type OperationsVTable = NotSupported;
+    type ValidityVTable = Self;
+    type VisitorVTable = Self;
+    type ComputeVTable = NotSupported;
+    type EncodeVTable = NotSupported;
+    type OperatorVTable = NotSupported;
+
+    fn id(&self) -> ArrayId {
+        self.vtable.id()
+    }
+
+    fn encoding(array: &Self::Array) -> ArrayVTable {
+        array.vtable.clone()
+    }
+
+    fn metadata(array: &Self::Array) -> VortexResult<Self::Metadata> {
+        let child_dtypes = array.children().iter().map(|c| c.dtype().clone()).collect();
+        Ok(ScalarFnMetadata {
+            scalar_fn: array.scalar_fn.clone(),
+            child_dtypes,
+        })
+    }
+
+    fn serialize(_metadata: Self::Metadata) -> VortexResult<Option<Vec<u8>>> {
+        // Not supported
+        Ok(None)
+    }
+
+    fn deserialize(_bytes: &[u8]) -> VortexResult<Self::Metadata> {
+        vortex_bail!("Deserialization of ScalarFnVTable metadata is not supported");
+    }
+
+    fn build(
+        &self,
+        dtype: &DType,
+        len: usize,
+        metadata: &ScalarFnMetadata,
+        _buffers: &[BufferHandle],
+        children: &dyn ArrayChildren,
+    ) -> VortexResult<Self::Array> {
+        let children: Vec<_> = metadata
+            .child_dtypes
+            .iter()
+            .enumerate()
+            .map(|(idx, child_dtype)| children.get(idx, child_dtype, len))
+            .try_collect()?;
+
+        #[cfg(debug_assertions)]
+        {
+            let child_dtypes: Vec<_> = children.iter().map(|c| c.dtype().clone()).collect();
+            vortex_error::vortex_ensure!(
+                &metadata.scalar_fn.return_dtype(&child_dtypes)? == dtype,
+                "Return dtype mismatch when building ScalarFnArray"
+            );
+        }
+
+        Ok(ScalarFnArray {
+            // This requires a new Arc, but we plan to remove this later anyway.
+            vtable: self.to_vtable(),
+            scalar_fn: metadata.scalar_fn.clone(),
+            dtype: dtype.clone(),
+            len,
+            children,
+            stats: Default::default(),
+        })
+    }
+
+    fn execute(array: &Self::Array, _ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {
+        let input_dtypes: Vec<_> = array.children().iter().map(|c| c.dtype().clone()).collect();
+        let input_datums = array
+            .children()
+            .iter()
+            .map(|child| child.execute())
+            .try_collect()?;
+        let ctx = functions::ExecutionCtx::new(
+            array.len(),
+            array.dtype.clone(),
+            input_dtypes,
+            input_datums,
+        );
+        Ok(array
+            .scalar_fn
+            .execute(&ctx)?
+            .into_vector()
+            .vortex_expect("Vector inputs should return vector outputs"))
+    }
+}

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

@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::ops::Range;
+
+use vortex_error::VortexExpect;
+use vortex_scalar::Scalar;
+use vortex_vector::Datum;
+
+use crate::ArrayRef;
+use crate::IntoArray;
+use crate::arrays::scalar_fn::array::ScalarFnArray;
+use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
+use crate::expr::functions::ExecutionCtx;
+use crate::vtable::OperationsVTable;
+
+impl OperationsVTable<ScalarFnVTable> for ScalarFnVTable {
+    fn slice(array: &ScalarFnArray, range: Range<usize>) -> ArrayRef {
+        let children: Vec<_> = array
+            .children()
+            .iter()
+            .map(|c| c.slice(range.clone()))
+            .collect();
+
+        ScalarFnArray {
+            vtable: array.vtable.clone(),
+            scalar_fn: array.scalar_fn.clone(),
+            dtype: array.dtype.clone(),
+            len: range.len(),
+            children,
+            stats: Default::default(),
+        }
+        .into_array()
+    }
+
+    fn scalar_at(array: &ScalarFnArray, index: usize) -> Scalar {
+        // TODO(ngates): we should evaluate the scalar function over the scalar inputs.
+        let input_datums: Vec<_> = array
+            .children()
+            .iter()
+            .map(|c| c.scalar_at(index))
+            .map(|scalar| Datum::from(scalar.to_vector_scalar()))
+            .collect();
+
+        let ctx = ExecutionCtx::new(
+            1,
+            array.dtype.clone(),
+            array.children().iter().map(|s| s.dtype().clone()).collect(),
+            input_datums,
+        );
+
+        let _result = array
+            .scalar_fn
+            .execute(&ctx)
+            .vortex_expect("Scalar function execution should be fallible")
+            .into_scalar()
+            .vortex_expect("Scalar function execution should return scalar");
+
+        // Convert the vector scalar back into a legacy Scalar for now.
+        todo!("Implement legacy scalar conversion")
+    }
+}

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

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexExpect;
+use vortex_mask::Mask;
+
+use crate::Array;
+use crate::arrays::scalar_fn::array::ScalarFnArray;
+use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
+use crate::expr::functions::NullHandling;
+use crate::vtable::ValidityVTable;
+
+impl ValidityVTable<ScalarFnVTable> for ScalarFnVTable {
+    fn is_valid(array: &ScalarFnArray, index: usize) -> bool {
+        array.scalar_at(index).is_valid()
+    }
+
+    fn all_valid(array: &ScalarFnArray) -> bool {
+        match array.scalar_fn.signature().null_handling() {
+            NullHandling::Propagate | NullHandling::AbsorbsNull => {
+                // Requires all children to guarantee all_valid
+                array.children().iter().all(|child| child.all_valid())
+            }
+            NullHandling::Custom => {
+                // We cannot guarantee that the array is all valid without evaluating the function
+                false
+            }
+        }
+    }
+
+    fn all_invalid(array: &ScalarFnArray) -> bool {
+        match array.scalar_fn.signature().null_handling() {
+            NullHandling::Propagate => {
+                // All null if any child is all null
+                array.children().iter().any(|child| child.all_invalid())
+            }
+            NullHandling::AbsorbsNull | NullHandling::Custom => {
+                // We cannot guarantee that the array is all valid without evaluating the function
+                false
+            }
+        }
+    }
+
+    fn validity_mask(array: &ScalarFnArray) -> Mask {
+        let vector = array
+            .execute()
+            .vortex_expect("Validity mask computation should be fallible");
+        Mask::from_buffer(vector.into_bool().into_bits())
+    }
+}