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::ArrayView;
19use crate::CanonicalView;
20use crate::ColumnarView;
21use crate::ExecutionCtx;
22use crate::arrays::Bool;
23use crate::arrays::Constant;
24use crate::arrays::Decimal;
25use crate::arrays::Extension;
26use crate::arrays::FixedSizeList;
27use crate::arrays::ListView;
28use crate::arrays::Null;
29use crate::arrays::Primitive;
30use crate::arrays::Struct;
31use crate::arrays::VarBinView;
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::new("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, target_dtype, ctx)? {
119                    Some(result) => Ok(result),
120                    None => vortex_bail!(
121                        "No CastKernel to cast canonical array {} from {} to {}",
122                        canonical.to_array_ref().encoding_id(),
123                        canonical.to_array_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/// [`CastKernel`] for each canonical encoding.
202///
203/// Canonical encodings that can manipulate validity directly all implement [`CastKernel`] —
204/// the kernel is the execution-time complement of their [`CastReduce`] rule and can compute
205/// statistics (e.g. min of the validity array) when the reduce rule had to give up.
206/// Encodings that delegate to scalars or storage (e.g. [`Null`], [`Constant`], [`Extension`])
207/// only implement [`CastReduce`] because they never need execution-level information.
208fn cast_canonical(
209    canonical: CanonicalView<'_>,
210    dtype: &DType,
211    ctx: &mut ExecutionCtx,
212) -> VortexResult<Option<ArrayRef>> {
213    match canonical {
214        CanonicalView::Null(a) => <Null as CastReduce>::cast(a, dtype),
215        CanonicalView::Bool(a) => <Bool as CastKernel>::cast(a, dtype, ctx),
216        CanonicalView::Primitive(a) => <Primitive as CastKernel>::cast(a, dtype, ctx),
217        CanonicalView::Decimal(a) => <Decimal as CastKernel>::cast(a, dtype, ctx),
218        CanonicalView::VarBinView(a) => <VarBinView as CastKernel>::cast(a, dtype, ctx),
219        CanonicalView::List(a) => <ListView as CastKernel>::cast(a, dtype, ctx),
220        CanonicalView::FixedSizeList(a) => <FixedSizeList as CastKernel>::cast(a, dtype, ctx),
221        CanonicalView::Struct(a) => <Struct as CastKernel>::cast(a, dtype, ctx),
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
239    use crate::IntoArray;
240    use crate::arrays::StructArray;
241    use crate::dtype::DType;
242    use crate::dtype::Nullability;
243    use crate::dtype::PType;
244    use crate::expr::Expression;
245    use crate::expr::cast;
246    use crate::expr::get_item;
247    use crate::expr::root;
248    use crate::expr::test_harness;
249
250    #[test]
251    fn dtype() {
252        let dtype = test_harness::struct_dtype();
253        assert_eq!(
254            cast(root(), DType::Bool(Nullability::NonNullable))
255                .return_dtype(&dtype)
256                .unwrap(),
257            DType::Bool(Nullability::NonNullable)
258        );
259    }
260
261    #[test]
262    fn replace_children() {
263        let expr = cast(root(), DType::Bool(Nullability::Nullable));
264        expr.with_children(vec![root()])
265            .vortex_expect("operation should succeed in test");
266    }
267
268    #[test]
269    fn evaluate() {
270        let test_array = StructArray::from_fields(&[
271            ("a", buffer![0i32, 1, 2].into_array()),
272            ("b", buffer![4i64, 5, 6].into_array()),
273        ])
274        .unwrap()
275        .into_array();
276
277        let expr: Expression = cast(
278            get_item("a", root()),
279            DType::Primitive(PType::I64, Nullability::NonNullable),
280        );
281        let result = test_array.apply(&expr).unwrap();
282
283        assert_eq!(
284            result.dtype(),
285            &DType::Primitive(PType::I64, Nullability::NonNullable)
286        );
287    }
288
289    #[test]
290    fn test_display() {
291        let expr = cast(
292            get_item("value", root()),
293            DType::Primitive(PType::I64, Nullability::NonNullable),
294        );
295        assert_eq!(expr.to_string(), "cast($.value as i64)");
296
297        let expr2 = cast(root(), DType::Bool(Nullability::Nullable));
298        assert_eq!(expr2.to_string(), "cast($ as bool?)");
299    }
300}
vortex_array/scalar_fn/fns/cast/mod.rs

vortex_array/scalar_fn/fns/cast/
mod.rs
