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

vortex_array/scalar_fn/fns/cast/
mod.rs
