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::Display;
7use std::fmt::Formatter;
8
9pub use kernel::*;
10use prost::Message;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_err;
14use vortex_proto::expr as pb;
15use vortex_session::VortexSession;
16use vortex_session::registry::CachedId;
17
18use crate::AnyColumnar;
19use crate::ArrayRef;
20use crate::ArrayView;
21use crate::CanonicalView;
22use crate::ColumnarView;
23use crate::ExecutionCtx;
24use crate::arrays::Bool;
25use crate::arrays::Constant;
26use crate::arrays::Decimal;
27use crate::arrays::Extension;
28use crate::arrays::FixedSizeList;
29use crate::arrays::ListView;
30use crate::arrays::Map;
31use crate::arrays::Null;
32use crate::arrays::Primitive;
33use crate::arrays::VarBinView;
34use crate::arrays::struct_::compute::cast::struct_cast;
35use crate::builtins::ArrayBuiltins;
36use crate::dtype::DType;
37use crate::expr::display::ExprDisplay;
38use crate::expr::expression::Expression;
39use crate::expr::lit;
40use crate::scalar_fn::Arity;
41use crate::scalar_fn::ChildName;
42use crate::scalar_fn::ExecutionArgs;
43use crate::scalar_fn::ReduceCtx;
44use crate::scalar_fn::ReduceNode;
45use crate::scalar_fn::ReduceNodeRef;
46use crate::scalar_fn::ScalarFnId;
47use crate::scalar_fn::ScalarFnVTable;
48use crate::scalar_fn::fns::literal::Literal;
49
50/// A cast expression that converts values to a target data type.
51#[derive(Clone)]
52pub struct Cast;
53
54impl ScalarFnVTable for Cast {
55    type Options = DType;
56
57    fn id(&self) -> ScalarFnId {
58        static ID: CachedId = CachedId::new("vortex.cast");
59        *ID
60    }
61
62    fn serialize(&self, dtype: &DType) -> VortexResult<Option<Vec<u8>>> {
63        Ok(Some(
64            pb::CastOpts {
65                target: Some(dtype.try_into()?),
66            }
67            .encode_to_vec(),
68        ))
69    }
70
71    fn deserialize(
72        &self,
73        _metadata: &[u8],
74        session: &VortexSession,
75    ) -> VortexResult<Self::Options> {
76        let proto = pb::CastOpts::decode(_metadata)?.target;
77        DType::from_proto(
78            proto
79                .as_ref()
80                .ok_or_else(|| vortex_err!("Missing target dtype in Cast expression"))?,
81            session,
82        )
83    }
84
85    fn arity(&self, _options: &DType) -> Arity {
86        Arity::Exact(1)
87    }
88
89    fn child_name(&self, _instance: &DType, child_idx: usize) -> ChildName {
90        match child_idx {
91            0 => ChildName::from("input"),
92            _ => unreachable!("Invalid child index {} for Cast expression", child_idx),
93        }
94    }
95
96    fn fmt_sql(
97        &self,
98        dtype: &DType,
99        expr: &dyn ExprDisplay,
100        f: &mut Formatter<'_>,
101    ) -> std::fmt::Result {
102        write!(f, "cast(")?;
103        Display::fmt(expr.display_child(0), f)?;
104        write!(f, " as {}", dtype)?;
105        write!(f, ")")
106    }
107
108    fn return_dtype(&self, dtype: &DType, _arg_dtypes: &[DType]) -> VortexResult<DType> {
109        Ok(dtype.clone())
110    }
111
112    fn execute(
113        &self,
114        target_dtype: &DType,
115        args: &dyn ExecutionArgs,
116        ctx: &mut ExecutionCtx,
117    ) -> VortexResult<ArrayRef> {
118        let input = args.get(0)?;
119
120        let Some(columnar) = input.as_opt::<AnyColumnar>() else {
121            return input.execute::<ArrayRef>(ctx)?.cast(target_dtype.clone());
122        };
123
124        match columnar {
125            ColumnarView::Canonical(canonical) => {
126                match cast_canonical(canonical, target_dtype, ctx)? {
127                    Some(result) => Ok(result),
128                    None => vortex_bail!(
129                        "No CastKernel to cast canonical array {} from {} to {}",
130                        canonical.to_array_ref().encoding_id(),
131                        canonical.to_array_ref().dtype(),
132                        target_dtype,
133                    ),
134                }
135            }
136            ColumnarView::Constant(constant) => match cast_constant(constant, target_dtype)? {
137                Some(result) => Ok(result),
138                None => vortex_bail!(
139                    "No CastReduce to cast constant array from {} to {}",
140                    constant.dtype(),
141                    target_dtype,
142                ),
143            },
144        }
145    }
146
147    fn reduce(
148        &self,
149        target_dtype: &DType,
150        node: &dyn ReduceNode,
151        _ctx: &dyn ReduceCtx,
152    ) -> VortexResult<Option<ReduceNodeRef>> {
153        // Collapse node if child is already the target type
154        let child = node.child(0);
155        if &child.node_dtype()? == target_dtype {
156            return Ok(Some(child));
157        }
158        Ok(None)
159    }
160
161    fn simplify_untyped(
162        &self,
163        target_dtype: &DType,
164        expr: &Expression,
165    ) -> VortexResult<Option<Expression>> {
166        let Some(scalar) = expr.child(0).as_opt::<Literal>() else {
167            return Ok(None);
168        };
169        // A failing cast (e.g. null to a non-nullable dtype) is left in place so the error
170        // surfaces at execution time rather than during optimization.
171        Ok(scalar.cast(target_dtype).ok().map(lit))
172    }
173
174    fn validity(&self, dtype: &DType, expression: &Expression) -> VortexResult<Option<Expression>> {
175        Ok(Some(if dtype.is_nullable() {
176            expression.child(0).validity()?
177        } else {
178            lit(true)
179        }))
180    }
181
182    fn is_strict(&self, _instance: &DType) -> bool {
183        // Cast options can pin a non-nullable output dtype instead of propagating nullability.
184        false
185    }
186}
187
188/// Cast a canonical array to the target dtype by dispatching to the appropriate
189/// [`CastKernel`] for each canonical encoding.
190///
191/// Canonical encodings that can manipulate validity directly all implement [`CastKernel`] —
192/// the kernel is the execution-time complement of their [`CastReduce`] rule and can compute
193/// statistics (e.g. min of the validity array) when the reduce rule had to give up.
194/// Encodings that delegate to scalars or storage (e.g. [`Null`], [`Constant`], [`Extension`])
195/// only implement [`CastReduce`] because they never need execution-level information.
196fn cast_canonical(
197    canonical: CanonicalView<'_>,
198    dtype: &DType,
199    ctx: &mut ExecutionCtx,
200) -> VortexResult<Option<ArrayRef>> {
201    match canonical {
202        CanonicalView::Null(a) => <Null as CastReduce>::cast(a, dtype),
203        CanonicalView::Bool(a) => <Bool as CastKernel>::cast(a, dtype, ctx),
204        CanonicalView::Primitive(a) => <Primitive as CastKernel>::cast(a, dtype, ctx),
205        CanonicalView::Decimal(a) => <Decimal as CastKernel>::cast(a, dtype, ctx),
206        CanonicalView::VarBinView(a) => <VarBinView as CastKernel>::cast(a, dtype, ctx),
207        CanonicalView::List(a) => <ListView as CastKernel>::cast(a, dtype, ctx),
208        CanonicalView::Map(a) => <Map as CastKernel>::cast(a, dtype, ctx),
209        CanonicalView::FixedSizeList(a) => <FixedSizeList as CastKernel>::cast(a, dtype, ctx),
210        CanonicalView::Struct(a) => struct_cast(a, dtype, ctx),
211        CanonicalView::Union(_) => {
212            todo!(
213                "TODO(connor)[Union]: implement Union casting with conformance coverage for outer \
214                 nullability changes, including validation of nullable-to-nonnullable casts"
215            )
216        }
217        CanonicalView::Extension(a) => <Extension as CastReduce>::cast(a, dtype),
218        CanonicalView::Variant(_) => {
219            vortex_bail!("Variant arrays don't support casting")
220        }
221    }
222}
223
224/// Cast a constant array by dispatching to its [`CastReduce`] implementation.
225fn cast_constant(array: ArrayView<Constant>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
226    <Constant as CastReduce>::cast(array, dtype)
227}
228
229#[cfg(test)]
230mod tests {
231    use vortex_buffer::buffer;
232    use vortex_error::VortexExpect as _;
233    use vortex_error::VortexResult;
234    use vortex_error::vortex_err;
235
236    use super::Cast;
237    use crate::IntoArray;
238    use crate::arrays::StructArray;
239    use crate::dtype::DType;
240    use crate::dtype::DecimalDType;
241    use crate::dtype::Nullability;
242    use crate::dtype::PType;
243    use crate::expr::Expression;
244    use crate::expr::cast;
245    use crate::expr::get_item;
246    use crate::expr::lit;
247    use crate::expr::root;
248    use crate::expr::test_harness;
249    use crate::scalar::DecimalValue;
250    use crate::scalar::Scalar;
251    use crate::scalar_fn::fns::literal::Literal;
252
253    #[test]
254    fn dtype() {
255        let dtype = test_harness::struct_dtype();
256        assert_eq!(
257            cast(root(), DType::Bool(Nullability::NonNullable))
258                .return_dtype(&dtype)
259                .unwrap(),
260            DType::Bool(Nullability::NonNullable)
261        );
262    }
263
264    #[test]
265    fn replace_children() {
266        let expr = cast(root(), DType::Bool(Nullability::Nullable));
267        expr.with_children(vec![root()])
268            .vortex_expect("operation should succeed in test");
269    }
270
271    #[test]
272    fn evaluate() {
273        let test_array = StructArray::from_fields(&[
274            ("a", buffer![0i32, 1, 2].into_array()),
275            ("b", buffer![4i64, 5, 6].into_array()),
276        ])
277        .unwrap()
278        .into_array();
279
280        let expr: Expression = cast(
281            get_item("a", root()),
282            DType::Primitive(PType::I64, Nullability::NonNullable),
283        );
284        let result = test_array.apply(&expr).unwrap();
285
286        assert_eq!(
287            result.dtype(),
288            &DType::Primitive(PType::I64, Nullability::NonNullable)
289        );
290    }
291
292    #[test]
293    fn simplify_folds_cast_of_literal() -> VortexResult<()> {
294        let expr = cast(
295            lit(3i32),
296            DType::Primitive(PType::F64, Nullability::NonNullable),
297        );
298        let optimized = expr.optimize(&test_harness::struct_dtype())?;
299
300        let scalar = optimized
301            .as_opt::<Literal>()
302            .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?;
303        assert_eq!(scalar, &Scalar::primitive(3.0f64, Nullability::NonNullable));
304        Ok(())
305    }
306
307    #[test]
308    fn simplify_folds_cast_of_decimal_literal() -> VortexResult<()> {
309        let decimal = Scalar::decimal(
310            DecimalValue::I128(319),
311            DecimalDType::new(3, 2),
312            Nullability::NonNullable,
313        );
314        let expr = cast(
315            lit(decimal),
316            DType::Primitive(PType::F64, Nullability::NonNullable),
317        );
318        let optimized = expr.optimize(&test_harness::struct_dtype())?;
319
320        let scalar = optimized
321            .as_opt::<Literal>()
322            .ok_or_else(|| vortex_err!("expected a bare literal, got {optimized}"))?;
323        assert_eq!(
324            scalar,
325            &Scalar::primitive(3.19f64, Nullability::NonNullable)
326        );
327        Ok(())
328    }
329
330    #[test]
331    fn simplify_leaves_failing_cast_unchanged() -> VortexResult<()> {
332        let target = DType::Primitive(PType::F64, Nullability::NonNullable);
333        let expr = cast(
334            lit(Scalar::null(DType::Primitive(
335                PType::I32,
336                Nullability::Nullable,
337            ))),
338            target.clone(),
339        );
340        let optimized = expr.optimize(&test_harness::struct_dtype())?;
341
342        assert!(optimized.as_opt::<Literal>().is_none());
343        assert_eq!(optimized.as_opt::<Cast>(), Some(&target));
344        Ok(())
345    }
346
347    #[test]
348    fn test_display() {
349        let expr = cast(
350            get_item("value", root()),
351            DType::Primitive(PType::I64, Nullability::NonNullable),
352        );
353        assert_eq!(expr.to_string(), "cast($.value as i64)");
354
355        let expr2 = cast(root(), DType::Bool(Nullability::Nullable));
356        assert_eq!(expr2.to_string(), "cast($ as bool?)");
357    }
358}
vortex_array/scalar_fn/fns/cast/mod.rs

vortex_array/scalar_fn/fns/cast/
mod.rs
