Test implementation of cast using ScalarFn (#5586)

* Defines a CastFn implementation of ScalarFn
* Updates the optimizer rules to support matcher instances (vs marker
structs) so we can match on runtime array IDs
* Implement CastReduce rule that for now just removes the cast if the
child array already has the same target DType.

---------

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

Author: gatesn

Cargo.lock

@@ -15,6 +15,7 @@ use crate::execution::BatchKernelRef;
 use crate::execution::BindCtx;
 use crate::execution::kernel;
 use crate::optimizer::rules::ArrayParentReduceRule;
+use crate::optimizer::rules::Exact;
 use crate::vtable::OperatorVTable;
 use crate::vtable::ValidityHelper;
 
@@ -47,7 +48,15 @@ impl OperatorVTable<BoolVTable> for BoolVTable {
 #[derive(Default, Debug)]
 pub struct BoolMaskedValidityRule;
 
-impl ArrayParentReduceRule<BoolVTable, MaskedVTable> for BoolMaskedValidityRule {
+impl ArrayParentReduceRule<Exact<BoolVTable>, Exact<MaskedVTable>> for BoolMaskedValidityRule {
+    fn child(&self) -> Exact<BoolVTable> {
+        Exact::from(&BoolVTable)
+    }
+
+    fn parent(&self) -> Exact<MaskedVTable> {
+        Exact::from(&MaskedVTable)
+    }
+
     fn reduce_parent(
         &self,
         array: &BoolArray,

vortex-array/src/arrays/bool/vtable/operator.rs

@@ -17,6 +17,7 @@ use crate::execution::BatchKernelRef;
 use crate::execution::BindCtx;
 use crate::execution::kernel;
 use crate::optimizer::rules::ArrayParentReduceRule;
+use crate::optimizer::rules::Exact;
 use crate::vtable::OperatorVTable;
 use crate::vtable::ValidityHelper;
 
@@ -54,7 +55,17 @@ impl OperatorVTable<DecimalVTable> for DecimalVTable {
 #[derive(Default, Debug)]
 pub struct DecimalMaskedValidityRule;
 
-impl ArrayParentReduceRule<DecimalVTable, MaskedVTable> for DecimalMaskedValidityRule {
+impl ArrayParentReduceRule<Exact<DecimalVTable>, Exact<MaskedVTable>>
+    for DecimalMaskedValidityRule
+{
+    fn child(&self) -> Exact<DecimalVTable> {
+        Exact::from(&DecimalVTable)
+    }
+
+    fn parent(&self) -> Exact<MaskedVTable> {
+        Exact::from(&MaskedVTable)
+    }
+
     fn reduce_parent(
         &self,
         array: &DecimalArray,

vortex-array/src/arrays/decimal/vtable/operator.rs

@@ -50,6 +50,7 @@ pub use listview::*;
 pub use masked::*;
 pub use null::*;
 pub use primitive::*;
+pub use scalar_fn::*;
 pub use struct_::*;
 pub use varbin::*;
 pub use varbinview::*;

vortex-array/src/arrays/mod.rs

@@ -17,6 +17,7 @@ use crate::execution::BatchKernelRef;
 use crate::execution::BindCtx;
 use crate::execution::kernel;
 use crate::optimizer::rules::ArrayParentReduceRule;
+use crate::optimizer::rules::Exact;
 use crate::vtable::OperatorVTable;
 use crate::vtable::ValidityHelper;
 
@@ -52,7 +53,17 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
 #[derive(Default, Debug)]
 pub struct PrimitiveMaskedValidityRule;
 
-impl ArrayParentReduceRule<PrimitiveVTable, MaskedVTable> for PrimitiveMaskedValidityRule {
+impl ArrayParentReduceRule<Exact<PrimitiveVTable>, Exact<MaskedVTable>>
+    for PrimitiveMaskedValidityRule
+{
+    fn child(&self) -> Exact<PrimitiveVTable> {
+        Exact::from(&PrimitiveVTable)
+    }
+
+    fn parent(&self) -> Exact<MaskedVTable> {
+        Exact::from(&MaskedVTable)
+    }
+
     fn reduce_parent(
         &self,
         array: &PrimitiveArray,

vortex-array/src/arrays/primitive/vtable/operator.rs

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

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

@@ -7,19 +7,28 @@ mod operations;
 mod validity;
 mod visitor;
 
+use std::marker::PhantomData;
+use std::ops::Deref;
+
 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_error::vortex_ensure;
 use vortex_vector::Vector;
 
 use crate::Array;
+use crate::ArrayRef;
+use crate::IntoArray;
 use crate::arrays::scalar_fn::array::ScalarFnArray;
 use crate::arrays::scalar_fn::metadata::ScalarFnMetadata;
 use crate::execution::ExecutionCtx;
 use crate::expr::functions;
+use crate::expr::functions::scalar::ScalarFn;
+use crate::optimizer::rules::MatchKey;
+use crate::optimizer::rules::Matcher;
 use crate::serde::ArrayChildren;
 use crate::vtable;
 use crate::vtable::ArrayId;
@@ -127,3 +136,112 @@ impl VTable for ScalarFnVTable {
             .vortex_expect("Vector inputs should return vector outputs"))
     }
 }
+
+/// Array factory functions for scalar functions.
+pub trait ScalarFnArrayExt: functions::VTable {
+    fn try_new_array(
+        &'static self,
+        len: usize,
+        options: Self::Options,
+        children: impl Into<Vec<ArrayRef>>,
+    ) -> VortexResult<ArrayRef> {
+        let scalar_fn = ScalarFn::new_static(self, options);
+
+        let children = children.into();
+        vortex_ensure!(
+            children.iter().all(|c| c.len() == len),
+            "All child arrays must have the same length as the scalar function array"
+        );
+
+        let child_dtypes = children.iter().map(|c| c.dtype().clone()).collect_vec();
+        let dtype = scalar_fn.return_dtype(&child_dtypes)?;
+
+        let array_vtable: ArrayVTable = ScalarFnVTable {
+            vtable: scalar_fn.vtable().clone(),
+        }
+        .into_vtable();
+
+        Ok(ScalarFnArray {
+            vtable: array_vtable,
+            scalar_fn,
+            dtype,
+            len,
+            children,
+            stats: Default::default(),
+        }
+        .into_array())
+    }
+}
+impl<V: functions::VTable> ScalarFnArrayExt for V {}
+
+/// A matcher that matches any scalar function expression.
+#[derive(Debug)]
+pub struct AnyScalarFn;
+impl Matcher for AnyScalarFn {
+    type View<'a> = &'a ScalarFnArray;
+
+    fn key(&self) -> MatchKey {
+        MatchKey::Any
+    }
+
+    fn try_match<'a>(&self, array: &'a ArrayRef) -> Option<Self::View<'a>> {
+        array.as_opt::<ScalarFnVTable>()
+    }
+}
+
+/// A matcher that matches a specific scalar function expression.
+#[derive(Debug)]
+pub struct ExactScalarFn<F: functions::VTable> {
+    id: ArrayId,
+    _phantom: PhantomData<F>,
+}
+
+impl<F: functions::VTable> From<&'static F> for ExactScalarFn<F> {
+    fn from(value: &'static F) -> Self {
+        Self {
+            id: value.id(),
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<F: functions::VTable> Matcher for ExactScalarFn<F> {
+    type View<'a> = ScalarFnArrayView<'a, F>;
+
+    fn key(&self) -> MatchKey {
+        MatchKey::Array(self.id.clone())
+    }
+
+    fn try_match<'a>(&self, array: &'a ArrayRef) -> Option<Self::View<'a>> {
+        let scalar_fn_array = array.as_opt::<ScalarFnVTable>()?;
+        let scalar_fn_vtable = scalar_fn_array
+            .scalar_fn
+            .vtable()
+            .as_any()
+            .downcast_ref::<F>()?;
+        let scalar_fn_options = scalar_fn_array
+            .scalar_fn
+            .options()
+            .as_any()
+            .downcast_ref::<F::Options>()?;
+        Some(ScalarFnArrayView {
+            array,
+            vtable: scalar_fn_vtable,
+            options: scalar_fn_options,
+        })
+    }
+}
+
+pub struct ScalarFnArrayView<'a, F: functions::VTable> {
+    array: &'a ArrayRef,
+    pub vtable: &'a F,
+    pub options: &'a F::Options,
+}
+
+impl<F: functions::VTable> Deref for ScalarFnArrayView<'_, F> {
+    type Target = ArrayRef;
+
+    fn deref(&self) -> &Self::Target {
+        self.array
+    }
+}

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

@@ -11,7 +11,7 @@ impl VisitorVTable<ScalarFnVTable> for ScalarFnVTable {
     fn visit_buffers(_array: &ScalarFnArray, _visitor: &mut dyn ArrayBufferVisitor) {}
 
     fn visit_children(array: &ScalarFnArray, visitor: &mut dyn ArrayChildVisitor) {
-        for (idx, child) in array.children().iter().enumerate() {
+        for (idx, child) in array.children.iter().enumerate() {
             let name = array.scalar_fn.signature().arg_name(idx);
             visitor.visit_child(name.as_ref(), child.as_ref())
         }