Allocate more kernel outputs through the session allocator

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

@@ -16,6 +16,7 @@ use vortex_array::arrays::filter::filter_chunked;
 use vortex_array::dtype::NativePType;
 use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_mask::MaskValues;
@@ -53,7 +54,7 @@ impl FilterKernel for BitPacked {
         let ptype = array.dtype().as_ptype();
         let validity = array.validity()?.filter(mask)?;
         let buffer = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |U| {
-            filter_values::<U>(array.data(), bit_width, values).into_byte_buffer()
+            filter_values::<U>(array.data(), bit_width, values, ctx.allocator()).into_byte_buffer()
         });
         let primitive = PrimitiveArray::from_byte_buffer(buffer, ptype, validity);
 
@@ -78,14 +79,15 @@ fn filter_values<T: NativePType + BitPacking>(
     array: &BitPackedData,
     bit_width: u8,
     mask: &MaskValues,
+    allocator: &BufferAllocatorRef,
 ) -> Buffer<T> {
     let bit_width = bit_width as usize;
     let chunks = PackedChunks {
         packed: array.packed_slice::<T>(),
         bit_width,
         packed_chunk_len: FL_CHUNK_SIZE * bit_width / (8 * size_of::<T>()),
     };
-    filter_chunked(&chunks, array.offset() as usize, mask)
+    filter_chunked(&chunks, array.offset() as usize, mask, allocator)
 }
 
 /// The packed FastLanes chunks of a bit-packed array.
@@ -151,6 +153,7 @@ mod tests {
     use vortex_array::validity::Validity;
     use vortex_buffer::BitBuffer;
     use vortex_buffer::Buffer;
+    use vortex_buffer::BufferAllocatorRef;
     use vortex_buffer::buffer;
     use vortex_error::VortexResult;
     use vortex_mask::Mask;
@@ -214,7 +217,13 @@ mod tests {
                 unreachable!("the indices select some but not all values")
             };
             assert_eq!(
-                filter_values::<T>(packed.data(), bit_width, &mask).as_slice(),
+                filter_values::<T>(
+                    packed.data(),
+                    bit_width,
+                    &mask,
+                    BufferAllocatorRef::static_ref()
+                )
+                .as_slice(),
                 expected
             );
         }
@@ -236,7 +245,8 @@ mod tests {
             unreachable!("the indices select some but not all values")
         };
         assert_eq!(
-            filter_values::<u32>(packed.data(), 0, &mask).as_slice(),
+            filter_values::<u32>(packed.data(), 0, &mask, BufferAllocatorRef::static_ref())
+                .as_slice(),
             &[0; 5]
         );
         Ok(())
@@ -510,7 +520,8 @@ mod tests {
             unreachable!("patterns select some but not all values")
         };
         assert_eq!(
-            filter_values::<T>(bitpacked.data(), 7, &mask).as_slice(),
+            filter_values::<T>(bitpacked.data(), 7, &mask, BufferAllocatorRef::static_ref())
+                .as_slice(),
             expected
         );
         Ok(())

encodings/runend/benches/run_end_filter.rs

@@ -21,6 +21,7 @@ use rand::SeedableRng;
 use rand::rngs::StdRng;
 use rand::seq::SliceRandom;
 use vortex_buffer::BitBuffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_runend::_benchmarking::filter_run_end_primitive;
 
 #[global_allocator]
@@ -112,6 +113,13 @@ fn filter_run_end(bencher: Bencher, args: FilterBenchArgs) {
     bencher
         .with_inputs(|| (run_ends.clone(), mask.clone()))
         .bench_refs(|(run_ends, mask)| {
-            filter_run_end_primitive::<u32>(run_ends, 0, length, mask).expect("filter")
+            filter_run_end_primitive::<u32>(
+                run_ends,
+                0,
+                length,
+                mask,
+                BufferAllocatorRef::static_ref(),
+            )
+            .expect("filter")
         });
 }

encodings/runend/src/compress.rs

@@ -28,6 +28,7 @@ use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::BitBufferMut;
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_buffer::buffer;
 use vortex_error::VortexExpect;
@@ -208,6 +209,7 @@ pub fn runend_decode_primitive(
                 validity_mask,
                 values.dtype().nullability(),
                 length,
+                ctx.allocator(),
             )
             .into_array()
         })
@@ -223,10 +225,11 @@ fn runend_decode_slice<T: Copy + Default>(
     values_validity: Mask,
     values_nullability: Nullability,
     length: usize,
+    allocator: &BufferAllocatorRef,
 ) -> (Buffer<T>, Validity) {
     match values_validity {
         Mask::AllTrue(_) => {
-            let mut decoded = RunWriter::new(length);
+            let mut decoded = RunWriter::new(length, allocator);
             for (end, &value) in run_ends.zip_eq(values) {
                 decoded.write_run(value, end);
             }
@@ -236,7 +239,7 @@ fn runend_decode_slice<T: Copy + Default>(
         Mask::Values(mask) => {
             // Nulls are usually rare and clustered, so start from all-valid and clear each
             // maximal streak of null runs once, rather than appending validity per run.
-            let mut decoded = RunWriter::new(length);
+            let mut decoded = RunWriter::new(length, allocator);
             let mut decoded_validity = BitBufferMut::full(true, length);
             let mut null_start = None;
             for (end, (&value, is_valid)) in
@@ -289,9 +292,9 @@ impl<T: Copy> RunWriter<T> {
         BLOCK_BYTES / size_of::<T>()
     };
 
-    fn new(length: usize) -> Self {
+    fn new(length: usize, allocator: &BufferAllocatorRef) -> Self {
         Self {
-            buffer: BufferMut::with_capacity(length + Self::BLOCK),
+            buffer: BufferMut::with_capacity_in(length + Self::BLOCK, allocator.clone()),
             len: 0,
             length,
         }
@@ -344,13 +347,15 @@ pub fn runend_decode_typed_primitive<T: NativePType>(
     values_validity: Mask,
     values_nullability: Nullability,
     length: usize,
+    allocator: &BufferAllocatorRef,
 ) -> PrimitiveArray {
     let (decoded, validity) = runend_decode_slice(
         run_ends,
         values,
         values_validity,
         values_nullability,
         length,
+        allocator,
     );
     PrimitiveArray::new(decoded, validity)
 }
@@ -381,6 +386,7 @@ pub fn runend_decode_decimal(
                 validity_mask,
                 values.dtype().nullability(),
                 length,
+                ctx.allocator(),
             )
         });
         DecimalArray::new(decoded, values.decimal_dtype(), validity).into_array()
@@ -413,6 +419,7 @@ pub fn runend_decode_varbinview(
             validity_mask,
             values.dtype().nullability(),
             length,
+            ctx.allocator(),
         )
     });
 

encodings/runend/src/compute/filter.rs

@@ -16,7 +16,8 @@ use vortex_array::dtype::NativePType;
 use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
-use vortex_buffer::buffer_mut;
+use vortex_buffer::BufferAllocatorRef;
+use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
@@ -72,6 +73,7 @@ impl FilterKernel for RunEnd {
                     array.offset() as u64,
                     array.len() as u64,
                     mask_values.bit_buffer(),
+                    ctx.allocator(),
                 )?
             });
         let filtered_values = array.values().filter(values_mask)?;
@@ -103,8 +105,9 @@ pub fn filter_run_end_primitive<R: NativePType + AddAssign + From<bool> + AsPrim
     offset: u64,
     length: u64,
     mask: &BitBuffer,
+    allocator: &BufferAllocatorRef,
 ) -> VortexResult<(PrimitiveArray, Mask)> {
-    let mut filtered_run_ends = buffer_mut![R::zero(); run_ends.len()];
+    let mut filtered_run_ends = BufferMut::<R>::zeroed_in(run_ends.len(), allocator.clone());
 
     let mut run_start = 0u64;
     let mut retained_run_count = 0;

encodings/sequence/src/compute/filter.rs

@@ -9,6 +9,7 @@ use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::filter::FilterKernel;
 use vortex_array::match_each_integer_ptype;
 use vortex_array::validity::Validity;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -22,12 +23,18 @@ impl FilterKernel for Sequence {
     fn filter(
         array: ArrayView<'_, Self>,
         mask: &Mask,
-        _ctx: &mut ExecutionCtx,
+        ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         let validity = Validity::from(array.dtype().nullability());
         match_each_integer_ptype!(array.dtype().as_ptype(), |O| {
             let (base, multiplier) = array.wrapping_parts::<O>()?;
-            Ok(Some(filter_impl(base, multiplier, mask, validity)))
+            Ok(Some(filter_impl(
+                base,
+                multiplier,
+                mask,
+                validity,
+                ctx.allocator(),
+            )))
         })
     }
 }
@@ -37,11 +44,12 @@ fn filter_impl<O: SequenceValue>(
     multiplier: O,
     mask: &Mask,
     validity: Validity,
+    allocator: &BufferAllocatorRef,
 ) -> ArrayRef {
     let mask_values = mask
         .values()
         .vortex_expect("FilterKernel precondition: mask is Mask::Values");
-    let mut buffer = BufferMut::<O>::with_capacity(mask_values.true_count());
+    let mut buffer = BufferMut::<O>::with_capacity_in(mask_values.true_count(), allocator.clone());
     buffer.extend(
         mask_values
             .indices()

vortex-array/src/arrays/constant/vtable/canonical.rs

@@ -10,7 +10,6 @@ use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_buffer::BufferString;
 use vortex_buffer::ByteBuffer;
-use vortex_buffer::buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_panic;
@@ -132,6 +131,7 @@ pub(crate) fn constant_canonicalize(
                 value,
                 array.dtype(),
                 array.len(),
+                ctx.allocator(),
             ))
         }
         DType::Binary(_) => {
@@ -140,6 +140,7 @@ pub(crate) fn constant_canonicalize(
                 value,
                 array.dtype(),
                 array.len(),
+                ctx.allocator(),
             ))
         }
         DType::List(..) => Canonical::List(constant_canonical_list_array(
@@ -225,10 +226,11 @@ fn constant_canonical_byte_view(
     scalar_bytes: Option<ByteBuffer>,
     dtype: &DType,
     len: usize,
+    allocator: &BufferAllocatorRef,
 ) -> VarBinViewArray {
     match scalar_bytes {
         None => {
-            let views = buffer![BinaryView::empty_view(); len];
+            let views = Buffer::full_in(BinaryView::empty_view(), len, allocator.clone());
 
             // SAFETY: for all-null the views and buffers are just zeroed, never accessed.
             unsafe {
@@ -251,7 +253,7 @@ fn constant_canonical_byte_view(
             }
 
             // Clone our constant view `len` times.
-            let views = buffer![view; len];
+            let views = Buffer::full_in(view, len, allocator.clone());
 
             // SAFETY: all the views are identical and point to a constant value.
             unsafe {

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

@@ -6,6 +6,7 @@
 use std::mem::MaybeUninit;
 
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
 use vortex_buffer::BufferMut;
 use vortex_mask::MaskValues;
 
@@ -92,8 +93,9 @@ pub fn filter_chunked<T: NativePType>(
     decoder: &impl ChunkDecoder<T>,
     offset: usize,
     mask: &MaskValues,
+    allocator: &BufferAllocatorRef,
 ) -> Buffer<T> {
-    let mut filter = ChunkedFilter::new(decoder, mask.true_count());
+    let mut filter = ChunkedFilter::new(decoder, mask.true_count(), allocator);
     if let Some(slices) = useful_cached_slices(mask) {
         filter.push_slices(offset, slices);
     } else if let Some(indices) = useful_cached_indices(mask) {
@@ -129,11 +131,14 @@ struct ChunkedFilter<'a, T, D> {
 }
 
 impl<'a, T: NativePType, D: ChunkDecoder<T>> ChunkedFilter<'a, T, D> {
-    fn new(decoder: &'a D, true_count: usize) -> Self {
+    fn new(decoder: &'a D, true_count: usize, allocator: &BufferAllocatorRef) -> Self {
         Self {
             decoder,
             // The SIMD kernels can store one vector past the last value that they write.
-            values: BufferMut::with_capacity(true_count + SLACK_BYTES / size_of::<T>()),
+            values: BufferMut::with_capacity_in(
+                true_count + SLACK_BYTES / size_of::<T>(),
+                allocator.clone(),
+            ),
             true_count,
             written: 0,
             scratch: Scratch([const { MaybeUninit::uninit() }; FILTER_CHUNK_LEN]),
@@ -423,6 +428,7 @@ mod tests {
     use rand::rngs::StdRng;
     use rstest::rstest;
     use vortex_buffer::BitBuffer;
+    use vortex_buffer::BufferAllocatorRef;
     use vortex_mask::Mask;
 
     use super::ChunkDecoder;
@@ -476,7 +482,12 @@ mod tests {
         let Mask::Values(mask) = mask else {
             unreachable!("the patterns select some but not all values")
         };
-        let filtered = filter_chunked(&SliceDecoder(&values), offset, &mask);
+        let filtered = filter_chunked(
+            &SliceDecoder(&values),
+            offset,
+            &mask,
+            BufferAllocatorRef::static_ref(),
+        );
         assert_eq!(filtered.as_slice(), expected);
     }