u

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs

@@ -106,16 +106,16 @@ where
     let len = array.len();
     let mut uninit_range = builder.uninit_range(len);
 
-    // SAFETY: We initialize all `len` values below via `decode_cast_into` and the patch loop.
+    // SAFETY: We initialize all `len` values below via `decode_map_into` and the patch loop.
     unsafe {
         uninit_range.append_mask(array.validity()?.execute_mask(len, ctx)?);
     }
 
-    // SAFETY: `decode_cast_into` writes a value to every slot in this range.
+    // SAFETY: `decode_map_into` writes a value to every slot in this range.
     let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, len) };
 
     let mut chunks = array.unpacked_chunks::<F>()?;
-    chunks.decode_cast_into(uninit_slice, |v: F| v.as_());
+    chunks.decode_map_into(uninit_slice, |v: F| v.as_());
 
     if let Some(patches) = array.patches() {
         apply_cast_patches_to_uninit_range::<F, T>(&mut uninit_range, &patches, ctx)?;

encodings/fastlanes/src/bitpacking/array/unpack_iter.rs

@@ -104,50 +104,6 @@ impl<T: BitPacked> BitUnpackedChunks<T> {
         )
     }
 
-    /// Decode all chunks (initial, full, and trailer), mapping each value through `f` and writing
-    /// the result into a differently-typed `output`.
-    ///
-    /// Unlike [`decode_into`](Self::decode_into), full chunks cannot be unpacked directly into the
-    /// output because the output element type `U` differs from the packed type `T`. Each chunk is
-    /// unpacked into the small internal scratch buffer (which stays resident in cache) and then
-    /// mapped value-by-value into `output`. This avoids materializing a full-length `T`-typed
-    /// intermediate buffer, which is the win for cast-on-decompression (e.g. bit-packed `u16` cast
-    /// to `u32`).
-    pub fn decode_cast_into<U: Copy>(&mut self, output: &mut [MaybeUninit<U>], f: impl Fn(T) -> U) {
-        let mut local_idx = 0;
-
-        if let Some(initial) = self.initial() {
-            for (dst, &src) in output[..initial.len()].iter_mut().zip(initial.iter()) {
-                dst.write(f(src));
-            }
-            local_idx += initial.len();
-        }
-
-        // `initial` already handled the only chunk when `num_chunks == 1`; mirror the guard in
-        // `decode_full_chunks_into_at` so we don't decode chunk 0 twice.
-        if self.num_chunks != 1 {
-            let mut chunks = self.full_chunks();
-            while let Some(chunk) = chunks.next() {
-                for (dst, &src) in output[local_idx..][..CHUNK_SIZE]
-                    .iter_mut()
-                    .zip(chunk.iter())
-                {
-                    dst.write(f(src));
-                }
-                local_idx += CHUNK_SIZE;
-            }
-        }
-
-        if let Some(trailer) = self.trailer() {
-            for (dst, &src) in output[local_idx..][..trailer.len()]
-                .iter_mut()
-                .zip(trailer.iter())
-            {
-                dst.write(f(src));
-            }
-        }
-    }
-
     pub fn full_chunks(&mut self) -> BitUnpackIterator<'_, T> {
         let elems_per_chunk = self.elems_per_chunk();
         let last_chunk_is_sliced = self.last_chunk_is_sliced() as usize;
@@ -161,6 +117,15 @@ impl<T: BitPacked> BitUnpackedChunks<T> {
             first_chunk_is_sliced,
         )
     }
+
+    /// Decode all chunks (initial, full, and trailer), mapping each value through `f` and writing
+    /// the result into a differently-typed `output`.
+    ///
+    /// Kept as a cast-oriented alias for callers that want the old name. Internal code can call
+    /// `decode_map_into` directly.
+    pub fn decode_cast_into<U: Copy>(&mut self, output: &mut [MaybeUninit<U>], f: impl Fn(T) -> U) {
+        self.decode_map_into(output, f);
+    }
 }
 
 impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
@@ -225,70 +190,79 @@ impl<T: PhysicalPType, S: UnpackStrategy<T>> UnpackedChunks<T, S> {
         })
     }
 
-    /// Decode all chunks (initial, full, and trailer) into the output range.
-    /// This consolidates the logic for handling all three chunk types in one place.
-    pub fn decode_into(&mut self, output: &mut [MaybeUninit<T>]) {
+    /// Decode all chunks (initial, full, and trailer), calling `write_chunk` for each concrete
+    /// unpacked chunk and its corresponding output range.
+    pub(crate) fn decode_chunks_into<U>(
+        &mut self,
+        output: &mut [MaybeUninit<U>],
+        mut write_chunk: impl FnMut(&[T], &mut [MaybeUninit<U>]),
+    ) {
+        debug_assert_eq!(output.len(), self.len);
         let mut local_idx = 0;
 
-        // Handle initial partial chunk if present
         if let Some(initial) = self.initial() {
-            local_idx = initial.len();
-
-            // TODO(connor): use `maybe_uninit_write_slice` feature when it gets stabilized.
-            // https://github.com/rust-lang/rust/issues/79995
-            // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
-            let init_initial: &[MaybeUninit<T>] = unsafe { mem::transmute(initial) };
-            output[..local_idx].copy_from_slice(init_initial);
+            let chunk_len = initial.len();
+            write_chunk(initial, &mut output[..chunk_len]);
+            local_idx += chunk_len;
         }
 
-        // Handle full chunks
-        local_idx = self.decode_full_chunks_into_at(output, local_idx);
+        if self.num_chunks != 1 {
+            let first_chunk_is_sliced = self.first_chunk_is_sliced();
+            let last_chunk_is_sliced = self.last_chunk_is_sliced();
+            let full_chunks_range =
+                (first_chunk_is_sliced as usize)..(self.num_chunks - last_chunk_is_sliced as usize);
+
+            let packed_slice: &[T::Physical] = buffer_as_slice(&self.packed);
+            let elems_per_chunk = self.elems_per_chunk();
+            for i in full_chunks_range {
+                let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk];
+
+                unsafe {
+                    let dst: &mut [MaybeUninit<T>] = &mut self.buffer;
+                    // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
+                    let dst: &mut [T::Physical] = mem::transmute(dst);
+                    self.strategy.unpack_chunk(self.bit_width, chunk, dst);
+                }
 
-        // Handle trailing partial chunk if present
-        if let Some(trailer) = self.trailer() {
-            // TODO(connor): use `maybe_uninit_write_slice` feature when it gets stabilized.
-            // https://github.com/rust-lang/rust/issues/79995
-            // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
-            let init_trailer: &[MaybeUninit<T>] = unsafe { mem::transmute(trailer) };
-            output[local_idx..][..init_trailer.len()].copy_from_slice(init_trailer);
+                // SAFETY: `unpack_chunk` initialized the whole scratch chunk above.
+                let unpacked: &[T] = unsafe { mem::transmute(&self.buffer[..]) };
+                write_chunk(unpacked, &mut output[local_idx..local_idx + CHUNK_SIZE]);
+                local_idx += CHUNK_SIZE;
+            }
         }
-    }
 
-    /// Unpack full chunks into output range starting at the given index.
-    /// Returns the next local index to write to.
-    fn decode_full_chunks_into_at(
-        &mut self,
-        output: &mut [MaybeUninit<T>],
-        start_idx: usize,
-    ) -> usize {
-        // If there's only one chunk it has been handled already by `initial` method
-        if self.num_chunks == 1 {
-            // Return the start_idx since initial already wrote everything.
-            return start_idx;
+        if let Some(trailer) = self.trailer() {
+            let chunk_len = trailer.len();
+            write_chunk(trailer, &mut output[local_idx..local_idx + chunk_len]);
+            local_idx += chunk_len;
         }
 
-        let first_chunk_is_sliced = self.first_chunk_is_sliced();
-
-        let last_chunk_is_sliced = self.last_chunk_is_sliced();
-        let full_chunks_range =
-            (first_chunk_is_sliced as usize)..(self.num_chunks - last_chunk_is_sliced as usize);
-
-        let mut local_idx = start_idx;
-
-        let packed_slice: &[T::Physical] = buffer_as_slice(&self.packed);
-        let elems_per_chunk = self.elems_per_chunk();
-        for i in full_chunks_range {
-            let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk];
+        debug_assert_eq!(local_idx, self.len);
+    }
 
-            unsafe {
-                let uninit_dst = &mut output[local_idx..local_idx + CHUNK_SIZE];
-                // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
-                let dst: &mut [T::Physical] = mem::transmute(uninit_dst);
-                self.strategy.unpack_chunk(self.bit_width, chunk, dst);
+    /// Decode all chunks (initial, full, and trailer), mapping each unpacked value through `f`.
+    pub(crate) fn decode_map_into<U>(
+        &mut self,
+        output: &mut [MaybeUninit<U>],
+        mut f: impl FnMut(T) -> U,
+    ) {
+        self.decode_chunks_into(output, |chunk, dst| {
+            for (dst, &src) in dst.iter_mut().zip(chunk.iter()) {
+                dst.write(f(src));
             }
-            local_idx += CHUNK_SIZE;
-        }
-        local_idx
+        });
+    }
+
+    /// Decode all chunks (initial, full, and trailer) into the output range.
+    /// This is the identity mapping of [`decode_map_into`](Self::decode_map_into).
+    pub fn decode_into(&mut self, output: &mut [MaybeUninit<T>]) {
+        self.decode_chunks_into(output, |chunk, dst| {
+            // TODO(connor): use `maybe_uninit_write_slice` feature when it gets stabilized.
+            // https://github.com/rust-lang/rust/issues/79995
+            // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
+            let initialized: &[MaybeUninit<T>] = unsafe { mem::transmute(chunk) };
+            dst.copy_from_slice(initialized);
+        });
     }
 
     /// Access last chunk of the array if the last chunk has fewer than 1024 due to slicing

encodings/fastlanes/src/bitpacking/compute/cast.rs

@@ -121,10 +121,15 @@ mod tests {
     use vortex_array::builtins::ArrayBuiltins;
     use vortex_array::compute::conformance::cast::test_cast_conformance;
     use vortex_array::dtype::DType;
+    use vortex_array::dtype::NativePType;
     use vortex_array::dtype::Nullability;
     use vortex_array::dtype::PType;
+    use vortex_array::match_each_integer_ptype;
+    use vortex_array::scalar_fn::fns::cast::CastKernel;
     use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
 
+    use crate::BitPacked;
     use crate::BitPackedArray;
     use crate::BitPackedData;
 
@@ -166,6 +171,65 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_cast_bitpacked_widening_integer_matrix() -> VortexResult<()> {
+        fn values<T: NativePType>(len: usize) -> PrimitiveArray {
+            PrimitiveArray::from_iter((0..len).map(|i| {
+                let value = if i % 17 == 0 { 31 } else { i % 8 };
+                <T as num_traits::FromPrimitive>::from_usize(value)
+                    .expect("test values fit every integer ptype")
+            }))
+        }
+
+        fn supported(src: PType, tgt: PType) -> bool {
+            src.is_int()
+                && tgt.is_int()
+                && tgt.byte_width() > src.byte_width()
+                && (src.is_unsigned_int() || tgt.is_signed_int())
+        }
+
+        let ptypes = [
+            PType::I8,
+            PType::I16,
+            PType::I32,
+            PType::I64,
+            PType::U8,
+            PType::U16,
+            PType::U32,
+            PType::U64,
+        ];
+        let lengths = [0, 1, 7, 1023, 1024, 1025, 2051];
+
+        for src in ptypes {
+            for tgt in ptypes {
+                if !supported(src, tgt) {
+                    continue;
+                }
+
+                for len in lengths {
+                    let source = match_each_integer_ptype!(src, |S| { values::<S>(len) });
+                    let source_ref = source.clone().into_array();
+                    let packed = bp(&source_ref, 3);
+                    let target = DType::Primitive(tgt, Nullability::NonNullable);
+
+                    let mut ctx = LEGACY_SESSION.create_execution_ctx();
+                    let casted =
+                        <BitPacked as CastKernel>::cast(packed.as_view(), &target, &mut ctx)?
+                            .expect(
+                                "supported widening integer cast should hit BitPacked CastKernel",
+                            );
+                    let reference = source_ref
+                        .cast(target)?
+                        .execute::<PrimitiveArray>(&mut ctx)?;
+
+                    assert_arrays_eq!(casted, reference);
+                }
+            }
+        }
+
+        Ok(())
+    }
+
     #[rstest]
     #[case(bp(&buffer![0u8, 10, 20, 30, 40, 50, 60, 63].into_array(), 6))]
     #[case(bp(&buffer![0u16, 100, 200, 300, 400, 500].into_array(), 9))]