vortex_array/scalar_fn/fns/cast/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod kernel;
5
6use std::fmt::Formatter;
7
8pub use kernel::*;
9use prost::Message;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_err;
13use vortex_proto::expr as pb;
14use vortex_session::VortexSession;
15
16use crate::AnyColumnar;
17use crate::ArrayRef;
18use crate::CanonicalView;
19use crate::ColumnarView;
20use crate::ExecutionCtx;
21use crate::arrays::BoolVTable;
22use crate::arrays::ConstantArray;
23use crate::arrays::ConstantVTable;
24use crate::arrays::DecimalVTable;
25use crate::arrays::ExtensionVTable;
26use crate::arrays::FixedSizeListVTable;
27use crate::arrays::ListViewVTable;
28use crate::arrays::NullVTable;
29use crate::arrays::PrimitiveVTable;
30use crate::arrays::StructVTable;
31use crate::arrays::VarBinViewVTable;
32use crate::builtins::ArrayBuiltins;
33use crate::dtype::DType;
34use crate::expr::StatsCatalog;
35use crate::expr::cast;
36use crate::expr::expression::Expression;
37use crate::expr::lit;
38use crate::expr::stats::Stat;
39use crate::scalar_fn::Arity;
40use crate::scalar_fn::ChildName;
41use crate::scalar_fn::ExecutionArgs;
42use crate::scalar_fn::ReduceCtx;
43use crate::scalar_fn::ReduceNode;
44use crate::scalar_fn::ReduceNodeRef;
45use crate::scalar_fn::ScalarFnId;
46use crate::scalar_fn::ScalarFnVTable;
47
48/// A cast expression that converts values to a target data type.
49#[derive(Clone)]
50pub struct Cast;
51
52impl ScalarFnVTable for Cast {
53    type Options = DType;
54
55    fn id(&self) -> ScalarFnId {
56        ScalarFnId::from("vortex.cast")
57    }
58
59    fn serialize(&self, dtype: &DType) -> VortexResult<Option<Vec<u8>>> {
60        Ok(Some(
61            pb::CastOpts {
62                target: Some(dtype.try_into()?),
63            }
64            .encode_to_vec(),
65        ))
66    }
67
68    fn deserialize(
69        &self,
70        _metadata: &[u8],
71        session: &VortexSession,
72    ) -> VortexResult<Self::Options> {
73        let proto = pb::CastOpts::decode(_metadata)?.target;
74        DType::from_proto(
75            proto
76                .as_ref()
77                .ok_or_else(|| vortex_err!("Missing target dtype in Cast expression"))?,
78            session,
79        )
80    }
81
82    fn arity(&self, _options: &DType) -> Arity {
83        Arity::Exact(1)
84    }
85
86    fn child_name(&self, _instance: &DType, child_idx: usize) -> ChildName {
87        match child_idx {
88            0 => ChildName::from("input"),
89            _ => unreachable!("Invalid child index {} for Cast expression", child_idx),
90        }
91    }
92
93    fn fmt_sql(&self, dtype: &DType, expr: &Expression, f: &mut Formatter<'_>) -> std::fmt::Result {
94        write!(f, "cast(")?;
95        expr.children()[0].fmt_sql(f)?;
96        write!(f, " as {}", dtype)?;
97        write!(f, ")")
98    }
99
100    fn return_dtype(&self, dtype: &DType, _arg_dtypes: &[DType]) -> VortexResult<DType> {
101        Ok(dtype.clone())
102    }
103
104    fn execute(
105        &self,
106        target_dtype: &DType,
107        args: &dyn ExecutionArgs,
108        ctx: &mut ExecutionCtx,
109    ) -> VortexResult<ArrayRef> {
110        let input = args.get(0)?;
111
112        let Some(columnar) = input.as_opt::<AnyColumnar>() else {
113            return input.execute::<ArrayRef>(ctx)?.cast(target_dtype.clone());
114        };
115
116        match columnar {
117            ColumnarView::Canonical(canonical) => {
118                match cast_canonical(canonical.clone(), target_dtype, ctx)? {
119                    Some(result) => Ok(result),
120                    None => vortex_bail!(
121                        "No CastKernel to cast canonical array {} from {} to {}",
122                        canonical.as_ref().encoding_id(),
123                        canonical.as_ref().dtype(),
124                        target_dtype,
125                    ),
126                }
127            }
128            ColumnarView::Constant(constant) => match cast_constant(constant, target_dtype)? {
129                Some(result) => Ok(result),
130                None => vortex_bail!(
131                    "No CastReduce to cast constant array from {} to {}",
132                    constant.dtype(),
133                    target_dtype,
134                ),
135            },
136        }
137    }
138
139    fn reduce(
140        &self,
141        target_dtype: &DType,
142        node: &dyn ReduceNode,
143        _ctx: &dyn ReduceCtx,
144    ) -> VortexResult<Option<ReduceNodeRef>> {
145        // Collapse node if child is already the target type
146        let child = node.child(0);
147        if &child.node_dtype()? == target_dtype {
148            return Ok(Some(child));
149        }
150        Ok(None)
151    }
152
153    fn stat_expression(
154        &self,
155        dtype: &DType,
156        expr: &Expression,
157        stat: Stat,
158        catalog: &dyn StatsCatalog,
159    ) -> Option<Expression> {
160        match stat {
161            Stat::IsConstant
162            | Stat::IsSorted
163            | Stat::IsStrictSorted
164            | Stat::NaNCount
165            | Stat::Sum
166            | Stat::UncompressedSizeInBytes => expr.child(0).stat_expression(stat, catalog),
167            Stat::Max | Stat::Min => {
168                // We cast min/max to the new type
169                expr.child(0)
170                    .stat_expression(stat, catalog)
171                    .map(|x| cast(x, dtype.clone()))
172            }
173            Stat::NullCount => {
174                // if !expr.data().is_nullable() {
175                // NOTE(ngates): we should decide on the semantics here. In theory, the null
176                //  count of something cast to non-nullable will be zero. But if we return
177                //  that we know this to be zero, then a pruning predicate may eliminate data
178                //  that would otherwise have caused the cast to error.
179                // return Some(lit(0u64));
180                // }
181                None
182            }
183        }
184    }
185
186    fn validity(&self, dtype: &DType, expression: &Expression) -> VortexResult<Option<Expression>> {
187        Ok(Some(if dtype.is_nullable() {
188            expression.child(0).validity()?
189        } else {
190            lit(true)
191        }))
192    }
193
194    // This might apply a nullability
195    fn is_null_sensitive(&self, _instance: &DType) -> bool {
196        true
197    }
198}
199
200/// Cast a canonical array to the target dtype by dispatching to the appropriate
201/// [`CastReduce`] or [`CastKernel`] for each canonical encoding.
202fn cast_canonical(
203    canonical: CanonicalView<'_>,
204    dtype: &DType,
205    ctx: &mut ExecutionCtx,
206) -> VortexResult<Option<ArrayRef>> {
207    match canonical {
208        CanonicalView::Null(a) => <NullVTable as CastReduce>::cast(a, dtype),
209        CanonicalView::Bool(a) => <BoolVTable as CastReduce>::cast(a, dtype),
210        CanonicalView::Primitive(a) => <PrimitiveVTable as CastKernel>::cast(a, dtype, ctx),
211        CanonicalView::Decimal(a) => <DecimalVTable as CastKernel>::cast(a, dtype, ctx),
212        CanonicalView::VarBinView(a) => <VarBinViewVTable as CastReduce>::cast(a, dtype),
213        CanonicalView::List(a) => <ListViewVTable as CastReduce>::cast(a, dtype),
214        CanonicalView::FixedSizeList(a) => <FixedSizeListVTable as CastReduce>::cast(a, dtype),
215        CanonicalView::Struct(a) => <StructVTable as CastKernel>::cast(a, dtype, ctx),
216        CanonicalView::Extension(a) => <ExtensionVTable as CastReduce>::cast(a, dtype),
217    }
218}
219
220/// Cast a constant array by dispatching to its [`CastReduce`] implementation.
221fn cast_constant(array: &ConstantArray, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
222    <ConstantVTable as CastReduce>::cast(array, dtype)
223}
224
225#[cfg(test)]
226mod tests {
227    use vortex_buffer::buffer;
228    use vortex_error::VortexExpect as _;
229
230    use crate::IntoArray;
231    use crate::arrays::StructArray;
232    use crate::dtype::DType;
233    use crate::dtype::Nullability;
234    use crate::dtype::PType;
235    use crate::expr::Expression;
236    use crate::expr::cast;
237    use crate::expr::get_item;
238    use crate::expr::root;
239    use crate::expr::test_harness;
240
241    #[test]
242    fn dtype() {
243        let dtype = test_harness::struct_dtype();
244        assert_eq!(
245            cast(root(), DType::Bool(Nullability::NonNullable))
246                .return_dtype(&dtype)
247                .unwrap(),
248            DType::Bool(Nullability::NonNullable)
249        );
250    }
251
252    #[test]
253    fn replace_children() {
254        let expr = cast(root(), DType::Bool(Nullability::Nullable));
255        expr.with_children(vec![root()])
256            .vortex_expect("operation should succeed in test");
257    }
258
259    #[test]
260    fn evaluate() {
261        let test_array = StructArray::from_fields(&[
262            ("a", buffer![0i32, 1, 2].into_array()),
263            ("b", buffer![4i64, 5, 6].into_array()),
264        ])
265        .unwrap()
266        .into_array();
267
268        let expr: Expression = cast(
269            get_item("a", root()),
270            DType::Primitive(PType::I64, Nullability::NonNullable),
271        );
272        let result = test_array.apply(&expr).unwrap();
273
274        assert_eq!(
275            result.dtype(),
276            &DType::Primitive(PType::I64, Nullability::NonNullable)
277        );
278    }
279
280    #[test]
281    fn test_display() {
282        let expr = cast(
283            get_item("value", root()),
284            DType::Primitive(PType::I64, Nullability::NonNullable),
285        );
286        assert_eq!(expr.to_string(), "cast($.value as i64)");
287
288        let expr2 = cast(root(), DType::Bool(Nullability::Nullable));
289        assert_eq!(expr2.to_string(), "cast($ as bool?)");
290    }
291}
vortex_array/scalar_fn/fns/cast/mod.rs

vortex_array/scalar_fn/fns/cast/
mod.rs
