vortex_array/arrays/bool/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::One;
5use num_traits::Zero;
6use vortex_buffer::BufferMut;
7use vortex_error::VortexResult;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::array::ArrayView;
13use crate::arrays::Bool;
14use crate::arrays::BoolArray;
15use crate::arrays::PrimitiveArray;
16use crate::arrays::bool::BoolArrayExt;
17use crate::dtype::DType;
18use crate::match_each_native_ptype;
19use crate::scalar_fn::fns::cast::CastKernel;
20use crate::scalar_fn::fns::cast::CastReduce;
21
22impl CastReduce for Bool {
23    fn cast(array: ArrayView<'_, Bool>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
24        if !dtype.is_boolean() {
25            return Ok(None);
26        }
27
28        let Some(new_validity) = array
29            .validity()?
30            .trivially_cast_nullability(dtype.nullability(), array.len())?
31        else {
32            return Ok(None);
33        };
34        Ok(Some(
35            BoolArray::new(array.to_bit_buffer(), new_validity).into_array(),
36        ))
37    }
38}
39
40impl CastKernel for Bool {
41    fn cast(
42        array: ArrayView<'_, Bool>,
43        dtype: &DType,
44        ctx: &mut ExecutionCtx,
45    ) -> VortexResult<Option<ArrayRef>> {
46        if dtype.is_boolean() {
47            let new_validity =
48                array
49                    .validity()?
50                    .cast_nullability(dtype.nullability(), array.len(), ctx)?;
51            return Ok(Some(
52                BoolArray::new(array.to_bit_buffer(), new_validity).into_array(),
53            ));
54        }
55
56        let DType::Primitive(new_ptype, new_nullability) = dtype else {
57            return Ok(None);
58        };
59
60        let new_validity =
61            array
62                .validity()?
63                .cast_nullability(*new_nullability, array.len(), ctx)?;
64
65        let bits = array.to_bit_buffer();
66        let len = bits.len();
67
68        Ok(Some(match_each_native_ptype!(*new_ptype, |T| {
69            let (one, zero) = (<T as One>::one(), <T as Zero>::zero());
70            let mut buffer = BufferMut::<T>::with_capacity(len);
71            buffer.extend(bits.iter().map(|v| if v { one } else { zero }));
72            PrimitiveArray::new(buffer.freeze(), new_validity).into_array()
73        })))
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use std::sync::LazyLock;
80
81    use rstest::rstest;
82    use vortex_buffer::buffer;
83    use vortex_session::VortexSession;
84
85    use crate::Canonical;
86    use crate::IntoArray;
87    use crate::VortexSessionExecute;
88    use crate::arrays::BoolArray;
89    use crate::assert_arrays_eq;
90    use crate::builtins::ArrayBuiltins;
91    use crate::compute::conformance::cast::test_cast_conformance;
92    use crate::dtype::DType;
93    use crate::dtype::Nullability;
94    use crate::dtype::PType;
95
96    static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
97
98    #[test]
99    fn try_cast_bool_success() {
100        let bool = BoolArray::from_iter(vec![Some(true), Some(false), Some(true)]);
101
102        let res = bool
103            .into_array()
104            .cast(DType::Bool(Nullability::NonNullable));
105        assert_arrays_eq!(
106            res.unwrap(),
107            BoolArray::from_iter([true, false, true]),
108            &mut SESSION.create_execution_ctx()
109        );
110    }
111
112    #[test]
113    fn try_cast_bool_fail() {
114        // When the validity array's min stat is not cached, the reduce rule defers and the
115        // failure surfaces during execution via the kernel (cast_nullability -> compute_min).
116        let bool = BoolArray::from_iter(vec![Some(true), Some(false), None]);
117        let mut ctx = SESSION.create_execution_ctx();
118        let result = bool
119            .into_array()
120            .cast(DType::Bool(Nullability::NonNullable))
121            .and_then(|a| a.execute::<Canonical>(&mut ctx).map(|c| c.into_array()));
122        assert!(result.is_err(), "Expected error, got: {result:?}");
123    }
124
125    #[rstest]
126    #[case(BoolArray::from_iter(vec![true, false, true, true, false]))]
127    #[case(BoolArray::from_iter(vec![Some(true), Some(false), None, Some(true), None]))]
128    #[case(BoolArray::from_iter(vec![true]))]
129    #[case(BoolArray::from_iter(vec![false, false]))]
130    fn test_cast_bool_conformance(#[case] array: BoolArray) {
131        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
132    }
133
134    #[rstest]
135    #[case(PType::I8)]
136    #[case(PType::I32)]
137    #[case(PType::I64)]
138    #[case(PType::U8)]
139    #[case(PType::U64)]
140    #[case(PType::F32)]
141    #[case(PType::F64)]
142    fn cast_bool_to_primitive(#[case] target: PType) {
143        let mut ctx = SESSION.create_execution_ctx();
144        let arr = BoolArray::from_iter(vec![true, false, true]).into_array();
145        let out = arr
146            .cast(DType::Primitive(target, Nullability::NonNullable))
147            .unwrap();
148        let expected = match target {
149            PType::I8 => buffer![1i8, 0, 1].into_array(),
150            PType::I32 => buffer![1i32, 0, 1].into_array(),
151            PType::I64 => buffer![1i64, 0, 1].into_array(),
152            PType::U8 => buffer![1u8, 0, 1].into_array(),
153            PType::U64 => buffer![1u64, 0, 1].into_array(),
154            PType::F32 => buffer![1f32, 0., 1.].into_array(),
155            PType::F64 => buffer![1f64, 0., 1.].into_array(),
156            _ => unreachable!("unexpected fixture type"),
157        };
158        assert_arrays_eq!(out, expected, &mut ctx);
159    }
160}
vortex_array/arrays/bool/compute/cast.rs

vortex_array/arrays/bool/compute/
cast.rs
