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

vortex_array/scalar_fn/fns/cast/
mod.rs
