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_in(
104                tgt_nullability,
105                array.len(),
106                ctx.allocator(),
107            );
108            match_each_integer_ptype!(src, |F| {
109                unpack_map_into_builder::<F, T, _>(array, &mut builder, ctx, |v: F| v.as_())?;
110            });
111            builder.finish_into_primitive().into_array()
112        });
113        Ok(Some(result))
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use std::sync::LazyLock;
120
121    use rstest::rstest;
122    use vortex_array::ArrayRef;
123    use vortex_array::IntoArray;
124    use vortex_array::VortexSessionExecute;
125    use vortex_array::arrays::PrimitiveArray;
126    use vortex_array::assert_arrays_eq;
127    use vortex_array::builtins::ArrayBuiltins;
128    use vortex_array::compute::conformance::cast::test_cast_conformance;
129    use vortex_array::dtype::DType;
130    use vortex_array::dtype::NativePType;
131    use vortex_array::dtype::Nullability;
132    use vortex_array::dtype::PType;
133    use vortex_array::match_each_integer_ptype;
134    use vortex_buffer::buffer;
135    use vortex_error::VortexResult;
136    use vortex_session::VortexSession;
137
138    use crate::BitPackedArray;
139    use crate::BitPackedData;
140
141    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
142        let session = vortex_array::array_session();
143        crate::initialize(&session);
144        session
145    });
146
147    fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedArray {
148        BitPackedData::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap()
149    }
150
151    #[test]
152    fn test_cast_bitpacked_u8_to_u32() {
153        let packed = bp(&buffer![10u8, 20, 30, 40, 50, 60].into_array(), 6);
154
155        let casted = packed
156            .into_array()
157            .cast(DType::Primitive(PType::U32, Nullability::NonNullable))
158            .unwrap();
159        assert_eq!(
160            casted.dtype(),
161            &DType::Primitive(PType::U32, Nullability::NonNullable)
162        );
163
164        assert_arrays_eq!(
165            casted,
166            PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]),
167            &mut SESSION.create_execution_ctx()
168        );
169    }
170
171    #[test]
172    fn test_cast_bitpacked_nullable() {
173        let values = PrimitiveArray::from_option_iter([Some(5u16), None, Some(10), Some(15), None]);
174        let packed = bp(&values.into_array(), 4);
175
176        let casted = packed
177            .into_array()
178            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
179            .unwrap();
180        assert_eq!(
181            casted.dtype(),
182            &DType::Primitive(PType::U32, Nullability::Nullable)
183        );
184    }
185
186    /// End-to-end check that the real engine path `array.cast(target).execute()` routes through the
187    /// bit-packed widening pushdown and matches a plain primitive cast over the same values, across
188    /// every supported integer pair, several chunk-boundary lengths, and a sliced (offset > 0) case.
189    #[test]
190    fn test_cast_bitpacked_widening_via_execute() -> VortexResult<()> {
191        fn values<T: NativePType>(len: usize) -> PrimitiveArray {
192            PrimitiveArray::from_iter((0..len).map(|i| {
193                let value = if i % 17 == 0 { 31 } else { i % 8 };
194                <T as num_traits::FromPrimitive>::from_usize(value)
195                    .expect("test values fit every integer ptype")
196            }))
197        }
198
199        fn supported(src: PType, tgt: PType) -> bool {
200            src.is_int()
201                && tgt.is_int()
202                && tgt.byte_width() > src.byte_width()
203                && (src.is_unsigned_int() || tgt.is_signed_int())
204        }
205
206        let ptypes = [
207            PType::I8,
208            PType::I16,
209            PType::I32,
210            PType::I64,
211            PType::U8,
212            PType::U16,
213            PType::U32,
214            PType::U64,
215        ];
216        // Lengths exercise empty, sub-chunk, exact chunk, chunk+1, and multi-chunk-with-trailer.
217        let lengths = [0, 1, 7, 1023, 1024, 1025, 2051];
218
219        for src in ptypes {
220            for tgt in ptypes {
221                if !supported(src, tgt) {
222                    continue;
223                }
224
225                for len in lengths {
226                    let source = match_each_integer_ptype!(src, |S| { values::<S>(len) });
227                    let source_ref = source.into_array();
228                    let target = DType::Primitive(tgt, Nullability::NonNullable);
229                    let mut ctx = SESSION.create_execution_ctx();
230
231                    // Reference: plain primitive cast of the same values.
232                    let reference = source_ref
233                        .clone()
234                        .cast(target.clone())?
235                        .execute::<PrimitiveArray>(&mut ctx)?;
236
237                    // Candidate: bit-pack, then cast through the real engine. This dispatches to
238                    // `BitPacked`'s `CastKernel` widening pushdown.
239                    let packed = bp(&source_ref, 3).into_array();
240                    let casted = packed
241                        .cast(target.clone())?
242                        .execute::<PrimitiveArray>(&mut ctx)?;
243                    assert_arrays_eq!(casted, reference, &mut ctx);
244
245                    // Also exercise the sliced/offset path (offset > 0, trailer present).
246                    if len >= 4 {
247                        let lo = len / 4;
248                        let hi = len - len / 4;
249                        let sliced = bp(&source_ref, 3).into_array().slice(lo..hi)?;
250                        let casted = sliced
251                            .cast(target.clone())?
252                            .execute::<PrimitiveArray>(&mut ctx)?;
253                        let reference = source_ref
254                            .clone()
255                            .slice(lo..hi)?
256                            .cast(target.clone())?
257                            .execute::<PrimitiveArray>(&mut ctx)?;
258                        assert_arrays_eq!(casted, reference, &mut ctx);
259                    }
260                }
261            }
262        }
263
264        Ok(())
265    }
266
267    #[rstest]
268    #[case(bp(&buffer![0u8, 10, 20, 30, 40, 50, 60, 63].into_array(), 6))]
269    #[case(bp(&buffer![0u16, 100, 200, 300, 400, 500].into_array(), 9))]
270    #[case(bp(&buffer![0u32, 1000, 2000, 3000, 4000].into_array(), 12))]
271    #[case(bp(&PrimitiveArray::from_option_iter([Some(1u32), None, Some(7), Some(15), None]).into_array(), 4))]
272    fn test_cast_bitpacked_conformance(#[case] array: BitPackedArray) {
273        test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx());
274    }
275}
vortex_fastlanes/bitpacking/compute/cast.rs

vortex_fastlanes/bitpacking/compute/
cast.rs
