f

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

Author: joseph-isaacs

vortex-array/src/arrays/primitive/compute/cast.rs

@@ -1,13 +1,19 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::mem::align_of;
+use std::mem::size_of;
+
 use num_traits::AsPrimitive;
 use num_traits::NumCast;
+use vortex_buffer::BitBuffer;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
+use vortex_buffer::lane_ops_indexed::ReinterpretSink;
 use vortex_buffer::lane_ops_indexed::map_no_validity;
 use vortex_buffer::lane_ops_indexed::try_map_no_validity;
 use vortex_buffer::lane_ops_indexed::try_map_with_mask;
+use vortex_buffer::lane_ops_indexed::try_map_with_mask_in_place;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_err;
@@ -147,9 +153,21 @@ where
     // Skip the fallible kernel when the conversion is infallible by type alone (widening) or
     // when cached min/max prove every value fits in `T`.
     let target_dtype = DType::Primitive(T::PTYPE, Nullability::NonNullable);
-    if casts_losslessly_to(F::PTYPE, T::PTYPE)
-        || cached_values_fit_in(array, &target_dtype) == Some(true)
+    let infallible = casts_losslessly_to(F::PTYPE, T::PTYPE)
+        || cached_values_fit_in(array, &target_dtype) == Some(true);
+
+    // Same-bit-width in-place fast path: when F and T have the same byte width and the
+    // buffer is uniquely owned, mutate in place and transmute the wrapper. Saves the
+    // output allocation. Falls through to the out-of-place path when the buffer is shared
+    // (the common case under the current borrow-based kernel API).
+    let same_bit_width = F::PTYPE.byte_width() == T::PTYPE.byte_width();
+    if same_bit_width
+        && let Ok(buffer_mut) = array.into_owned().try_into_buffer_mut::<F>()
     {
+        return cast_buffer_in_place::<F, T>(buffer_mut, array, new_validity, ctx, infallible);
+    }
+
+    if infallible {
         let mut buffer = BufferMut::<T>::with_capacity(values.len());
         // Truncating `as`-cast — safe here because stats prove every valid value fits.
         // Null lanes' underlying garbage gets truncated/wrapped (harmless: the result
@@ -204,6 +222,72 @@ where
     Ok(PrimitiveArray::new(buffer, new_validity).into_array())
 }
 
+/// In-place cast of an owned `BufferMut<F>` to `BufferMut<T>` when `F` and `T` have the
+/// same byte width. Each slot is read as `F`, converted, and written back as `T`-bits
+/// using `BufferMut`'s transmute family. Avoids allocating a second output buffer.
+///
+/// The caller has already verified `F::PTYPE.byte_width() == T::PTYPE.byte_width()`.
+fn cast_buffer_in_place<F, T>(
+    buffer: BufferMut<F>,
+    array: ArrayView<'_, Primitive>,
+    new_validity: Validity,
+    ctx: &mut ExecutionCtx,
+    infallible: bool,
+) -> VortexResult<ArrayRef>
+where
+    F: NativePType + AsPrimitive<T>,
+    T: NativePType,
+{
+    debug_assert_eq!(size_of::<F>(), size_of::<T>());
+    debug_assert_eq!(align_of::<F>(), align_of::<T>());
+
+    if infallible {
+        // `map_each_in_place` does the BufferMut<F> → BufferMut<T> transmute internally
+        // (same size + alignment for primitives of equal byte width) and walks each slot
+        // with the closure.
+        let result: BufferMut<T> = buffer.map_each_in_place(|v: F| v.as_());
+        return Ok(PrimitiveArray::new(result.freeze(), new_validity).into_array());
+    }
+
+    let mask = array.validity()?.execute_mask(array.len(), ctx)?;
+    let overflow = || {
+        vortex_err!(
+            Compute: "Cannot cast {} to {} — value exceeds target range",
+            F::PTYPE, T::PTYPE,
+        )
+    };
+
+    // All-null short-circuit: zero out the buffer and skip the conversion loop entirely.
+    if matches!(mask, Mask::AllFalse(_)) {
+        // SAFETY: same size + alignment by NativePType same-byte-width invariant.
+        let mut t_buf: BufferMut<T> = unsafe { buffer.transmute::<T>() };
+        t_buf.as_mut_slice().fill(T::zero());
+        return Ok(PrimitiveArray::new(t_buf.freeze(), new_validity).into_array());
+    }
+
+    let bit_buffer = match &mask {
+        Mask::AllTrue(n) => BitBuffer::new_set(*n),
+        Mask::AllFalse(_) => unreachable!("handled above"),
+        Mask::Values(m) => m.bit_buffer().clone(),
+    };
+
+    let mut buffer = buffer;
+    try_map_with_mask_in_place(
+        ReinterpretSink::<F, T>::new(buffer.as_mut_slice()),
+        &bit_buffer,
+        |f_val: F, valid| -> Option<T> {
+            <T as NumCast>::from(f_val).or_else(|| (!valid).then(T::zero))
+        },
+    )
+    .map_err(|_| overflow())?;
+
+    // SAFETY: same size + alignment for NativePType same-byte-width pairs. Every F-slot
+    // now holds a valid T-bit pattern because `ReinterpretSink::set_unchecked` wrote a
+    // real `T` at every visited lane.
+    let result: BufferMut<T> = unsafe { buffer.transmute::<T>() };
+    Ok(PrimitiveArray::new(result.freeze(), new_validity).into_array())
+}
+
 fn reinterpret(
     array: ArrayView<'_, Primitive>,
     new_ptype: PType,