refactor(array): allocate execution outputs through context

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-array/src/arrays/decimal/compute/between.rs

@@ -27,7 +27,7 @@ impl BetweenKernel for Decimal {
         lower: &ArrayRef,
         upper: &ArrayRef,
         options: &BetweenOptions,
-        _ctx: &mut ExecutionCtx,
+        ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         // NOTE: We know that the precision and scale were already checked to be equal by the main
         // `between` entrypoint function.
@@ -41,7 +41,7 @@ impl BetweenKernel for Decimal {
             arr.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability();
 
         match_each_decimal_value_type!(arr.values_type(), |D| {
-            between_unpack::<D>(arr, lower, upper, nullability, options)
+            between_unpack::<D>(arr, lower, upper, nullability, options, ctx)
         })
     }
 }
@@ -52,6 +52,7 @@ fn between_unpack<T: NativeDecimalType>(
     upper: Scalar,
     nullability: Nullability,
     options: &BetweenOptions,
+    ctx: &mut ExecutionCtx,
 ) -> VortexResult<Option<ArrayRef>> {
     let Some(lower_dv) = lower.as_decimal().decimal_value() else {
         // Null lower bound — fall back to canonical path.
@@ -119,6 +120,7 @@ fn between_unpack<T: NativeDecimalType>(
         nullability,
         lower_op,
         upper_op,
+        ctx,
     )))
 }
 
@@ -129,15 +131,20 @@ fn between_impl<T: NativeDecimalType>(
     nullability: Nullability,
     lower_op: impl Fn(T, T) -> bool,
     upper_op: impl Fn(T, T) -> bool,
+    ctx: &mut ExecutionCtx,
 ) -> ArrayRef {
     let buffer = arr.buffer::<T>();
     BoolArray::new(
-        BitBuffer::collect_bool_multiversioned(buffer.len(), |idx| {
-            // SAFETY: `collect_bool_multiversioned` invokes the predicate with indices
-            // `0..buffer.len()` only.
-            let value = unsafe { *buffer.get_unchecked(idx) };
-            lower.is_none_or(|l| lower_op(l, value)) & upper.is_none_or(|u| upper_op(value, u))
-        }),
+        BitBuffer::collect_bool_multiversioned_in(
+            buffer.len(),
+            |idx| {
+                // SAFETY: `collect_bool_multiversioned` invokes the predicate with indices
+                // `0..buffer.len()` only.
+                let value = unsafe { *buffer.get_unchecked(idx) };
+                lower.is_none_or(|l| lower_op(l, value)) & upper.is_none_or(|u| upper_op(value, u))
+            },
+            ctx.allocator().clone(),
+        ),
         arr.validity()
             .vortex_expect("validity should be derivable")
             .union_nullability(nullability),

vortex-array/src/arrays/filter/execute/buffer.rs

@@ -20,6 +20,7 @@
 use std::mem::size_of;
 
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_mask::MaskValues;
 
 use crate::arrays::filter::execute::byte_compress;
@@ -35,6 +36,7 @@ const MIN_SLICES_AVERAGE_RUN_LENGTH: usize = 8;
 /// Dense uniquely owned buffers are compacted in place; other buffers allocate a new output.
 pub(crate) fn filter_buffer<T: Copy>(buffer: Buffer<T>, mask: &MaskValues) -> Buffer<T> {
     assert_eq!(buffer.len(), mask.len());
+    let allocator = buffer.allocator().clone();
 
     let buffer = if mask.density() >= IN_PLACE_MIN_DENSITY {
         match buffer.try_into_mut() {
@@ -49,28 +51,32 @@ pub(crate) fn filter_buffer<T: Copy>(buffer: Buffer<T>, mask: &MaskValues) -> Bu
         buffer
     };
 
-    filter_slice(buffer.as_slice(), mask)
+    filter_slice(buffer.as_slice(), mask, allocator)
 }
 
-fn filter_slice<T: Copy>(values: &[T], mask: &MaskValues) -> Buffer<T> {
+fn filter_slice<T: Copy>(
+    values: &[T],
+    mask: &MaskValues,
+    allocator: BufferAllocatorRef,
+) -> Buffer<T> {
     if let Some(slices) = useful_cached_slices(mask) {
-        return slice::filter_slice_by_slices(values, slices, mask.true_count());
+        return slice::filter_slice_by_slices(values, slices, mask.true_count(), allocator);
     }
 
     if mask.density() <= CACHED_INDICES_MAX_DENSITY
         && let Some(indices) = mask.cached_indices()
     {
-        return slice::filter_slice_by_indices(values, indices);
+        return slice::filter_slice_by_indices(values, indices, allocator);
     }
 
-    if let Some(filtered) = simd_compress::filter_slice_by_bitmap(values, mask) {
+    if let Some(filtered) = simd_compress::filter_slice_by_bitmap(values, mask, allocator.clone()) {
         return filtered;
     }
 
     if mask.density() >= byte_compress_density_threshold::<T>() {
-        byte_compress::filter_buffer(values, mask)
+        byte_compress::filter_buffer(values, mask, allocator)
     } else {
-        slice::filter_slice_by_bitmap(values, mask)
+        slice::filter_slice_by_bitmap(values, mask, allocator)
     }
 }
 

vortex-array/src/arrays/filter/execute/byte_compress.rs

@@ -7,7 +7,9 @@
 //! permutation table compacts the selected bytes in a single indexed copy,
 //! avoiding the overhead of materializing indices or slices.
 
+use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_mask::MaskValues;
 
@@ -41,30 +43,35 @@ static BYTE_COMPRESS_LUT: &[([u8; 8], u8); 256] = &{
 ///
 /// Processes the mask one byte at a time (8 source elements per byte),
 /// using a precomputed permutation to compact selected elements.
-pub(crate) fn filter_buffer<T: Copy>(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer<T> {
+pub(crate) fn filter_buffer<T: Copy>(
+    buffer: impl AsRef<[T]>,
+    mask: &MaskValues,
+    allocator: BufferAllocatorRef,
+) -> Buffer<T> {
     let src = buffer.as_ref();
     debug_assert_eq!(src.len(), mask.len());
 
     let true_count = mask.true_count();
 
     if true_count == 0 {
-        return Buffer::empty();
+        return BufferMut::empty_aligned_in(Alignment::of::<T>(), allocator).freeze();
     }
 
     let mask_buffer = mask.bit_buffer();
     let mask_bytes = mask_buffer.inner().as_ref();
     let mask_offset = mask_buffer.offset();
 
-    filter_bitpacked(src, mask_bytes, mask_offset, true_count)
+    filter_bitpacked(src, mask_bytes, mask_offset, true_count, allocator)
 }
 
 fn filter_bitpacked<T: Copy>(
     src: &[T],
     mask_bytes: &[u8],
     mask_offset: usize,
     true_count: usize,
+    allocator: BufferAllocatorRef,
 ) -> Buffer<T> {
-    let mut out = BufferMut::<T>::with_capacity(true_count);
+    let mut out = BufferMut::<T>::with_capacity_in(true_count, allocator);
     let mut write_pos: usize = 0;
 
     if mask_offset == 0 {
@@ -165,6 +172,10 @@ mod tests {
 
     use super::*;
 
+    fn filter_buffer<T: Copy>(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer<T> {
+        super::filter_buffer(buffer, mask, BufferAllocatorRef::statically_allocated())
+    }
+
     fn mask_values(mask: &Mask) -> &MaskValues {
         match mask {
             Mask::Values(v) => v.as_ref(),

vortex-array/src/arrays/filter/execute/simd_compress/mod.rs

@@ -27,6 +27,7 @@
 use std::ptr;
 
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_mask::MaskValues;
 
@@ -50,12 +51,14 @@ type Kernel = unsafe fn(*const u8, *mut u8, &MaskValues) -> usize;
 pub(super) fn filter_slice_by_bitmap<T: Copy>(
     values: &[T],
     mask: &MaskValues,
+    allocator: BufferAllocatorRef,
 ) -> Option<Buffer<T>> {
     debug_assert_eq!(values.len(), mask.len());
     let kernel = select_kernel::<T, false>(mask)?;
 
     let true_count = mask.true_count();
-    let mut out = BufferMut::<T>::with_capacity(true_count + SLACK_BYTES / size_of::<T>());
+    let mut out =
+        BufferMut::<T>::with_capacity_in(true_count + SLACK_BYTES / size_of::<T>(), allocator);
     // SAFETY: `select_kernel` probed the kernel's target features; `values` holds `mask.len()`
     // elements and the output has capacity for every selected element plus a full vector of
     // slack, so each unmasked store stays in bounds.

vortex-array/src/arrays/filter/execute/simd_compress/tests.rs

@@ -11,6 +11,10 @@ use vortex_mask::MaskValues;
 use super::super::slice;
 use super::*;
 
+fn filter_slice_by_bitmap<T: Copy>(values: &[T], mask: &MaskValues) -> Option<Buffer<T>> {
+    super::filter_slice_by_bitmap(values, mask, BufferAllocatorRef::statically_allocated())
+}
+
 fn mask_values(mask: &Mask) -> Option<&MaskValues> {
     match mask {
         Mask::Values(values) => Some(values.as_ref()),
@@ -43,7 +47,8 @@ fn check<T: Copy + PartialEq + std::fmt::Debug>(values: &[T], mask: &Mask) {
     let Some(mask) = mask_values(mask) else {
         return;
     };
-    let expected = slice::filter_slice_by_bitmap(values, mask);
+    let expected =
+        slice::filter_slice_by_bitmap(values, mask, BufferAllocatorRef::statically_allocated());
 
     if let Some(actual) = filter_slice_by_bitmap(values, mask) {
         assert_eq!(actual.as_slice(), expected.as_slice());
@@ -114,7 +119,8 @@ fn avx2_kernels_match_scalar() {
         values: &[T],
         mask: &MaskValues,
     ) {
-        let expected = slice::filter_slice_by_bitmap(values, mask);
+        let expected =
+            slice::filter_slice_by_bitmap(values, mask, BufferAllocatorRef::statically_allocated());
 
         let mut out = vec![T::default(); mask.true_count() + SLACK_BYTES / size_of::<T>()];
         // SAFETY: AVX2 was detected above and the output has a vector of slack.

vortex-array/src/arrays/filter/execute/slice.rs

@@ -9,6 +9,7 @@
 use std::ptr;
 
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_mask::MaskValues;
 
@@ -53,15 +54,19 @@ pub(super) fn low_bits_mask(len: usize) -> u64 {
 }
 
 /// Filter a slice from the mask bitmap without materializing indices or ranges.
-pub(super) fn filter_slice_by_bitmap<T: Copy>(slice: &[T], mask: &MaskValues) -> Buffer<T> {
+pub(super) fn filter_slice_by_bitmap<T: Copy>(
+    slice: &[T],
+    mask: &MaskValues,
+    allocator: BufferAllocatorRef,
+) -> Buffer<T> {
     assert_eq!(
         mask.len(),
         slice.len(),
         "Selection mask length must equal the buffer length"
     );
 
     let output_len = mask.true_count();
-    let mut out = BufferMut::<T>::with_capacity(output_len);
+    let mut out = BufferMut::<T>::with_capacity_in(output_len, allocator);
     let src_ptr = slice.as_ptr();
     let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::<T>();
     let mut write_pos = 0;
@@ -98,8 +103,12 @@ pub(super) fn filter_slice_by_bitmap<T: Copy>(slice: &[T], mask: &MaskValues) ->
 }
 
 /// Filter a slice by a set of strictly increasing indices.
-pub(super) fn filter_slice_by_indices<T: Copy>(slice: &[T], indices: &[usize]) -> Buffer<T> {
-    let mut out = BufferMut::<T>::with_capacity(indices.len());
+pub(super) fn filter_slice_by_indices<T: Copy>(
+    slice: &[T],
+    indices: &[usize],
+    allocator: BufferAllocatorRef,
+) -> Buffer<T> {
+    let mut out = BufferMut::<T>::with_capacity_in(indices.len(), allocator);
     let src_ptr = slice.as_ptr();
     let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::<T>();
 
@@ -119,8 +128,9 @@ pub(super) fn filter_slice_by_slices<T: Copy>(
     slice: &[T],
     slices: &[(usize, usize)],
     output_len: usize,
+    allocator: BufferAllocatorRef,
 ) -> Buffer<T> {
-    let mut out = BufferMut::<T>::with_capacity(output_len);
+    let mut out = BufferMut::<T>::with_capacity_in(output_len, allocator);
     for (start, end) in slices {
         out.extend_from_slice(&slice[*start..*end]);
     }

vortex-array/src/arrays/filter/execute/take.rs

@@ -78,7 +78,8 @@ fn take_impl(
                 return array.child().filter(mask)?.cast(result_dtype);
             }
 
-            let translated = translate_indices(array.filter_mask(), indices, None)?;
+            let translated =
+                translate_indices(array.filter_mask(), indices, None, ctx.allocator().clone())?;
             let translated_indices =
                 PrimitiveArray::new(translated, indices.validity()?).into_array();
 
@@ -90,7 +91,12 @@ fn take_impl(
         )
         .into_array()),
         AllOr::Some(buf) => {
-            let translated = translate_indices(array.filter_mask(), indices, Some(buf))?;
+            let translated = translate_indices(
+                array.filter_mask(),
+                indices,
+                Some(buf),
+                ctx.allocator().clone(),
+            )?;
             let translated_indices =
                 PrimitiveArray::new(translated, indices.validity()?).into_array();