Scalar functions (#5561)

Introduces the traits required to define abstract scalar functions and
wrap them up as either expressions or arrays.

This is one of the types of functions that we will split the generic
ComputeFn trait into and allows us to lazily defer scalar compute in the
array tree.

We don't add them in this PR, but there may be some properties that are
useful for ScalarFns in some form:
```

    /// The identity element `e` where `f(e, x) = f(x, e) = x`.
    ///
    /// When an argument is the identity element, the function can be
    /// eliminated entirely, returning the other argument unchanged.
    ///
    /// # Examples
    /// - `AND`: `true` (AND(true, x) → x)
    /// - `OR`: `false` (OR(false, x) → x)
    /// - `+`: `0` (0 + x → x)
    /// - `*`: `1` (1 * x → x)
    /// - `COALESCE`: `NULL` (COALESCE(NULL, x) → x)
    fn identity_element(&self, options: &Self::Options) -> Option<Scalar> {
        _ = options;
        None
    }

    /// The absorbing element `a` where `f(a, x) = f(x, a) = a`.
    ///
    /// When any argument is the absorbing element, the function short-circuits
    /// immediately, returning that element without evaluating other arguments.
    /// Also known as the "annihilator" or "zero element".
    ///
    /// # Examples
    /// - `AND`: `false` (AND(false, x) → false)
    /// - `OR`: `true` (OR(true, x) → true)
    /// - `*`: `0` (0 * x → 0)
    fn absorbing_element(&self, options: &Self::Options) -> Option<Scalar> {
        _ = options;
        None
    }

    /// Whether argument order is irrelevant: `f(a, b) = f(b, a)`.
    ///
    /// Enables expression normalization (e.g., sorting arguments by column id)
    /// for better common subexpression elimination and pattern matching.
    ///
    /// # Examples
    /// - Commutative: `+`, `*`, `AND`, `OR`, `=`, `!=`, `MIN`, `MAX`
    /// - Non-commutative: `-`, `/`, `<`, `>`, `CONCAT`
    fn is_commutative(&self, options: &Self::Options) -> bool {
        _ = options;
        false
    }

    /// Whether `f(x, x) = x`.
    ///
    /// Enables simplification when the same expression appears multiple times
    /// as arguments to the function.
    ///
    /// # Examples
    /// - Idempotent: `AND`, `OR`, `MIN`, `MAX`
    /// - Non-idempotent: `+` (x + x = 2x), `*` (x * x = x²)
    fn is_idempotent(&self, options: &Self::Options) -> bool {
        _ = options;
        false
    }

    /// Whether `f(f(x)) = x` for unary functions.
    ///
    /// Enables cancellation of nested self-applications.
    ///
    /// # Examples
    /// - Involutions: `NOT`, `NEG` (for signed types), `REVERSE`
    /// - Non-involutions: `ABS`, `UPPER`, `LOWER`
    fn is_involution(&self, options: &Self::Options) -> bool {
        _ = options;
        false
    }

    /// How the function behaves when one or more arguments are NULL.
    ///
    /// Most functions propagate NULL (any NULL argument produces NULL output).
    /// Some functions have special NULL handling that can short-circuit
    /// evaluation or treat NULL as a meaningful value.
    ///
    /// Required for correct NULL semantics; may also enable optimizations
    /// when argument nullability is known from schema or statistics.
    fn null_handling(&self, options: &Self::Options) -> NullHandling {
        _ = options;
        NullHandling::default()
    }
    
    
/// How a function handles NULL arguments.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum NullHandling {
    /// NULL in any argument produces NULL output.
    ///
    /// This is standard SQL behavior for most scalar functions.
    /// Enables simplification when any argument is known to be NULL.
    Propagate,

    /// NULL is short-circuited when paired with the absorbing element.
    ///
    /// This is a special case where the absorbing element "wins" over NULL.
    ///
    /// # Examples
    /// - `AND_KLEENE(false, NULL)` → `false` (false absorbs NULL)
    /// - `OR_KLEENE(true, NULL)` → `true` (true absorbs NULL)
    AbsorbsNull,

    /// The function has special NULL semantics that don't follow
    /// simple propagation rules.
    ///
    /// This prevents any simplifications based on NULL arguments.
    ///
    /// # Examples
    /// - `IS NULL`, `IS NOT NULL`: NULL → true/false
    /// - `COALESCE`: returns first non-NULL argument
    /// - `NULLIF`: conditionally produces NULL
    #[default]
    Custom,
}
```

---------

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

Author: gatesn

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