Add cast expression (#3440)

Author: delta003

vortex-expr/src/cast.rs

@@ -0,0 +1,160 @@
+use std::any::Any;
+use std::fmt::Display;
+use std::sync::Arc;
+
+use vortex_array::compute::cast as compute_cast;
+use vortex_array::{Array, ArrayRef};
+use vortex_dtype::DType;
+use vortex_error::{VortexExpect, VortexResult};
+
+use crate::{ExprRef, VortexExpr};
+
+#[derive(Debug, Eq, Hash)]
+#[allow(clippy::derived_hash_with_manual_eq)]
+pub struct Cast {
+    target: DType,
+    child: ExprRef,
+}
+
+impl Cast {
+    pub fn new_expr(child: ExprRef, target: DType) -> ExprRef {
+        Arc::new(Self { target, child })
+    }
+}
+
+impl PartialEq for Cast {
+    fn eq(&self, other: &Self) -> bool {
+        self.target.eq(&other.target) && self.child.eq(&other.child)
+    }
+}
+
+impl Display for Cast {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "cast({}, {})", self.child, self.target)
+    }
+}
+
+#[cfg(feature = "proto")]
+pub(crate) mod proto {
+    use vortex_dtype::DType;
+    use vortex_error::{VortexResult, vortex_bail, vortex_err};
+    use vortex_proto::expr::kind;
+    use vortex_proto::expr::kind::Kind;
+
+    use crate::cast::Cast;
+    use crate::{ExprDeserialize, ExprRef, ExprSerializable, Id};
+
+    pub(crate) struct CastSerde;
+
+    impl Id for CastSerde {
+        fn id(&self) -> &'static str {
+            "cast"
+        }
+    }
+
+    impl ExprDeserialize for CastSerde {
+        fn deserialize(&self, kind: &Kind, children: Vec<ExprRef>) -> VortexResult<ExprRef> {
+            let Kind::Cast(kind::Cast { target }) = kind else {
+                vortex_bail!("wrong kind {:?}, want cast", kind)
+            };
+            let target: DType = target
+                .as_ref()
+                .ok_or_else(|| vortex_err!("empty target dtype"))?
+                .try_into()?;
+
+            Ok(Cast::new_expr(children[0].clone(), target))
+        }
+    }
+
+    impl ExprSerializable for Cast {
+        fn id(&self) -> &'static str {
+            CastSerde.id()
+        }
+
+        fn serialize_kind(&self) -> VortexResult<Kind> {
+            Ok(Kind::Cast(kind::Cast {
+                target: Some((&self.target).into()),
+            }))
+        }
+    }
+}
+
+impl VortexExpr for Cast {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn unchecked_evaluate(&self, batch: &dyn Array) -> VortexResult<ArrayRef> {
+        let array = self.child.evaluate(batch)?;
+        compute_cast(&array, &self.target)
+    }
+
+    fn children(&self) -> Vec<&ExprRef> {
+        vec![&self.child]
+    }
+
+    fn replacing_children(self: Arc<Self>, mut children: Vec<ExprRef>) -> ExprRef {
+        Self::new_expr(
+            children
+                .pop()
+                .vortex_expect("Cast::replacing_children should have one child"),
+            self.target.clone(),
+        )
+    }
+
+    fn return_dtype(&self, _scope_dtype: &DType) -> VortexResult<DType> {
+        Ok(self.target.clone())
+    }
+}
+
+pub fn cast(child: ExprRef, target: DType) -> ExprRef {
+    Cast::new_expr(child, target)
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_array::IntoArray;
+    use vortex_array::arrays::StructArray;
+    use vortex_buffer::buffer;
+    use vortex_dtype::{DType, Nullability, PType};
+
+    use crate::{ExprRef, cast, get_item, ident, test_harness};
+
+    #[test]
+    fn dtype() {
+        let dtype = test_harness::struct_dtype();
+        assert_eq!(
+            cast(ident(), DType::Bool(Nullability::NonNullable))
+                .return_dtype(&dtype)
+                .unwrap(),
+            DType::Bool(Nullability::NonNullable)
+        );
+    }
+
+    #[test]
+    fn replace_children() {
+        let expr = cast(ident(), DType::Bool(Nullability::Nullable));
+        let _ = expr.replacing_children(vec![ident()]);
+    }
+
+    #[test]
+    fn evaluate() {
+        let test_array = StructArray::from_fields(&[
+            ("a", buffer![0i32, 1, 2].into_array()),
+            ("b", buffer![4i64, 5, 6].into_array()),
+        ])
+        .unwrap()
+        .into_array();
+
+        let expr: ExprRef = cast(
+            get_item("a", ident()),
+            DType::Primitive(PType::I64, Nullability::NonNullable),
+        );
+        let result = expr.unchecked_evaluate(&test_array).unwrap();
+
+        assert_eq!(
+            result.dtype(),
+            &DType::Primitive(PType::I64, Nullability::NonNullable)
+        );
+    }
+}

vortex-expr/src/lib.rs

@@ -7,6 +7,7 @@ use dyn_hash::DynHash;
 mod binary;
 
 mod between;
+mod cast;
 mod field;
 pub mod forms;
 mod get_item;
@@ -27,6 +28,7 @@ pub mod traversal;
 
 pub use between::*;
 pub use binary::*;
+pub use cast::*;
 pub use get_item::*;
 pub use identity::*;
 pub use is_null::*;

vortex-proto/proto/expr.proto

@@ -3,6 +3,7 @@ syntax = "proto3";
 package vortex.expr;
 
 import "scalar.proto";
+import "dtype.proto";
 
 option java_package = "dev.vortex.proto";
 option java_outer_classname = "ExprProtos";
@@ -27,6 +28,7 @@ message Kind {
     Between between = 8;
     Like like = 9;
     IsNull is_null = 10;
+    Cast cast = 11;
   }
 
   message Literal {
@@ -77,6 +79,10 @@ message Kind {
   message IsNull {
 
   }
+
+  message Cast {
+    vortex.dtype.DType target = 1;
+  }
 }
 
 

vortex-proto/src/generated/vortex.expr.rs

@@ -12,7 +12,7 @@ pub struct Expr {
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Kind {
     /// This enum is very unstable, and will likely be replaced with something more extensible.
-    #[prost(oneof = "kind::Kind", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10")]
+    #[prost(oneof = "kind::Kind", tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11")]
     pub kind: ::core::option::Option<kind::Kind>,
 }
 /// Nested message and enum types in `Kind`.
@@ -54,6 +54,11 @@ pub mod kind {
     }
     #[derive(Clone, Copy, PartialEq, ::prost::Message)]
     pub struct IsNull {}
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct Cast {
+        #[prost(message, optional, tag = "1")]
+        pub target: ::core::option::Option<super::super::dtype::DType>,
+    }
     #[derive(
         Clone,
         Copy,
@@ -131,5 +136,7 @@ pub mod kind {
         Like(Like),
         #[prost(message, tag = "10")]
         IsNull(IsNull),
+        #[prost(message, tag = "11")]
+        Cast(Cast),
     }
 }