From 200cb9ff5774acec7fe2d3f4cc5545cb5a74cac2 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 19 Oct 2026 20:47:19 +0100 Subject: [PATCH] Allocate kernel outputs through the session allocator The filter of bit-packed, sequence and run-end arrays, run-end decoding and constant byte views allocated their output buffers from the global allocator. Allocate them through the allocator of the execution context instead, so a session can choose where decoded buffers come from. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5.5 --- .../src/bitpacking/compute/filter.rs | 21 ++++++++++++++----- encodings/runend/benches/run_end_filter.rs | 10 ++++++++- encodings/runend/src/compress.rs | 15 +++++++++---- encodings/runend/src/compute/filter.rs | 7 +++++-- encodings/sequence/src/compute/filter.rs | 14 ++++++++++--- .../src/arrays/constant/vtable/canonical.rs | 8 ++++--- .../src/arrays/filter/execute/chunked.rs | 19 +++++++++++++---- 7 files changed, 72 insertions(+), 22 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index b9aa2fdea79..9bee37f66f1 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/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::(array.data(), bit_width, values).into_byte_buffer() + filter_values::(array.data(), bit_width, values, ctx.allocator()).into_byte_buffer() }); let primitive = PrimitiveArray::from_byte_buffer(buffer, ptype, validity); @@ -78,6 +79,7 @@ fn filter_values( array: &BitPackedData, bit_width: u8, mask: &MaskValues, + allocator: &BufferAllocatorRef, ) -> Buffer { let bit_width = bit_width as usize; let chunks = PackedChunks { @@ -85,7 +87,7 @@ fn filter_values( bit_width, packed_chunk_len: FL_CHUNK_SIZE * bit_width / (8 * size_of::()), }; - 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::(packed.data(), bit_width, &mask).as_slice(), + filter_values::( + 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::(packed.data(), 0, &mask).as_slice(), + filter_values::(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::(bitpacked.data(), 7, &mask).as_slice(), + filter_values::(bitpacked.data(), 7, &mask, BufferAllocatorRef::static_ref()) + .as_slice(), expected ); Ok(()) diff --git a/encodings/runend/benches/run_end_filter.rs b/encodings/runend/benches/run_end_filter.rs index b012d43ca7d..2161bede77e 100644 --- a/encodings/runend/benches/run_end_filter.rs +++ b/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::(run_ends, 0, length, mask).expect("filter") + filter_run_end_primitive::( + run_ends, + 0, + length, + mask, + BufferAllocatorRef::static_ref(), + ) + .expect("filter") }); } diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index 6d4ebd98c33..a52d6e43a44 100644 --- a/encodings/runend/src/compress.rs +++ b/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( values_validity: Mask, values_nullability: Nullability, length: usize, + allocator: &BufferAllocatorRef, ) -> (Buffer, 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( 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 RunWriter { BLOCK_BYTES / size_of::() }; - 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,6 +347,7 @@ pub fn runend_decode_typed_primitive( values_validity: Mask, values_nullability: Nullability, length: usize, + allocator: &BufferAllocatorRef, ) -> PrimitiveArray { let (decoded, validity) = runend_decode_slice( run_ends, @@ -351,6 +355,7 @@ pub fn runend_decode_typed_primitive( 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(), ) }); diff --git a/encodings/runend/src/compute/filter.rs b/encodings/runend/src/compute/filter.rs index d4b4df3b2b0..b85465aa4f7 100644 --- a/encodings/runend/src/compute/filter.rs +++ b/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 + 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::::zeroed_in(run_ends.len(), allocator.clone()); let mut run_start = 0u64; let mut retained_run_count = 0; diff --git a/encodings/sequence/src/compute/filter.rs b/encodings/sequence/src/compute/filter.rs index 52da57654f5..dcdc63c454d 100644 --- a/encodings/sequence/src/compute/filter.rs +++ b/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> { let validity = Validity::from(array.dtype().nullability()); match_each_integer_ptype!(array.dtype().as_ptype(), |O| { let (base, multiplier) = array.wrapping_parts::()?; - 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( 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::::with_capacity(mask_values.true_count()); + let mut buffer = BufferMut::::with_capacity_in(mask_values.true_count(), allocator.clone()); buffer.extend( mask_values .indices() diff --git a/vortex-array/src/arrays/constant/vtable/canonical.rs b/vortex-array/src/arrays/constant/vtable/canonical.rs index 170ad88167b..1a4244a07c6 100644 --- a/vortex-array/src/arrays/constant/vtable/canonical.rs +++ b/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, 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 { diff --git a/vortex-array/src/arrays/filter/execute/chunked.rs b/vortex-array/src/arrays/filter/execute/chunked.rs index d6c2647a6de..d8f0edf02dd 100644 --- a/vortex-array/src/arrays/filter/execute/chunked.rs +++ b/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( decoder: &impl ChunkDecoder, offset: usize, mask: &MaskValues, + allocator: &BufferAllocatorRef, ) -> Buffer { - 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> 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::()), + values: BufferMut::with_capacity_in( + true_count + SLACK_BYTES / size_of::(), + 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); }