CastFn

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

Author: gatesn

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

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

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

@@ -7,6 +7,9 @@ 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;
@@ -24,6 +27,8 @@ 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;
@@ -168,3 +173,75 @@ pub trait ScalarFnArrayExt: functions::VTable {
     }
 }
 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/visitor.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())
         }

vortex-array/src/expr/functions/scalar.rs

@@ -196,9 +196,9 @@ pub struct ScalarFnOptions<'a> {
     pub(crate) options: &'a dyn Any,
 }
 
-impl ScalarFnOptions<'_> {
+impl<'a> ScalarFnOptions<'a> {
     /// Get the options as a `dyn Any`.
-    pub fn as_any(&self) -> &dyn Any {
+    pub fn as_any(&self) -> &'a dyn Any {
         self.options
     }
 }

vortex-array/src/expr/functions/vtable.rs

@@ -167,6 +167,8 @@ pub enum NullHandling {
 
 /// An object-safe vtable for scalar functions that dispatches to the non-object-safe vtable.
 pub(crate) trait DynScalarFnVTable: 'static + Send + Sync {
+    fn as_any(&self) -> &dyn Any;
+
     fn id(&self) -> FunctionId;
 
     fn options_serialize(&self, options: &dyn Any) -> VortexResult<Option<Vec<u8>>>;
@@ -202,6 +204,10 @@ pub(crate) trait DynScalarFnVTable: 'static + Send + Sync {
 #[repr(transparent)]
 pub struct ScalarFnVTableAdapter<V>(V);
 impl<V: VTable> DynScalarFnVTable for ScalarFnVTableAdapter<V> {
+    fn as_any(&self) -> &dyn Any {
+        &self.0
+    }
+
     fn id(&self) -> FunctionId {
         V::id(&self.0)
     }
@@ -308,6 +314,10 @@ impl ScalarFnVTable {
         self.0.id()
     }
 
+    pub fn as_any(&self) -> &dyn Any {
+        self.0.deref().as_any()
+    }
+
     pub fn deserialize(&self, bytes: &[u8]) -> VortexResult<ScalarFn> {
         let options = self.0.options_deserialize(bytes)?;
         // SAFETY: options were created by this vtable.

vortex-array/src/scalar_fns/cast/array.rs

@@ -1,7 +1,59 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-// TODO(ngates): can't yet write the rule matching on scalar function vtable
+use vortex_error::VortexResult;
+
+use crate::Array;
+use crate::ArrayRef;
+use crate::ArrayVisitor;
+use crate::arrays::ExactScalarFn;
+use crate::arrays::ScalarFnArrayView;
+use crate::optimizer::rules::ArrayReduceRule;
+use crate::scalar_fns::cast::CastFn;
+
 #[derive(Debug)]
-#[allow(dead_code)]
-pub struct CastReduce;
+pub(crate) struct CastArrayReduce;
+
+impl ArrayReduceRule<ExactScalarFn<CastFn>> for CastArrayReduce {
+    fn matcher(&self) -> ExactScalarFn<CastFn> {
+        ExactScalarFn::from(&CastFn)
+    }
+
+    fn reduce(&self, array: ScalarFnArrayView<'_, CastFn>) -> VortexResult<Option<ArrayRef>> {
+        let target_dtype = array.options;
+
+        // If the array is already of the target dtype, then return the input node as-is.
+        if array.dtype() == target_dtype {
+            return Ok(Some(array.children()[0].clone()));
+        }
+
+        Ok(None)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use vortex_error::VortexResult;
+
+    use super::CastArrayReduce;
+    use crate::Array;
+    use crate::array::IntoArray;
+    use crate::arrays::ConstantArray;
+    use crate::arrays::ConstantVTable;
+    use crate::optimizer::ArrayOptimizer;
+    use crate::scalar_fns::BuiltinScalarFns;
+
+    #[test]
+    fn test_same_dtype() -> VortexResult<()> {
+        let mut optimizer = ArrayOptimizer::default();
+        optimizer.register_reduce_rule(CastArrayReduce);
+
+        let array = ConstantArray::new(true, 10).into_array();
+        let cast_same_dtype = array.cast(array.dtype().clone())?;
+
+        let optimized = optimizer.optimize_array(cast_same_dtype)?;
+        assert!(optimized.is::<ConstantVTable>());
+
+        Ok(())
+    }
+}

vortex-array/src/scalar_fns/cast/mod.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-mod array;
+pub(crate) mod array;
 
 use prost::Message;
 use vortex_dtype::DType;

vortex-array/src/scalar_fns/mod.rs

@@ -17,9 +17,8 @@ use crate::ArrayRef;
 use crate::arrays::ScalarFnArrayExt;
 use crate::expr::Expression;
 use crate::expr::ScalarFnExprExt;
-use crate::scalar_fns::cast::CastFn;
 
-mod cast;
+pub mod cast;
 
 /// A collection of built-in scalar functions that can be applied to expressions or arrays.
 pub trait BuiltinScalarFns: Sized {
@@ -29,12 +28,12 @@ pub trait BuiltinScalarFns: Sized {
 
 impl BuiltinScalarFns for Expression {
     fn cast(&self, dtype: DType) -> VortexResult<Expression> {
-        CastFn.try_new_expr(dtype, [self.clone()])
+        cast::CastFn.try_new_expr(dtype, [self.clone()])
     }
 }
 
 impl BuiltinScalarFns for ArrayRef {
     fn cast(&self, dtype: DType) -> VortexResult<Self> {
-        CastFn.try_new_array(self.len(), dtype, [self.clone()])
+        cast::CastFn.try_new_array(self.len(), dtype, [self.clone()])
     }
 }

vortex-array/src/session/mod.rs

@@ -23,6 +23,7 @@ use crate::arrays::StructVTable;
 use crate::arrays::VarBinVTable;
 use crate::arrays::VarBinViewVTable;
 use crate::optimizer::ArrayOptimizer;
+use crate::scalar_fns::cast::array::CastArrayReduce;
 use crate::vtable::ArrayVTable;
 use crate::vtable::ArrayVTableExt;
 
@@ -97,6 +98,9 @@ impl Default for ArraySession {
         optimizer.register_parent_rule(PrimitiveMaskedValidityRule);
         optimizer.register_parent_rule(DecimalMaskedValidityRule);
 
+        // Scalar function rules
+        optimizer.register_reduce_rule(CastArrayReduce);
+
         session
     }
 }