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

vortex_array/scalar_fn/fns/cast/
mod.rs
