vortex_fastlanes/bitpacking/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::AsPrimitive;
5use vortex_array::ArrayRef;
6use vortex_array::ArrayView;
7use vortex_array::ExecutionCtx;
8use vortex_array::IntoArray;
9use vortex_array::builders::PrimitiveBuilder;
10use vortex_array::builtins::ArrayBuiltins;
11use vortex_array::dtype::DType;
12use vortex_array::dtype::PType;
13use vortex_array::match_each_integer_ptype;
14use vortex_array::scalar_fn::fns::cast::CastKernel;
15use vortex_array::scalar_fn::fns::cast::CastReduce;
16use vortex_array::validity::Validity;
17use vortex_error::VortexResult;
18
19use crate::bitpacking::BitPacked;
20use crate::bitpacking::array::BitPackedArrayExt;
21use crate::bitpacking::array::bitpack_decompress::unpack_map_into_builder;
22
23/// Returns `true` if casting `src` to `tgt` is a widening integer cast for which every value a
24/// bit-packed array can hold is guaranteed to be representable in `tgt` (so no per-value bounds
25/// check is needed). This holds when `tgt` is strictly wider and either the source is unsigned
26/// (always non-negative, fits in any wider type) or the target is also signed (sign-extension).
27fn is_widening_int_cast(src: PType, tgt: PType) -> bool {
28    src.is_int()
29        && tgt.is_int()
30        && tgt.byte_width() > src.byte_width()
31        && (src.is_unsigned_int() || tgt.is_signed_int())
32}
33
34fn build_with_validity(
35    array: ArrayView<'_, BitPacked>,
36    dtype: &DType,
37    new_validity: Validity,
38) -> VortexResult<ArrayRef> {
39    Ok(BitPacked::try_new(
40        array.packed().clone(),
41        dtype.as_ptype(),
42        new_validity,
43        array
44            .patches()
45            .map(|patches| patches.map_values(|values| values.cast(dtype.clone())))
46            .transpose()?,
47        array.bit_width(),
48        array.len(),
49        array.offset(),
50    )?
51    .into_array())
52}
53
54impl CastReduce for BitPacked {
55    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
56        if !array.dtype().eq_ignore_nullability(dtype) {
57            return Ok(None);
58        }
59        let Some(new_validity) = array
60            .validity()?
61            .trivially_cast_nullability(dtype.nullability(), array.len())?
62        else {
63            return Ok(None);
64        };
65        build_with_validity(array, dtype, new_validity).map(Some)
66    }
67}
68
69impl CastKernel for BitPacked {
70    fn cast(
71        array: ArrayView<'_, Self>,
72        dtype: &DType,
73        ctx: &mut ExecutionCtx,
74    ) -> VortexResult<Option<ArrayRef>> {
75        // Nullability-only change: keep the values bit-packed, just adjust validity.
76        if array.dtype().eq_ignore_nullability(dtype) {
77            let new_validity =
78                array
79                    .validity()?
80                    .cast_nullability(dtype.nullability(), array.len(), ctx)?;
81            return build_with_validity(array, dtype, new_validity).map(Some);
82        }
83
84        // Widening integer cast: unpack each FastLanes chunk into a cache-resident scratch buffer
85        // and cast-copy straight into the wide output, avoiding a full-length intermediate buffer
86        // and the generic cast kernel's bounds-check scan (unnecessary when widening).
87        let DType::Primitive(tgt, tgt_nullability) = dtype else {
88            return Ok(None);
89        };
90        let (tgt, tgt_nullability) = (*tgt, *tgt_nullability);
91        let src = array.dtype().as_ptype();
92        if !is_widening_int_cast(src, tgt) {
93            return Ok(None);
94        }
95
96        // Surface the standard error if a nullable source with nulls is cast to a non-nullable
97        // type; on success the per-value validity is handled inside the unpack below.
98        array
99            .validity()?
100            .cast_nullability(tgt_nullability, array.len(), ctx)?;
101
102        let result = match_each_integer_ptype!(tgt, |T| {
103            let mut builder = PrimitiveBuilder::<T>::with_capacity(tgt_nullability, array.len());
104            match_each_integer_ptype!(src, |F| {
105                unpack_map_into_builder::<F, T, _>(array, &mut builder, ctx, |v: F| v.as_())?;
106            });
107            builder.finish_into_primitive().into_array()
108        });
109        Ok(Some(result))
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use std::sync::LazyLock;
116
117    use rstest::rstest;
118    use vortex_array::ArrayRef;
119    use vortex_array::IntoArray;
120    use vortex_array::VortexSessionExecute;
121    use vortex_array::arrays::PrimitiveArray;
122    use vortex_array::assert_arrays_eq;
123    use vortex_array::builtins::ArrayBuiltins;
124    use vortex_array::compute::conformance::cast::test_cast_conformance;
125    use vortex_array::dtype::DType;
126    use vortex_array::dtype::NativePType;
127    use vortex_array::dtype::Nullability;
128    use vortex_array::dtype::PType;
129    use vortex_array::match_each_integer_ptype;
130    use vortex_buffer::buffer;
131    use vortex_error::VortexResult;
132    use vortex_session::VortexSession;
133
134    use crate::BitPackedArray;
135    use crate::BitPackedData;
136
137    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
138        let session = vortex_array::array_session();
139        crate::initialize(&session);
140        session
141    });
142
143    fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray {
144        BitPackedData::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap()
145    }
146
147    #[test]
148    fn test_cast_bitpacked_u8_to_u32() {
149        let packed = bp(&buffer![10u8, 20, 30, 40, 50, 60].into_array(), 6);
150
151        let casted = packed
152            .into_array()
153            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
154            .unwrap();
155        assert_eq!(
156            casted.dtype(),
157            &DType::Primitive(PType::U32, Nullability::NonNullable)
158        );
159
160        assert_arrays_eq!(
161            casted,
162            PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]),
163            &mut SESSION.create_execution_ctx()
164        );
165    }
166
167    #[test]
168    fn test_cast_bitpacked_nullable() {
169        let values = PrimitiveArray::from_option_iter([Some(5u16), None, Some(10), Some(15), None]);
170        let packed = bp(&values.into_array(), 4);
171
172        let casted = packed
173            .into_array()
174            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
175            .unwrap();
176        assert_eq!(
177            casted.dtype(),
178            &DType::Primitive(PType::U32, Nullability::Nullable)
179        );
180    }
181
182    /// End-to-end check that the real engine path `array.cast(target).execute()` routes through the
183    /// bit-packed widening pushdown and matches a plain primitive cast over the same values, across
184    /// every supported integer pair, several chunk-boundary lengths, and a sliced (offset > 0) case.
185    #[test]
186    fn test_cast_bitpacked_widening_via_execute() -> VortexResult<()> {
187        fn values<T: NativePType>(len: usize) -> PrimitiveArray {
188            PrimitiveArray::from_iter((0..len).map(|i| {
189                let value = if i % 17 == 0 { 31 } else { i % 8 };
190                <T as num_traits::FromPrimitive>::from_usize(value)
191                    .expect("test values fit every integer ptype")
192            }))
193        }
194
195        fn supported(src: PType, tgt: PType) -> bool {
196            src.is_int()
197                && tgt.is_int()
198                && tgt.byte_width() > src.byte_width()
199                && (src.is_unsigned_int() || tgt.is_signed_int())
200        }
201
202        let ptypes = [
203            PType::I8,
204            PType::I16,
205            PType::I32,
206            PType::I64,
207            PType::U8,
208            PType::U16,
209            PType::U32,
210            PType::U64,
211        ];
212        // Lengths exercise empty, sub-chunk, exact chunk, chunk+1, and multi-chunk-with-trailer.
213        let lengths = [0, 1, 7, 1023, 1024, 1025, 2051];
214
215        for src in ptypes {
216            for tgt in ptypes {
217                if !supported(src, tgt) {
218                    continue;
219                }
220
221                for len in lengths {
222                    let source = match_each_integer_ptype!(src, |S| { values::<S>(len) });
223                    let source_ref = source.into_array();
224                    let target = DType::Primitive(tgt, Nullability::NonNullable);
225                    let mut ctx = SESSION.create_execution_ctx();
226
227                    // Reference: plain primitive cast of the same values.
228                    let reference = source_ref
229                        .clone()
230                        .cast(target.clone())?
231                        .execute::<PrimitiveArray>(&mut ctx)?;
232
233                    // Candidate: bit-pack, then cast through the real engine. This dispatches to
234                    // `BitPacked`'s `CastKernel` widening pushdown.
235                    let packed = bp(&source_ref, 3).into_array();
236                    let casted = packed
237                        .cast(target.clone())?
238                        .execute::<PrimitiveArray>(&mut ctx)?;
239                    assert_arrays_eq!(casted, reference, &mut ctx);
240
241                    // Also exercise the sliced/offset path (offset > 0, trailer present).
242                    if len >= 4 {
243                        let lo = len / 4;
244                        let hi = len - len / 4;
245                        let sliced = bp(&source_ref, 3).into_array().slice(lo..hi)?;
246                        let casted = sliced
247                            .cast(target.clone())?
248                            .execute::<PrimitiveArray>(&mut ctx)?;
249                        let reference = source_ref
250                            .clone()
251                            .slice(lo..hi)?
252                            .cast(target.clone())?
253                            .execute::<PrimitiveArray>(&mut ctx)?;
254                        assert_arrays_eq!(casted, reference, &mut ctx);
255                    }
256                }
257            }
258        }
259
260        Ok(())
261    }
262
263    #[rstest]
264    #[case(bp(&buffer![0u8, 10, 20, 30, 40, 50, 60, 63].into_array(), 6))]
265    #[case(bp(&buffer![0u16, 100, 200, 300, 400, 500].into_array(), 9))]
266    #[case(bp(&buffer![0u32, 1000, 2000, 3000, 4000].into_array(), 12))]
267    #[case(bp(&PrimitiveArray::from_option_iter([Some(1u32), None, Some(7), Some(15), None]).into_array(), 4))]
268    fn test_cast_bitpacked_conformance(#[case] array: BitPackedArray) {
269        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
270    }
271}
vortex_fastlanes/bitpacking/compute/cast.rs

vortex_fastlanes/bitpacking/compute/
cast.rs
