From e6a89230d8a8f4cdee8a89c9e93d3a40d90a22bc Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Thu, 27 Aug 2026 13:08:17 -0400 Subject: [PATCH 1/2] refactor(array): allocate execution outputs through context Signed-off-by: Nicholas Gates --- .../src/arrays/decimal/compute/between.rs | 23 ++- .../src/arrays/filter/execute/buffer.rs | 20 ++- .../arrays/filter/execute/byte_compress.rs | 19 ++- .../filter/execute/simd_compress/mod.rs | 5 +- .../filter/execute/simd_compress/tests.rs | 10 +- .../src/arrays/filter/execute/slice.rs | 20 ++- .../src/arrays/filter/execute/take.rs | 10 +- .../arrays/filter/execute/take/fixed_width.rs | 82 +++++++--- .../src/arrays/filter/execute/take/rank.rs | 30 ++-- vortex-array/src/arrays/fixed_width/filter.rs | 6 +- .../src/arrays/fixed_width/take/avx2/mod.rs | 23 ++- .../src/arrays/fixed_width/take/mod.rs | 9 +- .../src/arrays/fixed_width/take/records.rs | 6 +- .../src/arrays/fixed_width/take/scalar.rs | 4 +- .../src/arrays/fixed_width/take/slices.rs | 17 ++- .../src/arrays/fixed_width/take/tests.rs | 25 +++- .../src/arrays/interleave/execute/bool.rs | 20 ++- .../arrays/interleave/execute/primitive.rs | 47 +++--- .../src/arrays/primitive/compute/between.rs | 23 ++- vortex-array/src/builders/dict/bytes.rs | 55 +++++-- vortex-array/src/builders/dict/mod.rs | 22 ++- vortex-array/src/builders/dict/primitive.rs | 35 ++++- vortex-array/src/patches.rs | 32 ++-- .../src/scalar_fn/fns/binary/boolean.rs | 53 +++++-- .../scalar_fn/fns/binary/compare/boolean.rs | 22 ++- .../src/scalar_fn/fns/binary/compare/bytes.rs | 141 ++++++++++++------ .../scalar_fn/fns/binary/compare/decimal.rs | 53 ++++--- .../src/scalar_fn/fns/binary/compare/mod.rs | 12 +- .../scalar_fn/fns/binary/compare/nested.rs | 28 +++- .../scalar_fn/fns/binary/compare/primitive.rs | 47 +++--- .../src/scalar_fn/fns/binary/compare/tests.rs | 21 +++ .../scalar_fn/fns/binary/numeric/checked.rs | 6 +- .../scalar_fn/fns/binary/numeric/decimal.rs | 48 +++--- vortex-array/src/scalar_fn/fns/like/mod.rs | 126 ++++++++++------ .../src/scalar_fn/fns/list_contains/mod.rs | 26 ++-- vortex-buffer/src/bit/buf.rs | 50 ++++++- vortex-buffer/src/bit/buf_mut.rs | 34 ++++- 37 files changed, 870 insertions(+), 340 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/between.rs b/vortex-array/src/arrays/decimal/compute/between.rs index 50623b077ee..91c9cadd6c4 100644 --- a/vortex-array/src/arrays/decimal/compute/between.rs +++ b/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> { // 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::(arr, lower, upper, nullability, options) + between_unpack::(arr, lower, upper, nullability, options, ctx) }) } } @@ -52,6 +52,7 @@ fn between_unpack( upper: Scalar, nullability: Nullability, options: &BetweenOptions, + ctx: &mut ExecutionCtx, ) -> VortexResult> { 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( nullability, lower_op, upper_op, + ctx, ))) } @@ -129,15 +131,20 @@ fn between_impl( nullability: Nullability, lower_op: impl Fn(T, T) -> bool, upper_op: impl Fn(T, T) -> bool, + ctx: &mut ExecutionCtx, ) -> ArrayRef { let buffer = arr.buffer::(); 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), diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index b8c73863f11..323521fda69 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/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(buffer: Buffer, mask: &MaskValues) -> Buffer { 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(buffer: Buffer, mask: &MaskValues) -> Bu buffer }; - filter_slice(buffer.as_slice(), mask) + filter_slice(buffer.as_slice(), mask, allocator) } -fn filter_slice(values: &[T], mask: &MaskValues) -> Buffer { +fn filter_slice( + values: &[T], + mask: &MaskValues, + allocator: BufferAllocatorRef, +) -> Buffer { 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::() { - 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) } } diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs index 032b11a56b3..909de49bf66 100644 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ b/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,21 +43,25 @@ 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(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer { +pub(crate) fn filter_buffer( + buffer: impl AsRef<[T]>, + mask: &MaskValues, + allocator: BufferAllocatorRef, +) -> Buffer { 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::(), 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( @@ -63,8 +69,9 @@ fn filter_bitpacked( mask_bytes: &[u8], mask_offset: usize, true_count: usize, + allocator: BufferAllocatorRef, ) -> Buffer { - let mut out = BufferMut::::with_capacity(true_count); + let mut out = BufferMut::::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(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer { + super::filter_buffer(buffer, mask, BufferAllocatorRef::statically_allocated()) + } + fn mask_values(mask: &Mask) -> &MaskValues { match mask { Mask::Values(v) => v.as_ref(), diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs b/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs index bb9ebea63a7..3e1978a29e6 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/mod.rs +++ b/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( values: &[T], mask: &MaskValues, + allocator: BufferAllocatorRef, ) -> Option> { debug_assert_eq!(values.len(), mask.len()); let kernel = select_kernel::(mask)?; let true_count = mask.true_count(); - let mut out = BufferMut::::with_capacity(true_count + SLACK_BYTES / size_of::()); + let mut out = + BufferMut::::with_capacity_in(true_count + SLACK_BYTES / size_of::(), 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. diff --git a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs b/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs index e79b70a2f94..1b2ea904f8d 100644 --- a/vortex-array/src/arrays/filter/execute/simd_compress/tests.rs +++ b/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(values: &[T], mask: &MaskValues) -> Option> { + 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(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::()]; // SAFETY: AVX2 was detected above and the output has a vector of slack. diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index ea428564dcb..fca32616acf 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/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,7 +54,11 @@ 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(slice: &[T], mask: &MaskValues) -> Buffer { +pub(super) fn filter_slice_by_bitmap( + slice: &[T], + mask: &MaskValues, + allocator: BufferAllocatorRef, +) -> Buffer { assert_eq!( mask.len(), slice.len(), @@ -61,7 +66,7 @@ pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> ); let output_len = mask.true_count(); - let mut out = BufferMut::::with_capacity(output_len); + let mut out = BufferMut::::with_capacity_in(output_len, allocator); let src_ptr = slice.as_ptr(); let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); let mut write_pos = 0; @@ -98,8 +103,12 @@ pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> } /// Filter a slice by a set of strictly increasing indices. -pub(super) fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { - let mut out = BufferMut::::with_capacity(indices.len()); +pub(super) fn filter_slice_by_indices( + slice: &[T], + indices: &[usize], + allocator: BufferAllocatorRef, +) -> Buffer { + let mut out = BufferMut::::with_capacity_in(indices.len(), allocator); let src_ptr = slice.as_ptr(); let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); @@ -119,8 +128,9 @@ pub(super) fn filter_slice_by_slices( slice: &[T], slices: &[(usize, usize)], output_len: usize, + allocator: BufferAllocatorRef, ) -> Buffer { - let mut out = BufferMut::::with_capacity(output_len); + let mut out = BufferMut::::with_capacity_in(output_len, allocator); for (start, end) in slices { out.extend_from_slice(&slice[*start..*end]); } diff --git a/vortex-array/src/arrays/filter/execute/take.rs b/vortex-array/src/arrays/filter/execute/take.rs index b8d8d2da1b6..e44dbe06d76 100644 --- a/vortex-array/src/arrays/filter/execute/take.rs +++ b/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(); diff --git a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs index 6518293db41..e926f501500 100644 --- a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs +++ b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs @@ -3,6 +3,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::AllOr; @@ -95,6 +96,7 @@ where let ranks = indices.as_slice::

(); let ranks_validity = indices.validity()?; let indices_validity = ranks_validity.execute_mask(indices.len(), ctx)?; + let allocator = ctx.allocator().clone(); match indices_validity.bit_buffer() { AllOr::All => { @@ -102,30 +104,43 @@ where { child.slice(start..end) } else { - take_filtered_values::(&child, filter, ranks, None)? + take_filtered_values::(&child, filter, ranks, None, allocator.clone())? }; let output_validity = if child_validity.definitely_no_nulls() { ranks_validity.union_nullability(child_validity.nullability()) } else { - let translated_indices = - PrimitiveArray::new(translate_ranks(filter, ranks, None)?, ranks_validity) - .into_array(); + let translated_indices = PrimitiveArray::new( + translate_ranks(filter, ranks, None, allocator)?, + ranks_validity, + ) + .into_array(); child_validity.take(&translated_indices)? }; Ok((taken, output_validity)) } - AllOr::None => Ok((Buffer::zeroed(ranks.len()), Validity::AllInvalid)), + AllOr::None => Ok(( + Buffer::zeroed_in(ranks.len(), allocator), + Validity::AllInvalid, + )), AllOr::Some(buf) => { - let taken = take_filtered_values(child.as_slice(), filter, ranks, Some(buf))?; + let taken = take_filtered_values( + child.as_slice(), + filter, + ranks, + Some(buf), + allocator.clone(), + )?; let output_validity = if child_validity.definitely_no_nulls() { ranks_validity.union_nullability(child_validity.nullability()) } else { - let translated_indices = - PrimitiveArray::new(translate_ranks(filter, ranks, Some(buf))?, ranks_validity) - .into_array(); + let translated_indices = PrimitiveArray::new( + translate_ranks(filter, ranks, Some(buf), allocator)?, + ranks_validity, + ) + .into_array(); child_validity.take(&translated_indices)? }; @@ -139,6 +154,7 @@ fn take_filtered_values( filter: &Mask, ranks: &[P], indices_validity: Option<&BitBuffer>, + allocator: BufferAllocatorRef, ) -> VortexResult> where T: Copy + Default, @@ -147,32 +163,49 @@ where let filtered_len = filter.true_count(); if let Some(start) = contiguous_filter_start(filter) { return if let Some(indices_validity) = indices_validity { - take_values_by_rank_nullable(values, ranks, indices_validity, filtered_len, |idx| { - start + idx - }) + take_values_by_rank_nullable( + values, + ranks, + indices_validity, + filtered_len, + allocator, + |idx| start + idx, + ) } else { - take_values_by_rank(values, ranks, filtered_len, |idx| start + idx) + take_values_by_rank(values, ranks, filtered_len, allocator, |idx| start + idx) }; } if ranks.len() <= small_take_rank_lookup_len(filter) { return if let Some(indices_validity) = indices_validity { - take_values_by_rank_nullable(values, ranks, indices_validity, filtered_len, |idx| { + take_values_by_rank_nullable( + values, + ranks, + indices_validity, + filtered_len, + allocator, + |idx| filter.rank(idx), + ) + } else { + take_values_by_rank(values, ranks, filtered_len, allocator, |idx| { filter.rank(idx) }) - } else { - take_values_by_rank(values, ranks, filtered_len, |idx| filter.rank(idx)) }; } match filter.indices() { AllOr::All => { if let Some(indices_validity) = indices_validity { - take_values_by_rank_nullable(values, ranks, indices_validity, filtered_len, |idx| { - idx - }) + take_values_by_rank_nullable( + values, + ranks, + indices_validity, + filtered_len, + allocator, + |idx| idx, + ) } else { - take_values_by_rank(values, ranks, filtered_len, |idx| idx) + take_values_by_rank(values, ranks, filtered_len, allocator, |idx| idx) } } AllOr::None => unreachable!("empty filters are handled by the filter short circuit"), @@ -183,10 +216,11 @@ where ranks, indices_validity, filtered_len, + allocator, |idx| unsafe { *indices.get_unchecked(idx) }, ) } else { - take_values_by_rank(values, ranks, filtered_len, |idx| unsafe { + take_values_by_rank(values, ranks, filtered_len, allocator, |idx| unsafe { *indices.get_unchecked(idx) }) } @@ -199,6 +233,7 @@ fn take_values_by_rank_nullable( ranks: &[P], ranks_validity: &BitBuffer, translated_len: usize, + allocator: BufferAllocatorRef, translate: L, ) -> VortexResult> where @@ -206,7 +241,7 @@ where P: IntegerPType, L: Fn(usize) -> usize, { - let mut out = BufferMut::::with_capacity(ranks.len()); + let mut out = BufferMut::::with_capacity_in(ranks.len(), allocator); let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); for (idx, rank) in ranks.iter().enumerate() { let value = if ranks_validity.value(idx) { @@ -232,6 +267,7 @@ fn take_values_by_rank( values: &[T], ranks: &[P], translated_len: usize, + allocator: BufferAllocatorRef, translate: L, ) -> VortexResult> where @@ -239,7 +275,7 @@ where P: IntegerPType, L: Fn(usize) -> usize, { - let mut out = BufferMut::::with_capacity(ranks.len()); + let mut out = BufferMut::::with_capacity_in(ranks.len(), allocator); let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); for (idx, rank) in ranks.iter().enumerate() { let rank = validate_rank(*rank, translated_len)?; diff --git a/vortex-array/src/arrays/filter/execute/take/rank.rs b/vortex-array/src/arrays/filter/execute/take/rank.rs index 08170ab85d3..f36f3086aec 100644 --- a/vortex-array/src/arrays/filter/execute/take/rank.rs +++ b/vortex-array/src/arrays/filter/execute/take/rank.rs @@ -3,6 +3,7 @@ use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -20,9 +21,10 @@ pub(in crate::arrays::filter) fn translate_indices( filter: &Mask, indices: &PrimitiveArray, indices_validity: Option<&BitBuffer>, + allocator: BufferAllocatorRef, ) -> VortexResult> { match_each_integer_ptype!(indices.ptype(), |P| { - translate_ranks(filter, indices.as_slice::

(), indices_validity) + translate_ranks(filter, indices.as_slice::

(), indices_validity, allocator) }) } @@ -68,27 +70,34 @@ pub(in crate::arrays::filter) fn translate_ranks( filter: &Mask, ranks: &[P], ranks_validity: Option<&BitBuffer>, + allocator: BufferAllocatorRef, ) -> VortexResult> { let filtered_len = filter.true_count(); if let Some(start) = contiguous_filter_start(filter) { - return translate_ranks_with(ranks, ranks_validity, filtered_len, |rank| start + rank); + return translate_ranks_with(ranks, ranks_validity, filtered_len, allocator, |rank| { + start + rank + }); } if ranks.len() <= small_take_rank_lookup_len(filter) { - return translate_ranks_with(ranks, ranks_validity, filtered_len, |rank| { + return translate_ranks_with(ranks, ranks_validity, filtered_len, allocator, |rank| { filter.rank(rank) }); } match filter.indices() { - AllOr::All => translate_ranks_with(ranks, ranks_validity, filtered_len, |rank| rank), - AllOr::None => unreachable!("empty filters are handled by the filter short circuit"), - AllOr::Some(filter_indices) => { - translate_ranks_with(ranks, ranks_validity, filtered_len, |rank| unsafe { - *filter_indices.get_unchecked(rank) - }) + AllOr::All => { + translate_ranks_with(ranks, ranks_validity, filtered_len, allocator, |rank| rank) } + AllOr::None => unreachable!("empty filters are handled by the filter short circuit"), + AllOr::Some(filter_indices) => translate_ranks_with( + ranks, + ranks_validity, + filtered_len, + allocator, + |rank| unsafe { *filter_indices.get_unchecked(rank) }, + ), } } @@ -96,13 +105,14 @@ fn translate_ranks_with( ranks: &[P], ranks_validity: Option<&BitBuffer>, filtered_len: usize, + allocator: BufferAllocatorRef, translate: L, ) -> VortexResult> where P: IntegerPType, L: Fn(usize) -> usize, { - let mut translated = BufferMut::::with_capacity(ranks.len()); + let mut translated = BufferMut::::with_capacity_in(ranks.len(), allocator); let translated_ptr = translated.spare_capacity_mut().as_mut_ptr().cast::(); for (idx, rank) in ranks.iter().enumerate() { diff --git a/vortex-array/src/arrays/fixed_width/filter.rs b/vortex-array/src/arrays/fixed_width/filter.rs index 7d7eec9d635..0ce702b47de 100644 --- a/vortex-array/src/arrays/fixed_width/filter.rs +++ b/vortex-array/src/arrays/fixed_width/filter.rs @@ -34,6 +34,7 @@ pub(crate) fn filter(array: &Array, mask: &MaskValuesRef) fn filter_records(values: ByteBuffer, byte_width: usize, mask: &MaskValues) -> ByteBuffer { let alignment = values.alignment(); + let allocator = values.allocator().clone(); match_each_record_width!( byte_width, @@ -57,7 +58,10 @@ fn filter_records(values: ByteBuffer, byte_width: usize, mask: &MaskValues) -> B values.freeze().into_byte_buffer().aligned(alignment) } Err(values) => { - let mut filtered = BufferMut::with_capacity(mask.true_count() * byte_width); + let mut filtered = BufferMut::with_capacity_in( + mask.true_count() * byte_width, + allocator, + ); mask.bit_buffer().for_each_set_index(|index| { let start = index * byte_width; filtered.extend_from_slice(&values[start..start + byte_width]); diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs index 6a0ba762552..c77bb2daecd 100644 --- a/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs @@ -19,6 +19,7 @@ use std::arch::x86_64::_mm256_set1_epi32; use vortex_buffer::Alignment; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use self::gather::Avx2Gather; @@ -46,13 +47,14 @@ use crate::match_each_unsigned_integer_ptype; pub(super) unsafe fn take_avx2( buffer: &[V], indices: &[I], + allocator: BufferAllocatorRef, ) -> Buffer { if buffer.is_empty() { assert!( indices.is_empty(), "cannot take a non-empty set of indices from an empty buffer" ); - return Buffer::empty(); + return BufferMut::empty_aligned_in(Alignment::of::(), allocator).freeze(); } // Dispatch on the gather lane width. The index type must still be concretized to select the @@ -63,7 +65,7 @@ pub(super) unsafe fn take_avx2( // SAFETY: `Idx` has the same `PTYPE` as `I`, so this is a no-op reinterpret of the // index slice into the concrete type the gather impl is keyed on. let indices = unsafe { std::mem::transmute::<&[I], &[Idx]>(indices) }; - exec_take::(buffer, indices) + exec_take::(buffer, indices, allocator.clone()) }) }}; } @@ -72,12 +74,12 @@ pub(super) unsafe fn take_avx2( // The i32 gather interprets u32 lanes as signed offsets. A valid high u32 index needs the // scalar path when the values slice exceeds the non-negative i32 addressable range. 4 if I::PTYPE == PType::U32 && !i32_gather_can_address(buffer.len()) => { - take_values_scalar(buffer, indices) + take_values_scalar(buffer, indices, allocator) } 4 => dispatch!(u32), 8 => dispatch!(u64), // 1/2-byte and >8-byte values have no AVX2 gather lane, so fall back to scalar. - _ => take_values_scalar(buffer, indices), + _ => take_values_scalar(buffer, indices, allocator), } } @@ -95,7 +97,11 @@ const fn i32_gather_can_address(values_len: usize) -> bool { /// `Out`; invalid lanes are masked and cause a panic before the output buffer is initialized. /// Gather instructions tolerate the source's potentially weaker alignment. #[inline(always)] -fn exec_take(values: &[Out], indices: &[Idx]) -> Buffer +fn exec_take( + values: &[Out], + indices: &[Idx], + allocator: BufferAllocatorRef, +) -> Buffer where Out: FixedWidthTakeValue, Idx: UnsignedPType, @@ -111,8 +117,11 @@ where // The length is an exclusive upper bound on valid indices. `None` means the bound does not // fit in the index type, so every representable index is in-bounds. let max_index = Idx::from(values.len()); - let mut buffer = - BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); + let mut buffer = BufferMut::::with_capacity_aligned_in( + indices_len, + Alignment::of::<__m256i>(), + allocator, + ); let buf_uninit = buffer.spare_capacity_mut(); let mut offset = 0; diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs index d6e2f744bdd..2999fdd2023 100644 --- a/vortex-array/src/arrays/fixed_width/take/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -13,6 +13,7 @@ mod tests; use std::sync::LazyLock; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; @@ -79,16 +80,17 @@ unsafe impl FixedWidthTakeValue for [u8; N] {} pub(crate) fn take_values( values: &[T], indices: &[I], + allocator: BufferAllocatorRef, ) -> Buffer { #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] if *HAS_AVX2 { // SAFETY: AVX2 was detected above and `FixedWidthTakeValue` guarantees an initialized byte // representation. The AVX2 dispatcher retains Primitive's existing scalar fallbacks and // out-of-bounds behavior for every value width. - return unsafe { avx2::take_avx2(values, indices) }; + return unsafe { avx2::take_avx2(values, indices, allocator) }; } - take_values_scalar(values, indices) + take_values_scalar(values, indices, allocator) } pub(crate) fn take( @@ -142,6 +144,7 @@ pub(crate) fn take( V::byte_width(array), array.len(), indices.as_slice::(), + ctx.allocator().clone(), ) })?; Ok(Some( @@ -173,6 +176,7 @@ fn take_contiguous_ranges( starts.as_slice::(), length, output_len, + ctx.allocator().clone(), ) }) } @@ -187,6 +191,7 @@ fn take_contiguous_ranges( starts.as_slice::(), lengths.as_slice::(), output_len, + ctx.allocator().clone(), ) }) }) diff --git a/vortex-array/src/arrays/fixed_width/take/records.rs b/vortex-array/src/arrays/fixed_width/take/records.rs index 6cb47957991..d3a8cf43387 100644 --- a/vortex-array/src/arrays/fixed_width/take/records.rs +++ b/vortex-array/src/arrays/fixed_width/take/records.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -16,6 +17,7 @@ pub(super) fn take_byte_records( byte_width: usize, record_count: usize, indices: &[I], + allocator: BufferAllocatorRef, ) -> VortexResult { let alignment = values.alignment(); @@ -24,7 +26,7 @@ pub(super) fn take_byte_records( |W| { let records = Buffer::<[u8; W]>::from_byte_buffer(values.clone()); debug_assert_eq!(records.len(), record_count); - Ok(take_values(records.as_slice(), indices) + Ok(take_values(records.as_slice(), indices, allocator) .into_byte_buffer() .aligned(alignment)) }, @@ -33,7 +35,7 @@ pub(super) fn take_byte_records( .len() .checked_mul(byte_width) .ok_or_else(|| vortex_err!("Fixed-width take output length overflows usize"))?; - let mut result = BufferMut::::with_capacity(output_len); + let mut result = BufferMut::::with_capacity_in(output_len, allocator); for index in indices { let index = index.as_(); assert!( diff --git a/vortex-array/src/arrays/fixed_width/take/scalar.rs b/vortex-array/src/arrays/fixed_width/take/scalar.rs index 7a1d1865be3..341e58830f7 100644 --- a/vortex-array/src/arrays/fixed_width/take/scalar.rs +++ b/vortex-array/src/arrays/fixed_width/take/scalar.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use crate::dtype::IntegerPType; @@ -10,10 +11,11 @@ use crate::dtype::IntegerPType; pub(crate) fn take_values_scalar( values: &[T], indices: &[I], + allocator: BufferAllocatorRef, ) -> Buffer { // The explicit pointer loop keeps the source length in a register and avoids a capacity check // for every output value. - let mut result = BufferMut::with_capacity(indices.len()); + let mut result = BufferMut::with_capacity_in(indices.len(), allocator); let result_ptr = result.spare_capacity_mut().as_mut_ptr().cast::(); for (output_index, index) in indices.iter().enumerate() { diff --git a/vortex-array/src/arrays/fixed_width/take/slices.rs b/vortex-array/src/arrays/fixed_width/take/slices.rs index a4d9ceda84a..248a29a3762 100644 --- a/vortex-array/src/arrays/fixed_width/take/slices.rs +++ b/vortex-array/src/arrays/fixed_width/take/slices.rs @@ -4,6 +4,7 @@ use std::ptr; use itertools::Itertools as _; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -19,12 +20,20 @@ pub(super) fn take_slices( starts: &[S], lengths: &[L], output_len: usize, + allocator: BufferAllocatorRef, ) -> VortexResult { let slices = starts .iter() .zip_eq(lengths) .map(|(&start, &length)| (start.as_(), length.as_())); - copy_slices(values, byte_width, record_count, slices, output_len) + copy_slices( + values, + byte_width, + record_count, + slices, + output_len, + allocator, + ) } pub(super) fn take_slices_constant_length( @@ -34,6 +43,7 @@ pub(super) fn take_slices_constant_length( starts: &[S], length: usize, output_len: usize, + allocator: BufferAllocatorRef, ) -> VortexResult { let computed_len = starts .len() @@ -49,6 +59,7 @@ pub(super) fn take_slices_constant_length( record_count, starts.iter().map(|start| (start.as_(), length)), output_len, + allocator, ) } @@ -58,6 +69,7 @@ fn copy_slices( record_count: usize, slices: impl IntoIterator, output_len: usize, + allocator: BufferAllocatorRef, ) -> VortexResult { let input_byte_len = record_count .checked_mul(byte_width) @@ -70,7 +82,8 @@ fn copy_slices( let output_byte_len = output_len .checked_mul(byte_width) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - let mut result = BufferMut::::with_capacity_aligned(output_byte_len, values.alignment()); + let mut result = + BufferMut::::with_capacity_aligned_in(output_byte_len, values.alignment(), allocator); let spare = &mut result.spare_capacity_mut()[..output_byte_len]; let mut cursor = 0usize; diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs index ff167d27499..1d24ab7a45d 100644 --- a/vortex-array/src/arrays/fixed_width/take/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -3,6 +3,7 @@ use rstest::rstest; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -25,16 +26,20 @@ use crate::dtype::DecimalDType; use crate::dtype::i256; use crate::validity::Validity; +fn allocator() -> BufferAllocatorRef { + BufferAllocatorRef::statically_allocated() +} + #[test] fn take_four_byte_records() { let values = [[1u8, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]; - let taken = take_values(&values, &[2u32, 0]); + let taken = take_values(&values, &[2u32, 0], allocator()); assert_eq!(taken.as_slice(), &[[9, 10, 11, 12], [1, 2, 3, 4]]); } #[test] fn take_eight_byte_values() { - let taken = take_values(&[10i64, 20, 30], &[1u16, 2, 0]); + let taken = take_values(&[10i64, 20, 30], &[1u16, 2, 0], allocator()); assert_eq!(taken.as_slice(), &[20, 30, 10]); } @@ -54,7 +59,13 @@ fn take_runtime_width_records(#[case] byte_width: usize) -> VortexResult<()> { .chain(&values[..byte_width]) .copied() .collect::>(); - let taken = take_byte_records(&values.into_byte_buffer(), byte_width, 3, &[2u32, 0])?; + let taken = take_byte_records( + &values.into_byte_buffer(), + byte_width, + 3, + &[2u32, 0], + allocator(), + )?; assert_eq!(taken.as_slice(), expected); Ok(()) } @@ -63,13 +74,13 @@ fn take_runtime_width_records(#[case] byte_width: usize) -> VortexResult<()> { #[should_panic(expected = "take index 3 out of bounds for length 3")] fn fallback_take_rejects_out_of_bounds_index() { let values = Buffer::from_iter((0u8..).take(9)).into_byte_buffer(); - drop(take_byte_records(&values, 3, 3, &[3u32])); + drop(take_byte_records(&values, 3, 3, &[3u32], allocator())); } #[test] fn take_variable_length_slices() -> VortexResult<()> { let values = buffer![10u8, 11, 12, 13, 14].into_byte_buffer(); - let taken = take_slices(&values, 1, 5, &[1u32, 3], &[2u32, 1], 3)?; + let taken = take_slices(&values, 1, 5, &[1u32, 3], &[2u32, 1], 3, allocator())?; assert_eq!(taken.as_slice(), &[11, 12, 13]); Ok(()) } @@ -77,13 +88,13 @@ fn take_variable_length_slices() -> VortexResult<()> { #[test] fn variable_length_slices_validate_output_length() { let values = buffer![10u8, 11, 12, 13].into_byte_buffer(); - assert!(take_slices(&values, 1, 4, &[0u32, 2], &[1u32, 1], 3).is_err()); + assert!(take_slices(&values, 1, 4, &[0u32, 2], &[1u32, 1], 3, allocator(),).is_err()); } #[test] fn take_constant_length_slices() -> VortexResult<()> { let values = buffer![10u8, 11, 12, 13, 14].into_byte_buffer(); - let taken = take_slices_constant_length(&values, 1, 5, &[0u32, 3], 2, 4)?; + let taken = take_slices_constant_length(&values, 1, 5, &[0u32, 3], 2, 4, allocator())?; assert_eq!(taken.as_slice(), &[10, 11, 13, 14]); Ok(()) } diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index fde5b161dfd..372818d189f 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -6,6 +6,7 @@ use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -25,7 +26,7 @@ use crate::require_child; /// each selected bit into the output position it routes to. pub(super) fn execute( array: Array, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let num_values = array.num_values(); @@ -54,6 +55,7 @@ pub(super) fn execute( &value_bits, array_indices.as_slice::(), row_indices.as_slice::(), + ctx.allocator().clone(), )? }) }); @@ -70,12 +72,13 @@ fn gather, R: AsPrimitive>( value_bits: &[BitBuffer], branches: &[A], rows: &[R], + allocator: BufferAllocatorRef, ) -> VortexResult { let len = validate_selectors(value_bits, branches, rows)?; // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every // `i < len` that `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()`. - Ok(unsafe { gather_bits(len, value_bits, branches, rows) }) + Ok(unsafe { gather_bits(len, value_bits, branches, rows, allocator) }) } /// Validates the per-row selector bounds, returning the output length (`branches.len()`). @@ -125,11 +128,16 @@ unsafe fn gather_bits, R: AsPrimitive>( bits: &[BitBuffer], branches: &[A], rows: &[R], + allocator: BufferAllocatorRef, ) -> BitBufferMut { // SAFETY: `collect_bool` calls this for `i < len`, and the caller guarantees `branches[i]` and // `rows[i]` are in bounds for `bits` / the selected buffer. - BitBufferMut::collect_bool(len, |i| unsafe { - bits.get_unchecked(branches.get_unchecked(i).as_()) - .value_unchecked(rows.get_unchecked(i).as_()) - }) + BitBufferMut::collect_bool_in( + len, + |i| unsafe { + bits.get_unchecked(branches.get_unchecked(i).as_()) + .value_unchecked(rows.get_unchecked(i).as_()) + }, + allocator, + ) } diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index e9af70a6e92..cd4253ce1b8 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -5,8 +5,8 @@ use num_traits::AsPrimitive; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; -use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -28,7 +28,7 @@ use crate::require_child; pub(super) fn execute( mut array: Array, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { let num_values = array.num_values(); array = require_child!(array, array.array_indices(), 0 => Primitive); @@ -39,7 +39,7 @@ pub(super) fn execute( let validity = array.as_ref().validity()?; let output = match_each_native_ptype!(array.dtype().as_ptype(), |T| { - let values = gather_values::(&array)?; + let values = gather_values::(&array, ctx.allocator().clone())?; VortexResult::Ok(PrimitiveArray::new(values, validity)) })?; @@ -54,7 +54,10 @@ struct PrimitiveSource { row_mask: usize, } -fn gather_values(array: &Array) -> VortexResult> { +fn gather_values( + array: &Array, + allocator: BufferAllocatorRef, +) -> VortexResult> { let values = (0..array.num_values()) .map(|i| { let value = array.value(i); @@ -67,7 +70,7 @@ fn gather_values(array: &Array) -> VortexResult() .unwrap_or_default(); PrimitiveSource { - data: buffer![payload], + data: Buffer::full_in(payload, 1, allocator.clone()), len, row_mask: 0, } @@ -85,7 +88,12 @@ fn gather_values(array: &Array) -> VortexResult(), rows.as_slice::()) + gather( + &values, + branches.as_slice::(), + rows.as_slice::(), + allocator, + ) }) }) } @@ -94,6 +102,7 @@ fn gather( values: &[PrimitiveSource], branches: &[A], rows: &[R], + allocator: BufferAllocatorRef, ) -> VortexResult> where T: NativePType, @@ -108,31 +117,33 @@ where rows.len() ); - let output = - BufferMut::try_from_trusted_len_iter(branches.iter().zip(rows).map(|(branch, row)| { - let Some(source) = values.get((*branch).as_()) else { - vortex_bail!("interleave array index out of bounds"); - }; - let row = (*row).as_(); - vortex_ensure!(row < source.len, "interleave row index out of bounds"); - Ok(source.data[row & source.row_mask]) - }))?; + let mut output = BufferMut::with_capacity_in(branches.len(), allocator); + for (branch, row) in branches.iter().zip(rows) { + let Some(source) = values.get((*branch).as_()) else { + vortex_bail!("interleave array index out of bounds"); + }; + let row = (*row).as_(); + vortex_ensure!(row < source.len, "interleave row index out of bounds"); + output.push(source.data[row & source.row_mask]); + } Ok(output.freeze()) } #[cfg(test)] mod tests { use super::*; + use crate::memory::BufferAllocatorRef; #[test] fn rejects_out_of_bounds_selectors() { let values = [PrimitiveSource { - data: buffer![1u32], + data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()), len: 1, row_mask: 0, }]; - assert!(gather(&values, &[1u8], &[0u8]).is_err()); - assert!(gather(&values, &[0u8], &[1u8]).is_err()); + let allocator = BufferAllocatorRef::statically_allocated(); + assert!(gather(&values, &[1u8], &[0u8], allocator.clone()).is_err()); + assert!(gather(&values, &[0u8], &[1u8], allocator).is_err()); } } diff --git a/vortex-array/src/arrays/primitive/compute/between.rs b/vortex-array/src/arrays/primitive/compute/between.rs index 85f2b3d42e3..7012cb1be5f 100644 --- a/vortex-array/src/arrays/primitive/compute/between.rs +++ b/vortex-array/src/arrays/primitive/compute/between.rs @@ -24,7 +24,7 @@ impl BetweenKernel for Primitive { lower: &ArrayRef, upper: &ArrayRef, options: &BetweenOptions, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let (Some(lower), Some(upper)) = (lower.as_constant(), upper.as_constant()) else { return Ok(None); @@ -43,6 +43,7 @@ impl BetweenKernel for Primitive { P::try_from(&upper)?, nullability, options, + ctx, ) }))) } @@ -54,6 +55,7 @@ fn between_impl( upper: T, nullability: Nullability, options: &BetweenOptions, + ctx: &mut ExecutionCtx, ) -> ArrayRef { match (options.lower_strict, options.upper_strict) { // Note: these comparisons are explicitly passed in to allow function impl inlining @@ -64,6 +66,7 @@ fn between_impl( upper, NativePType::is_lt, nullability, + ctx, ), (StrictComparison::Strict, StrictComparison::NonStrict) => between_impl_( arr, @@ -72,6 +75,7 @@ fn between_impl( upper, NativePType::is_le, nullability, + ctx, ), (StrictComparison::NonStrict, StrictComparison::Strict) => between_impl_( arr, @@ -80,6 +84,7 @@ fn between_impl( upper, NativePType::is_lt, nullability, + ctx, ), (StrictComparison::NonStrict, StrictComparison::NonStrict) => between_impl_( arr, @@ -88,6 +93,7 @@ fn between_impl( upper, NativePType::is_le, nullability, + ctx, ), } } @@ -99,17 +105,22 @@ fn between_impl_( upper: T, upper_fn: impl Fn(T, T) -> bool, nullability: Nullability, + ctx: &mut ExecutionCtx, ) -> ArrayRef where T: NativePType + Copy, { let slice = arr.as_slice::(); BoolArray::new( - BitBuffer::collect_bool_multiversioned(slice.len(), |idx| { - // We only iterate upto arr len and |arr| == |slice|. - let i = unsafe { *slice.get_unchecked(idx) }; - lower_fn(lower, i) & upper_fn(i, upper) - }), + BitBuffer::collect_bool_multiversioned_in( + slice.len(), + |idx| { + // We only iterate upto arr len and |arr| == |slice|. + let i = unsafe { *slice.get_unchecked(idx) }; + lower_fn(lower, i) & upper_fn(i, upper) + }, + ctx.allocator().clone(), + ), arr.validity() .vortex_expect("validity should be derivable") .union_nullability(nullability), diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs index 361735536ed..6116a716ec8 100644 --- a/vortex-array/src/builders/dict/bytes.rs +++ b/vortex-array/src/builders/dict/bytes.rs @@ -8,7 +8,9 @@ use std::sync::Arc; use num_traits::AsPrimitive; use vortex_array::ExecutionCtx; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; @@ -50,29 +52,44 @@ pub struct BytesDictBuilder { dtype: DType, max_dict_bytes: usize, max_dict_len: usize, + allocator: BufferAllocatorRef, } -pub fn bytes_dict_builder(dtype: DType, constraints: &DictConstraints) -> Box { +pub fn bytes_dict_builder( + dtype: DType, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, +) -> Box { match constraints.max_len as u64 { - max if max <= u8::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u16::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - max if max <= u32::MAX as u64 => Box::new(BytesDictBuilder::::new(dtype, constraints)), - _ => Box::new(BytesDictBuilder::::new(dtype, constraints)), + max if max <= u8::MAX as u64 => { + Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)) + } + max if max <= u16::MAX as u64 => { + Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)) + } + max if max <= u32::MAX as u64 => { + Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)) + } + _ => Box::new(BytesDictBuilder::::new(dtype, constraints, allocator)), } } impl BytesDictBuilder { - pub fn new(dtype: DType, constraints: &DictConstraints) -> Self { + pub fn new(dtype: DType, constraints: &DictConstraints, allocator: BufferAllocatorRef) -> Self { Self { lookup: Some(HashTable::new()), - views: BufferMut::::empty(), + views: BufferMut::::empty_aligned_in( + Alignment::of::(), + allocator.clone(), + ), null_code: OnceCell::new(), - values: BufferMut::empty(), - values_nulls: BitBufferMut::empty(), + values: BufferMut::empty_aligned_in(Alignment::of::(), allocator.clone()), + values_nulls: BitBufferMut::empty_in(allocator.clone()), hasher: DefaultHashBuilder::default(), dtype, max_dict_bytes: constraints.max_bytes, max_dict_len: constraints.max_len, + allocator, } } @@ -172,7 +189,7 @@ impl BytesDictBuilder { F: FnMut(usize) -> &'a [u8], { let mut local_lookup = self.lookup.take().vortex_expect("Must have a lookup dict"); - let mut codes: BufferMut = BufferMut::with_capacity(len); + let mut codes = BufferMut::::with_capacity_in(len, self.allocator.clone()); match validity_mask.bit_buffer() { AllOr::All => { @@ -286,9 +303,21 @@ impl DictEncoder for BytesDictBuilder { } fn reset(&mut self) -> ArrayRef { - let views = mem::take(&mut self.views).freeze(); - let buffer = mem::take(&mut self.values).freeze(); - let value_nulls = mem::take(&mut self.values_nulls).freeze(); + let views = mem::replace( + &mut self.views, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + let buffer = mem::replace( + &mut self.values, + BufferMut::empty_aligned_in(Alignment::of::(), self.allocator.clone()), + ) + .freeze(); + let value_nulls = mem::replace( + &mut self.values_nulls, + BitBufferMut::empty_in(self.allocator.clone()), + ) + .freeze(); // SAFETY: we build the views explicitly and the bytes should be checked before feeding // to the encoder. diff --git a/vortex-array/src/builders/dict/mod.rs b/vortex-array/src/builders/dict/mod.rs index 94834f83cc4..6dc696d5b25 100644 --- a/vortex-array/src/builders/dict/mod.rs +++ b/vortex-array/src/builders/dict/mod.rs @@ -3,6 +3,7 @@ use bytes::bytes_dict_builder; use primitive::primitive_dict_builder; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_panic; @@ -45,14 +46,27 @@ pub trait DictEncoder: Send { } pub fn dict_encoder(array: &ArrayRef, constraints: &DictConstraints) -> Box { + dict_encoder_in( + array, + constraints, + BufferAllocatorRef::statically_allocated(), + ) +} + +/// Creates a dictionary encoder using the provided allocator. +pub fn dict_encoder_in( + array: &ArrayRef, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, +) -> Box { let dict_builder: Box = if let Some(pa) = array.as_opt::() { match_each_native_ptype!(pa.ptype(), |P| { - primitive_dict_builder::

(pa.dtype().nullability(), constraints) + primitive_dict_builder::

(pa.dtype().nullability(), constraints, allocator) }) } else if let Some(vbv) = array.as_opt::() { - bytes_dict_builder(vbv.dtype().clone(), constraints) + bytes_dict_builder(vbv.dtype().clone(), constraints, allocator) } else if let Some(vb) = array.as_opt::() { - bytes_dict_builder(vb.dtype().clone(), constraints) + bytes_dict_builder(vb.dtype().clone(), constraints, allocator) } else { vortex_panic!("Can only encode primitive or varbin/view arrays") }; @@ -67,7 +81,7 @@ pub fn dict_encode_with_constraints( constraints: &DictConstraints, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut encoder = dict_encoder(array, constraints); + let mut encoder = dict_encoder_in(array, constraints, ctx.allocator().clone()); let codes = encoder.encode(array, ctx)?.narrow(ctx)?; // SAFETY: The encoding process will produce a value set of codes and values // All values in the dictionary are guaranteed to be referenced by at least one code diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs index dea2932367b..7300be6cc2d 100644 --- a/vortex-array/src/builders/dict/primitive.rs +++ b/vortex-array/src/builders/dict/primitive.rs @@ -6,7 +6,9 @@ use std::hash::Hash; use std::mem; use rustc_hash::FxBuildHasher; +use vortex_buffer::Alignment; use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -31,6 +33,7 @@ use crate::validity::Validity; pub fn primitive_dict_builder( nullability: Nullability, constraints: &DictConstraints, + allocator: BufferAllocatorRef, ) -> Box where NativeValue: Hash + Eq, @@ -44,20 +47,25 @@ where width => vortex_panic!("invalid bit_width: {width}"), }); match max_possible_len { - max if max <= u8::MAX as u64 => { - Box::new(PrimitiveDictBuilder::::new(nullability, constraints)) - } + max if max <= u8::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( + nullability, + constraints, + allocator, + )), max if max <= u16::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( nullability, constraints, + allocator, )), max if max <= u32::MAX as u64 => Box::new(PrimitiveDictBuilder::::new( nullability, constraints, + allocator, )), _ => Box::new(PrimitiveDictBuilder::::new( nullability, constraints, + allocator, )), } } @@ -68,15 +76,20 @@ where NativeValue: Hash + Eq, Code: UnsignedPType, { - pub fn new(nullability: Nullability, constraints: &DictConstraints) -> Self { + pub fn new( + nullability: Nullability, + constraints: &DictConstraints, + allocator: BufferAllocatorRef, + ) -> Self { let max_dict_len = constraints .max_len .min(constraints.max_bytes / T::PTYPE.byte_width()); Self { lookup: HashMap::with_hasher(FxBuildHasher), null_code: OnceCell::new(), - values: BufferMut::::empty(), - values_nulls: BitBufferMut::empty(), + values: BufferMut::::empty_aligned_in(Alignment::of::(), allocator.clone()), + values_nulls: BitBufferMut::empty_in(allocator.clone()), + allocator, nullability, max_dict_len, } @@ -135,6 +148,7 @@ pub struct PrimitiveDictBuilder { values_nulls: BitBufferMut, nullability: Nullability, max_dict_len: usize, + allocator: BufferAllocatorRef, } impl DictEncoder for PrimitiveDictBuilder @@ -144,7 +158,7 @@ where Code: UnsignedPType, { fn encode(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let mut codes = BufferMut::::with_capacity(array.len()); + let mut codes = BufferMut::::with_capacity_in(array.len(), self.allocator.clone()); let prim = array.clone().execute::(ctx)?; match prim.validity()?.execute_mask(array.len(), ctx)? { @@ -183,9 +197,14 @@ where } fn reset(&mut self) -> ArrayRef { + let nulls = mem::replace( + &mut self.values_nulls, + BitBufferMut::empty_in(self.allocator.clone()), + ) + .freeze(); PrimitiveArray::new( self.values.clone(), - Validity::from_bit_buffer(mem::take(&mut self.values_nulls).freeze(), self.nullability), + Validity::from_bit_buffer(nulls, self.nullability), ) .into_array() } diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index 641efd61404..0c68de52307 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -648,6 +648,7 @@ impl Patches { self.offset(), self.values(), mask_indices, + ctx.allocator().clone(), ) }) } @@ -678,11 +679,15 @@ impl Patches { let patch_indices = self.indices().clone().execute::(ctx)?; match_each_unsigned_integer_ptype!(patch_indices.ptype(), |P| { let patch_indices = patch_indices.as_slice::

(); - Mask::from_buffer(BitBuffer::collect_bool(patch_indices.len(), |i| { - #[allow(clippy::cast_possible_truncation)] - let idx = (patch_indices[i] as usize) - self.offset; - !masked.value(idx) - })) + Mask::from_buffer(BitBuffer::collect_bool_in( + patch_indices.len(), + |i| { + #[allow(clippy::cast_possible_truncation)] + let idx = (patch_indices[i] as usize) - self.offset; + !masked.value(idx) + }, + ctx.allocator().clone(), + )) }) } }; @@ -855,6 +860,7 @@ impl Patches { .validity()? .execute_mask(take_indices.as_ref().len(), ctx)?, include_nulls, + ctx.allocator().clone(), |take_idx| { self.search_index_chunked_batch( patch_indices_slice, @@ -873,6 +879,7 @@ impl Patches { .validity()? .execute_mask(take_indices.as_ref().len(), ctx)?, include_nulls, + ctx.allocator().clone(), |take_idx| { let Some(offset) = ::from(self.offset) else { // If the offset cannot be converted to T, it's larger than all values in this array. @@ -939,6 +946,7 @@ impl Patches { min_index, max_index, include_nulls, + ctx.allocator().clone(), )? }) }) @@ -1028,6 +1036,7 @@ fn take_map, T: NativePType>( min_index: usize, max_index: usize, include_nulls: bool, + allocator: vortex_buffer::BufferAllocatorRef, ) -> VortexResult> where usize: TryFrom, @@ -1043,8 +1052,9 @@ where .map(|(value_index, sparse_index)| (sparse_index, value_index)) .collect(); - let mut new_sparse_indices = BufferMut::::with_capacity(take_indices.len()); - let mut value_indices = BufferMut::::with_capacity(take_indices.len()); + let mut new_sparse_indices = + BufferMut::::with_capacity_in(take_indices.len(), allocator.clone()); + let mut value_indices = BufferMut::::with_capacity_in(take_indices.len(), allocator); for (idx_in_take, &take_idx) in take_indices.iter().enumerate() { let ti = usize::try_from(take_idx) @@ -1094,9 +1104,10 @@ fn filter_patches_with_mask( offset: usize, patch_values: &ArrayRef, mask_indices: &[usize], + allocator: vortex_buffer::BufferAllocatorRef, ) -> VortexResult> { let true_count = mask_indices.len(); - let mut new_patch_indices = BufferMut::::with_capacity(true_count); + let mut new_patch_indices = BufferMut::::with_capacity_in(true_count, allocator); let mut new_mask_indices = Vec::with_capacity(true_count); // Attempt to move the window by `STRIDE` elements on each iteration. This assumes that @@ -1192,10 +1203,11 @@ fn take_indices_with_search_fn< take_indices: &[T], take_validity: Mask, include_nulls: bool, + allocator: vortex_buffer::BufferAllocatorRef, search_fn: F, ) -> VortexResult<(BufferMut, BufferMut)> { - let mut values_indices = BufferMut::with_capacity(take_indices.len()); - let mut new_indices = BufferMut::with_capacity(take_indices.len()); + let mut values_indices = BufferMut::with_capacity_in(take_indices.len(), allocator.clone()); + let mut new_indices = BufferMut::with_capacity_in(take_indices.len(), allocator); for (new_patch_idx, &take_idx) in take_indices.iter().enumerate() { if !take_validity.value(new_patch_idx) { diff --git a/vortex-array/src/scalar_fn/fns/binary/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/boolean.rs index 78f683eddf2..9e445440ed0 100644 --- a/vortex-array/src/scalar_fn/fns/binary/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/boolean.rs @@ -4,6 +4,7 @@ use std::iter::repeat_n; use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::read_u64_le; use vortex_error::VortexResult; @@ -249,6 +250,7 @@ pub fn kleene_boolean_buffers( &rhs_valid, operator, nullability, + ctx.allocator().clone(), ) } @@ -285,14 +287,14 @@ pub fn kleene_boolean_buffer_scalar( .execute_mask(len, ctx)? .bitand_not(&Mask::from_buffer(values)); BoolArray::try_new( - BitBuffer::new_unset(len), + BitBuffer::new_unset_in(len, ctx.allocator().clone()), Validity::from_mask(valid, nullability), )? } (Operator::Or, None) => { let valid = validity.execute_mask(len, ctx)? & &Mask::from_buffer(values); BoolArray::try_new( - BitBuffer::new_set(len), + BitBuffer::new_set_in(len, ctx.allocator().clone()), Validity::from_mask(valid, nullability), )? } @@ -302,6 +304,10 @@ pub fn kleene_boolean_buffer_scalar( Ok(result.into_array()) } +#[expect( + clippy::too_many_arguments, + reason = "four buffers plus result options" +)] fn fused_boolean_buffers( len: usize, lhs_values: &BitBuffer, @@ -310,6 +316,7 @@ fn fused_boolean_buffers( rhs_validity: &Mask, operator: Operator, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult { if let Some(result) = fused_boolean_buffers_aligned( len, @@ -319,6 +326,7 @@ fn fused_boolean_buffers( rhs_validity, operator, nullability, + allocator.clone(), )? { return Ok(result); } @@ -335,6 +343,7 @@ fn fused_boolean_buffers( $rhs_valid_words, operator, nullability, + allocator.clone(), ) }; } @@ -389,6 +398,10 @@ impl WordSource<'_> { } } +#[expect( + clippy::too_many_arguments, + reason = "four buffers plus result options" +)] fn fused_boolean_buffers_aligned( len: usize, lhs_values: &BitBuffer, @@ -397,6 +410,7 @@ fn fused_boolean_buffers_aligned( rhs_validity: &Mask, operator: Operator, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult> { let Some(lhs_values) = word_source_from_bit_buffer(lhs_values) else { return Ok(None); @@ -419,6 +433,7 @@ fn fused_boolean_buffers_aligned( rhs_validity, operator, nullability, + allocator, )?)) } @@ -434,6 +449,10 @@ fn word_source_from_mask(mask: &Mask) -> Option> { } } +#[expect( + clippy::too_many_arguments, + reason = "four word sources plus result options" +)] fn fused_boolean_word_sources( len: usize, lhs_words: WordSource<'_>, @@ -442,6 +461,7 @@ fn fused_boolean_word_sources( rhs_valid_words: WordSource<'_>, operator: Operator, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult { match operator { Operator::And => fused_boolean_and_word_sources( @@ -451,6 +471,7 @@ fn fused_boolean_word_sources( lhs_valid_words, rhs_valid_words, nullability, + allocator, ), Operator::Or => fused_boolean_or_word_sources( len, @@ -459,6 +480,7 @@ fn fused_boolean_word_sources( lhs_valid_words, rhs_valid_words, nullability, + allocator, ), other => vortex_bail!("Not a boolean operator: {other}"), } @@ -471,12 +493,13 @@ fn fused_boolean_and_word_sources( lhs_valid_words: WordSource<'_>, rhs_valid_words: WordSource<'_>, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult { let n_bytes = len.div_ceil(8); let n_words = n_bytes.div_ceil(8); let full_bytes = n_bytes - n_bytes % 8; - let mut values = BufferMut::::with_capacity(n_words); - let mut validity = BufferMut::::with_capacity(n_words); + let mut values = BufferMut::::with_capacity_in(n_words, allocator.clone()); + let mut validity = BufferMut::::with_capacity_in(n_words, allocator); for byte_offset in (0..full_bytes).step_by(8) { let lhs = lhs_words.word_at(byte_offset, 8); @@ -518,12 +541,13 @@ fn fused_boolean_or_word_sources( lhs_valid_words: WordSource<'_>, rhs_valid_words: WordSource<'_>, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult { let n_bytes = len.div_ceil(8); let n_words = n_bytes.div_ceil(8); let full_bytes = n_bytes - n_bytes % 8; - let mut values = BufferMut::::with_capacity(n_words); - let mut validity = BufferMut::::with_capacity(n_words); + let mut values = BufferMut::::with_capacity_in(n_words, allocator.clone()); + let mut validity = BufferMut::::with_capacity_in(n_words, allocator); for byte_offset in (0..full_bytes).step_by(8) { let lhs = lhs_words.word_at(byte_offset, 8); @@ -579,6 +603,10 @@ fn finish_fused_boolean_words( .into_array()) } +#[expect( + clippy::too_many_arguments, + reason = "four word iterators plus result options" +)] fn fused_boolean_words( len: usize, lhs_words: L, @@ -587,6 +615,7 @@ fn fused_boolean_words( rhs_valid_words: RV, operator: Operator, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult where L: Iterator, @@ -602,6 +631,7 @@ where lhs_valid_words, rhs_valid_words, nullability, + allocator, ), Operator::Or => fused_boolean_or_words( len, @@ -610,6 +640,7 @@ where lhs_valid_words, rhs_valid_words, nullability, + allocator, ), other => vortex_bail!("Not a boolean operator: {other}"), } @@ -622,6 +653,7 @@ fn fused_boolean_and_words( lhs_valid_words: LV, rhs_valid_words: RV, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult where L: Iterator, @@ -630,8 +662,8 @@ where RV: Iterator, { let n_words = len.div_ceil(64); - let mut values = BufferMut::::with_capacity(n_words); - let mut validity = BufferMut::::with_capacity(n_words); + let mut values = BufferMut::::with_capacity_in(n_words, allocator.clone()); + let mut validity = BufferMut::::with_capacity_in(n_words, allocator); for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words .zip(rhs_words) @@ -658,6 +690,7 @@ fn fused_boolean_or_words( lhs_valid_words: LV, rhs_valid_words: RV, nullability: Nullability, + allocator: BufferAllocatorRef, ) -> VortexResult where L: Iterator, @@ -666,8 +699,8 @@ where RV: Iterator, { let n_words = len.div_ceil(64); - let mut values = BufferMut::::with_capacity(n_words); - let mut validity = BufferMut::::with_capacity(n_words); + let mut values = BufferMut::::with_capacity_in(n_words, allocator.clone()); + let mut validity = BufferMut::::with_capacity_in(n_words, allocator); for (((lhs, rhs), lhs_valid), rhs_valid) in lhs_words .zip(rhs_words) diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs index 15bd486d341..f544c308af4 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs @@ -4,6 +4,7 @@ //! Native comparison of boolean arrays using word-wise bit operations. use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -77,16 +78,16 @@ pub(super) fn compare_bool( compare_bits(l, r, op) } (BoolOperand::Array { bits, .. }, BoolOperand::Constant { value, .. }) => { - compare_bits_constant(bits, value, op) + compare_bits_constant(bits, value, op, ctx.allocator().clone()) } (BoolOperand::Constant { value, .. }, BoolOperand::Array { bits, .. }) => { - compare_bits_constant(bits, value, op.swap()) + compare_bits_constant(bits, value, op.swap(), ctx.allocator().clone()) } (BoolOperand::Constant { value: l, .. }, BoolOperand::Constant { value: r, .. }) => { // Unreachable through `execute_compare` (constant-constant is folded there), but // cheap to answer anyway. let result = super::ordering_predicate(op)(l.cmp(&r)); - BitBuffer::full(result, len) + BitBuffer::full_in(result, len, ctx.allocator().clone()) } }; @@ -109,7 +110,12 @@ fn compare_bits(lhs: BitBuffer, rhs: BitBuffer, op: CompareOperator) -> BitBuffe } /// Compare array bits against a non-null constant: `bits value`. -fn compare_bits_constant(bits: BitBuffer, value: bool, op: CompareOperator) -> BitBuffer { +fn compare_bits_constant( + bits: BitBuffer, + value: bool, + op: CompareOperator, + allocator: BufferAllocatorRef, +) -> BitBuffer { let len = bits.len(); match (op, value) { (CompareOperator::Eq, true) @@ -120,7 +126,11 @@ fn compare_bits_constant(bits: BitBuffer, value: bool, op: CompareOperator) -> B | (CompareOperator::NotEq, true) | (CompareOperator::Lt, true) | (CompareOperator::Lte, false) => !bits, - (CompareOperator::Lt, false) | (CompareOperator::Gt, true) => BitBuffer::new_unset(len), - (CompareOperator::Lte, true) | (CompareOperator::Gte, false) => BitBuffer::new_set(len), + (CompareOperator::Lt, false) | (CompareOperator::Gt, true) => { + BitBuffer::new_unset_in(len, allocator) + } + (CompareOperator::Lte, true) | (CompareOperator::Gte, false) => { + BitBuffer::new_set_in(len, allocator) + } } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs index 70c56411d08..6e3b3799f42 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs @@ -11,6 +11,7 @@ use std::cmp::Ordering; use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -213,25 +214,44 @@ pub(super) fn compare_bytes( let bits = match (&lhs, &rhs) { (BytesOperand::Array { values: l, .. }, BytesOperand::Array { values: r, .. }) => { - compare_views(&ViewsSide::new(l), &ViewsSide::new(r), op) + compare_views( + &ViewsSide::new(l), + &ViewsSide::new(r), + op, + ctx.allocator().clone(), + ) } (BytesOperand::Array { values, .. }, BytesOperand::Constant { value, .. }) => { - compare_views_constant(&ViewsSide::new(values), value, op) + compare_views_constant(&ViewsSide::new(values), value, op, ctx.allocator().clone()) } (BytesOperand::Constant { value, .. }, BytesOperand::Array { values, .. }) => { - compare_views_constant(&ViewsSide::new(values), value, op.swap()) + compare_views_constant( + &ViewsSide::new(values), + value, + op.swap(), + ctx.allocator().clone(), + ) } (BytesOperand::Constant { value: l, .. }, BytesOperand::Constant { value: r, .. }) => { // Unreachable through `execute_compare` (constant-constant is folded there), but // cheap to answer anyway. - BitBuffer::full(ordering_predicate(op)(l.as_slice().cmp(r.as_slice())), len) + BitBuffer::full_in( + ordering_predicate(op)(l.as_slice().cmp(r.as_slice())), + len, + ctx.allocator().clone(), + ) } }; Ok(BoolArray::try_new(bits, validity)?.into_array()) } -fn compare_views(lhs: &ViewsSide<'_>, rhs: &ViewsSide<'_>, op: CompareOperator) -> BitBuffer { +fn compare_views( + lhs: &ViewsSide<'_>, + rhs: &ViewsSide<'_>, + op: CompareOperator, + allocator: BufferAllocatorRef, +) -> BitBuffer { let len = lhs.len(); // The unchecked view accesses below index both sides with `i < len`, so this must hold even // in release builds. @@ -239,18 +259,26 @@ fn compare_views(lhs: &ViewsSide<'_>, rhs: &ViewsSide<'_>, op: CompareOperator) // Dispatch the operator outside the lane loop so each predicate inlines into its own loop; // a shared `fn(Ordering) -> bool` pointer would cost an indirect call per lane. match op { - CompareOperator::Eq => BitBuffer::collect_bool(len, |i| { - // SAFETY: `collect_bool` yields i < len == views.len() for both sides. - unsafe { view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } - }), - CompareOperator::NotEq => BitBuffer::collect_bool(len, |i| { - // SAFETY: `collect_bool` yields i < len == views.len() for both sides. - unsafe { !view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } - }), - CompareOperator::Gt => collect_ordering_bits(lhs, rhs, Ordering::is_gt), - CompareOperator::Gte => collect_ordering_bits(lhs, rhs, Ordering::is_ge), - CompareOperator::Lt => collect_ordering_bits(lhs, rhs, Ordering::is_lt), - CompareOperator::Lte => collect_ordering_bits(lhs, rhs, Ordering::is_le), + CompareOperator::Eq => BitBuffer::collect_bool_in( + len, + |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + unsafe { view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } + }, + allocator, + ), + CompareOperator::NotEq => BitBuffer::collect_bool_in( + len, + |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + unsafe { !view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } + }, + allocator, + ), + CompareOperator::Gt => collect_ordering_bits(lhs, rhs, Ordering::is_gt, allocator), + CompareOperator::Gte => collect_ordering_bits(lhs, rhs, Ordering::is_ge, allocator), + CompareOperator::Lt => collect_ordering_bits(lhs, rhs, Ordering::is_lt, allocator), + CompareOperator::Lte => collect_ordering_bits(lhs, rhs, Ordering::is_le, allocator), } } @@ -259,16 +287,26 @@ fn collect_ordering_bits( lhs: &ViewsSide<'_>, rhs: &ViewsSide<'_>, predicate: impl Fn(Ordering) -> bool, + allocator: BufferAllocatorRef, ) -> BitBuffer { let len = lhs.len(); assert_eq!(len, rhs.len(), "compared views must have equal lengths"); - BitBuffer::collect_bool(len, |i| { - // SAFETY: `collect_bool` yields i < len == views.len() for both sides. - predicate(unsafe { view_cmp(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) }) - }) + BitBuffer::collect_bool_in( + len, + |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + predicate(unsafe { view_cmp(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) }) + }, + allocator, + ) } -fn compare_views_constant(lhs: &ViewsSide<'_>, constant: &[u8], op: CompareOperator) -> BitBuffer { +fn compare_views_constant( + lhs: &ViewsSide<'_>, + constant: &[u8], + op: CompareOperator, + allocator: BufferAllocatorRef, +) -> BitBuffer { let len = lhs.len(); // The same head/prefix/tail words a view stores, precomputed once for the constant. let mut prefix_bytes = [0u8; 4]; @@ -290,22 +328,31 @@ fn compare_views_constant(lhs: &ViewsSide<'_>, constant: &[u8], op: CompareOpera u128::from(constant_head) | (u128::from(u64::from_le_bytes(tail_bytes)) << 64); match op { - CompareOperator::Eq => BitBuffer::collect_bool(len, |i| { - // SAFETY: `collect_bool` yields i < len == views.len(). - let view = unsafe { lhs.view_unchecked(i) }; - constant_eq(lhs, view, constant, constant_head, constant_inlined) - }), - CompareOperator::NotEq => BitBuffer::collect_bool(len, |i| { - // SAFETY: `collect_bool` yields i < len == views.len(). - let view = unsafe { lhs.view_unchecked(i) }; - !constant_eq(lhs, view, constant, constant_head, constant_inlined) - }), + CompareOperator::Eq => BitBuffer::collect_bool_in( + len, + |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + constant_eq(lhs, view, constant, constant_head, constant_inlined) + }, + allocator, + ), + CompareOperator::NotEq => BitBuffer::collect_bool_in( + len, + |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + !constant_eq(lhs, view, constant, constant_head, constant_inlined) + }, + allocator, + ), CompareOperator::Gt => collect_constant_ordering_bits( lhs, constant, constant_prefix, constant_tail, Ordering::is_gt, + allocator, ), CompareOperator::Gte => collect_constant_ordering_bits( lhs, @@ -313,6 +360,7 @@ fn compare_views_constant(lhs: &ViewsSide<'_>, constant: &[u8], op: CompareOpera constant_prefix, constant_tail, Ordering::is_ge, + allocator, ), CompareOperator::Lt => collect_constant_ordering_bits( lhs, @@ -320,6 +368,7 @@ fn compare_views_constant(lhs: &ViewsSide<'_>, constant: &[u8], op: CompareOpera constant_prefix, constant_tail, Ordering::is_lt, + allocator, ), CompareOperator::Lte => collect_constant_ordering_bits( lhs, @@ -327,6 +376,7 @@ fn compare_views_constant(lhs: &ViewsSide<'_>, constant: &[u8], op: CompareOpera constant_prefix, constant_tail, Ordering::is_le, + allocator, ), } } @@ -338,18 +388,23 @@ fn collect_constant_ordering_bits( constant_prefix: u32, constant_tail: u64, predicate: impl Fn(Ordering) -> bool, + allocator: BufferAllocatorRef, ) -> BitBuffer { - BitBuffer::collect_bool(lhs.len(), |i| { - // SAFETY: `collect_bool` yields i < len == views.len(). - let view = unsafe { lhs.view_unchecked(i) }; - predicate(constant_cmp( - lhs, - view, - constant, - constant_prefix, - constant_tail, - )) - }) + BitBuffer::collect_bool_in( + lhs.len(), + |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + predicate(constant_cmp( + lhs, + view, + constant, + constant_prefix, + constant_tail, + )) + }, + allocator, + ) } /// Compare a view against a constant for equality using the constant's precomputed head and diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs index eea6589645e..78d5827ad12 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs @@ -10,6 +10,7 @@ //! [`DecimalDType`]: crate::dtype::DecimalDType use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -87,19 +88,23 @@ pub(super) fn compare_decimal( let bits = match (lhs, rhs) { (DecimalOperand::Array { values: l, .. }, DecimalOperand::Array { values: r, .. }) => { - compare_decimal_values(&l, &r, op) + compare_decimal_values(&l, &r, op, ctx.allocator().clone()) } (DecimalOperand::Array { values, .. }, DecimalOperand::Constant { value, .. }) => { - compare_decimal_constant(&values, value, op) + compare_decimal_constant(&values, value, op, ctx.allocator().clone()) } (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { values, .. }) => { - compare_decimal_constant(&values, value, op.swap()) + compare_decimal_constant(&values, value, op.swap(), ctx.allocator().clone()) } (DecimalOperand::Constant { value: l, .. }, DecimalOperand::Constant { value: r, .. }) => { // Unreachable through `execute_compare` (constant-constant is folded there), but // cheap to answer anyway. let ordering = l.as_i256().cmp(&r.as_i256()); - BitBuffer::full(super::ordering_predicate(op)(ordering), len) + BitBuffer::full_in( + super::ordering_predicate(op)(ordering), + len, + ctx.allocator().clone(), + ) } }; @@ -110,12 +115,13 @@ fn compare_decimal_values( lhs: &DecimalArray, rhs: &DecimalArray, op: CompareOperator, + allocator: BufferAllocatorRef, ) -> BitBuffer { let common = lhs.values_type().max(rhs.values_type()); match_each_decimal_value_type!(common, |W| { let lhs = widened_buffer::(lhs); let rhs = widened_buffer::(rhs); - compare_slices::(&lhs, &rhs, op) + compare_slices::(&lhs, &rhs, op, allocator) }) } @@ -123,10 +129,11 @@ fn compare_decimal_constant( array: &DecimalArray, constant: DecimalValue, op: CompareOperator, + allocator: BufferAllocatorRef, ) -> BitBuffer { match_each_decimal_value_type!(array.values_type(), |T| { match constant.cast::() { - Some(value) => compare_slice_constant::(&array.buffer::(), value, op), + Some(value) => compare_slice_constant::(&array.buffer::(), value, op, allocator), None => { // The constant does not fit the array's storage type, so it is either greater // than every possible array value or less than every possible array value; the @@ -139,20 +146,25 @@ fn compare_decimal_constant( CompareOperator::Lt | CompareOperator::Lte => constant_greater, CompareOperator::Gt | CompareOperator::Gte => !constant_greater, }; - BitBuffer::full(result, array.len()) + BitBuffer::full_in(result, array.len(), allocator) } } }) } -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { +fn compare_slices( + lhs: &[T], + rhs: &[T], + op: CompareOperator, + allocator: BufferAllocatorRef, +) -> BitBuffer { match op { - CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a == b), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| a != b), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, |a: T, b: T| a > b), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, |a: T, b: T| a >= b), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, |a: T, b: T| a < b), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, |a: T, b: T| a <= b), + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a == b, allocator), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| a != b, allocator), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, |a: T, b: T| a > b, allocator), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, |a: T, b: T| a >= b, allocator), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, |a: T, b: T| a < b, allocator), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, |a: T, b: T| a <= b, allocator), } } @@ -160,13 +172,14 @@ fn compare_slice_constant( values: &[T], constant: T, op: CompareOperator, + allocator: BufferAllocatorRef, ) -> BitBuffer { match op { - CompareOperator::Eq => collect_bits(values, |a: T| a == constant), - CompareOperator::NotEq => collect_bits(values, |a: T| a != constant), - CompareOperator::Gt => collect_bits(values, |a: T| a > constant), - CompareOperator::Gte => collect_bits(values, |a: T| a >= constant), - CompareOperator::Lt => collect_bits(values, |a: T| a < constant), - CompareOperator::Lte => collect_bits(values, |a: T| a <= constant), + CompareOperator::Eq => collect_bits(values, |a: T| a == constant, allocator), + CompareOperator::NotEq => collect_bits(values, |a: T| a != constant, allocator), + CompareOperator::Gt => collect_bits(values, |a: T| a > constant, allocator), + CompareOperator::Gte => collect_bits(values, |a: T| a >= constant, allocator), + CompareOperator::Lt => collect_bits(values, |a: T| a < constant, allocator), + CompareOperator::Lte => collect_bits(values, |a: T| a <= constant, allocator), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs index 8b5a4522b07..5b096fefd98 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -14,6 +14,7 @@ use std::cmp::Ordering; use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_compute::lane_kernels::LaneZip; @@ -294,17 +295,22 @@ pub(super) fn collect_zip_bits( lhs: &[T], rhs: &[T], f: impl Fn(T, T) -> bool, + allocator: BufferAllocatorRef, ) -> BitBuffer { let len = lhs.len(); - let mut words = BufferMut::::zeroed(len.div_ceil(64)); + let mut words = BufferMut::::zeroed_in(len.div_ceil(64), allocator); LaneZip::new(lhs, rhs).map_bits_into(words.as_mut_slice(), |(a, b)| f(a, b)); bit_buffer_from_words(words, len) } /// Bit-pack the predicate `f(values[i])` over a slice into a [`BitBuffer`]. -pub(super) fn collect_bits(values: &[T], f: impl Fn(T) -> bool) -> BitBuffer { +pub(super) fn collect_bits( + values: &[T], + f: impl Fn(T) -> bool, + allocator: BufferAllocatorRef, +) -> BitBuffer { let len = values.len(); - let mut words = BufferMut::::zeroed(len.div_ceil(64)); + let mut words = BufferMut::::zeroed_in(len.div_ceil(64), allocator); values.map_bits_into(words.as_mut_slice(), f); bit_buffer_from_words(words, len) } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs index 7fd71262030..628f55cf17a 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -16,6 +16,7 @@ use std::cmp::Ordering; use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; @@ -75,12 +76,24 @@ pub(super) fn compare_nested( // Dispatch the operator outside the row loop so the predicate inlines into each loop; the // comparator call itself stays virtual. let bits = match op { - CompareOperator::Eq => collect_ordering_bits(len, &comparator, Ordering::is_eq), - CompareOperator::NotEq => collect_ordering_bits(len, &comparator, Ordering::is_ne), - CompareOperator::Gt => collect_ordering_bits(len, &comparator, Ordering::is_gt), - CompareOperator::Gte => collect_ordering_bits(len, &comparator, Ordering::is_ge), - CompareOperator::Lt => collect_ordering_bits(len, &comparator, Ordering::is_lt), - CompareOperator::Lte => collect_ordering_bits(len, &comparator, Ordering::is_le), + CompareOperator::Eq => { + collect_ordering_bits(len, &comparator, Ordering::is_eq, ctx.allocator().clone()) + } + CompareOperator::NotEq => { + collect_ordering_bits(len, &comparator, Ordering::is_ne, ctx.allocator().clone()) + } + CompareOperator::Gt => { + collect_ordering_bits(len, &comparator, Ordering::is_gt, ctx.allocator().clone()) + } + CompareOperator::Gte => { + collect_ordering_bits(len, &comparator, Ordering::is_ge, ctx.allocator().clone()) + } + CompareOperator::Lt => { + collect_ordering_bits(len, &comparator, Ordering::is_lt, ctx.allocator().clone()) + } + CompareOperator::Lte => { + collect_ordering_bits(len, &comparator, Ordering::is_le, ctx.allocator().clone()) + } }; Ok(BoolArray::try_new(bits, validity)?.into_array()) @@ -91,8 +104,9 @@ fn collect_ordering_bits( len: usize, comparator: &RowComparator, predicate: impl Fn(Ordering) -> bool, + allocator: BufferAllocatorRef, ) -> BitBuffer { - BitBuffer::collect_bool(len, |i| predicate(comparator(i, i))) + BitBuffer::collect_bool_in(len, |i| predicate(comparator(i, i)), allocator) } /// The validity mask of a recursively canonical array. diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs index 1247358dce1..9d9d61a1198 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -4,6 +4,7 @@ //! Native comparison of primitive arrays via bit-packing lane kernels. use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -65,22 +66,22 @@ fn compare_primitive_typed( ( PrimitiveOperand::Array { values: lhs, .. }, PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slices(lhs, rhs, op), + ) => compare_slices(lhs, rhs, op, ctx.allocator().clone()), ( PrimitiveOperand::Array { values: lhs, .. }, PrimitiveOperand::Constant { value: rhs, .. }, - ) => compare_slice_constant(lhs, *rhs, op), + ) => compare_slice_constant(lhs, *rhs, op, ctx.allocator().clone()), ( PrimitiveOperand::Constant { value: lhs, .. }, PrimitiveOperand::Array { values: rhs, .. }, - ) => compare_slice_constant(rhs, *lhs, op.swap()), + ) => compare_slice_constant(rhs, *lhs, op.swap(), ctx.allocator().clone()), ( PrimitiveOperand::Constant { value: lhs, .. }, PrimitiveOperand::Constant { value: rhs, .. }, ) => { // Unreachable through `execute_compare` (constant-constant is folded there), but // cheap to answer anyway. - BitBuffer::full(apply_op(*lhs, *rhs, op), len) + BitBuffer::full_in(apply_op(*lhs, *rhs, op), len, ctx.allocator().clone()) } (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { return Ok( @@ -105,26 +106,36 @@ fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { } } -fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { +fn compare_slices( + lhs: &[T], + rhs: &[T], + op: CompareOperator, + allocator: BufferAllocatorRef, +) -> BitBuffer { // Dispatch the operator outside the lane loop so each instantiation vectorizes a single // branch-free predicate. match op { - CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), - CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), - CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), - CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), - CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), - CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b), allocator), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b), allocator), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt, allocator), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge, allocator), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt, allocator), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le, allocator), } } -fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { +fn compare_slice_constant( + lhs: &[T], + rhs: T, + op: CompareOperator, + allocator: BufferAllocatorRef, +) -> BitBuffer { match op { - CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), - CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), - CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), - CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), - CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), - CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), + CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs), allocator), + CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs), allocator), + CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs), allocator), + CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs), allocator), + CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs), allocator), + CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs), allocator), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs index 9831a963354..fc18d20e744 100644 --- a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -3,8 +3,10 @@ use std::sync::Arc; +use allocator_api2::alloc::Global; use rstest::rstest; use vortex_buffer::BitBuffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -24,6 +26,7 @@ use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinArray; use crate::arrays::VarBinViewArray; +use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; use crate::builders::ArrayBuilder; use crate::builders::MapBuilder; @@ -37,6 +40,7 @@ use crate::dtype::PType; use crate::extension::datetime::TimeUnit; use crate::extension::datetime::Timestamp; use crate::extension::datetime::TimestampOptions; +use crate::memory::MemorySessionExt; use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::scalar_fn::fns::binary::scalar_cmp; @@ -359,6 +363,23 @@ fn execute_compare_test(lhs: ArrayRef, rhs: ArrayRef, op: Operator) -> ArrayRef lhs.binary(rhs, op).unwrap() } +#[test] +fn comparison_uses_execution_allocator() -> VortexResult<()> { + let allocator = BufferAllocatorRef::new(Global); + let mut ctx = array_session() + .with_allocator(allocator.clone()) + .create_execution_ctx(); + let result = buffer![1i32, 2, 3] + .into_array() + .binary(buffer![1i32, 0, 3].into_array(), Operator::Eq)? + .execute::(&mut ctx)?; + let output = result.to_bit_buffer(); + let output_allocator = output.inner().allocator(); + + assert!(output_allocator.ptr_eq(&allocator)); + Ok(()) +} + #[rstest] #[case(Operator::Eq, [false, true, false, false])] #[case(Operator::NotEq, [true, false, true, true])] diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index 1901f0260c0..a3bad7c87ef 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -9,6 +9,7 @@ //! scanning the finished output. use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::IndexedSource; use vortex_compute::lane_kernels::IndexedSourceExt; @@ -32,6 +33,7 @@ pub(super) fn checked_lanes( source: S, valid_rows: &Mask, apply: Apply, + allocator: BufferAllocatorRef, ) -> Result, usize> where S: IndexedSource, @@ -43,11 +45,11 @@ where let valid_bits = match valid_rows.bit_buffer() { AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), + AllOr::None => return Ok(Buffer::zeroed_in(len, allocator)), AllOr::Some(valid_bits) => Some(valid_bits), }; - let mut values = BufferMut::::with_capacity(len); + let mut values = BufferMut::::with_capacity_in(len, allocator); let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs index fa9ffb0b5f1..0c02590e59f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/decimal.rs @@ -23,6 +23,8 @@ use num_traits::CheckedDiv; use num_traits::CheckedMul; use num_traits::CheckedSub; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; use vortex_compute::lane_kernels::LaneZip; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -99,6 +101,7 @@ pub(super) fn execute_numeric_decimal( validity, &valid_rows, &constants, + ctx.allocator().clone(), ) }; } @@ -318,6 +321,10 @@ impl CheckedDecimalOp for CheckedDecimalDiv { } } +#[expect( + clippy::too_many_arguments, + reason = "typed decimal execution parameters" +)] fn execute_decimal_typed( lhs: &DecimalOperand, rhs: &DecimalOperand, @@ -326,6 +333,7 @@ fn execute_decimal_typed( validity: Validity, valid_rows: &Mask, constants: &DecimalOpConstants, + allocator: BufferAllocatorRef, ) -> VortexResult where W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, @@ -336,24 +344,30 @@ where let values = match (lhs, rhs) { (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Array { values: rhs, .. }) => { - checked_decimal_arrays::(lhs, rhs, constants, valid_rows) + checked_decimal_arrays::(lhs, rhs, constants, valid_rows, allocator.clone()) } (DecimalOperand::Array { values: lhs, .. }, DecimalOperand::Constant { value, .. }) => { let rhs = typed_constant::(value); match_each_decimal_value_type!(lhs.values_type(), |L| { let lhs = lhs.buffer::(); - checked_lanes(lhs.as_slice(), valid_rows, |lhs| { - Op::apply(::from(lhs)?, rhs, constants) - }) + checked_lanes( + lhs.as_slice(), + valid_rows, + |lhs| Op::apply(::from(lhs)?, rhs, constants), + allocator.clone(), + ) }) } (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { values: rhs, .. }) => { let lhs = typed_constant::(value); match_each_decimal_value_type!(rhs.values_type(), |R| { let rhs = rhs.buffer::(); - checked_lanes(rhs.as_slice(), valid_rows, |rhs| { - Op::apply(lhs, ::from(rhs)?, constants) - }) + checked_lanes( + rhs.as_slice(), + valid_rows, + |rhs| Op::apply(lhs, ::from(rhs)?, constants), + allocator.clone(), + ) }) } ( @@ -380,6 +394,7 @@ where values, result_decimal_dtype, validity.union_nullability(result_dtype.nullability()), + allocator, )) } @@ -390,6 +405,7 @@ fn decimal_array_narrowed( values: Buffer, decimal_dtype: DecimalDType, validity: Validity, + allocator: BufferAllocatorRef, ) -> ArrayRef { let target = DecimalType::smallest_decimal_value_type(&decimal_dtype); if target == W::DECIMAL_TYPE { @@ -397,16 +413,12 @@ fn decimal_array_narrowed( } match_each_decimal_value_type!(target, |O| { - let narrowed: Buffer = values - .as_slice() - .iter() - .copied() - .map(|value| { - ::from(value) - .vortex_expect("precision-checked decimal result must fit the output width") - }) - .collect(); - DecimalArray::new(narrowed, decimal_dtype, validity).into_array() + let mut narrowed = BufferMut::with_capacity_in(values.len(), allocator); + narrowed.extend(values.as_slice().iter().copied().map(|value| { + ::from(value) + .vortex_expect("precision-checked decimal result must fit the output width") + })); + DecimalArray::new(narrowed.freeze(), decimal_dtype, validity).into_array() }) } @@ -415,6 +427,7 @@ fn checked_decimal_arrays( rhs: &DecimalArray, constants: &DecimalOpConstants, valid_rows: &Mask, + allocator: BufferAllocatorRef, ) -> Result, usize> where W: NativeDecimalType + CheckedAdd + CheckedSub + CheckedMul + CheckedDiv + Mul, @@ -435,6 +448,7 @@ where constants, ) }, + allocator, ) }) }) diff --git a/vortex-array/src/scalar_fn/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index 0b9af18d3c6..7fde38e9193 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -12,6 +12,8 @@ pub use kernel::*; use pattern::LikePattern; use prost::Message; use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferAllocatorRef; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -234,7 +236,12 @@ pub(crate) fn execute_like( options.case_insensitive, ascii_haystack, )?; - let bits = eval_pattern(&haystack, &compiled, options.negated); + let bits = eval_pattern( + &haystack, + &compiled, + options.negated, + ctx.allocator().clone(), + ); let validity = values.validity()?.union_nullability(nullability); return Ok(BoolArray::new(bits, validity).into_array()); } @@ -246,10 +253,10 @@ pub(crate) fn execute_like( let pattern_views = ResolvedViews::new(&patterns); let ascii_haystack = options.case_insensitive && haystack.is_ascii(); - let mut bits = Vec::with_capacity(len); // Reuse the previous row's compiled pattern while the pattern bytes repeat, so runs of // identical patterns (the common case for non-constant pattern children) compile once. let mut cached: Option<(&[u8], LikePattern)> = None; + let mut bits = BitBufferMut::with_capacity_in(len, ctx.allocator().clone()); for i in 0..len { let pattern_bytes = pattern_views.bytes(i); let compiled = match &cached { @@ -265,13 +272,13 @@ pub(crate) fn execute_like( &cached.insert((pattern_bytes, compiled)).1 } }; - bits.push(compiled.matches(haystack.bytes(i)) != options.negated); + bits.append(compiled.matches(haystack.bytes(i)) != options.negated); } let validity = values .validity()? .and(patterns.validity()?)? .union_nullability(nullability); - Ok(BoolArray::new(BitBuffer::from_iter(bits), validity).into_array()) + Ok(BoolArray::new(bits.freeze(), validity).into_array()) } /// Resolved views over a canonical [`VarBinViewArray`]: the view structs plus borrowed slices @@ -340,27 +347,38 @@ impl<'a> ResolvedViews<'a> { /// The equality, prefix, and suffix patterns exploit the view layout: a view stores the value /// length and its first four bytes inline, which settles most elements without touching the /// data buffers (values of up to 12 bytes are stored entirely inline). -fn eval_pattern(haystack: &ResolvedViews<'_>, pattern: &LikePattern, negated: bool) -> BitBuffer { +fn eval_pattern( + haystack: &ResolvedViews<'_>, + pattern: &LikePattern, + negated: bool, + allocator: BufferAllocatorRef, +) -> BitBuffer { let len = haystack.views.len(); match pattern { LikePattern::Eq(needle) if needle.len() <= BinaryView::MAX_INLINED_SIZE => { // The needle fits in a view, so equality is a single 16-byte comparison: a view // of a different length or prefix can never share the same bit pattern. let needle_view = BinaryView::new_inlined(needle).as_u128(); - BitBuffer::collect_bool(len, |i| { - (haystack.views[i].as_u128() == needle_view) != negated - }) + BitBuffer::collect_bool_in( + len, + |i| (haystack.views[i].as_u128() == needle_view) != negated, + allocator, + ) } LikePattern::Eq(needle) => { // Compare the view head (length plus 4-byte prefix) first; only views that agree // on both dereference their data buffer for the remaining bytes. let needle_head = needle_head(needle); - BitBuffer::collect_bool(len, |i| { - let view = &haystack.views[i]; - let matched = - view_head(view) == needle_head && haystack.bytes(i)[4..] == needle[4..]; - matched != negated - }) + BitBuffer::collect_bool_in( + len, + |i| { + let view = &haystack.views[i]; + let matched = + view_head(view) == needle_head && haystack.bytes(i)[4..] == needle[4..]; + matched != negated + }, + allocator, + ) } LikePattern::StartsWith(needle) => { // A branch-free masked comparison of the view's inline 4-byte prefix rejects @@ -378,42 +396,62 @@ fn eval_pattern(haystack: &ResolvedViews<'_>, pattern: &LikePattern, negated: bo } else { (1u32 << (8 * prefix_len)) - 1 }; - BitBuffer::collect_bool(len, |i| { - let view = &haystack.views[i]; - let matched = view.len() as usize >= needle_len - && (view_prefix(view) & prefix_mask) == needle_prefix - && (needle_len <= 4 || haystack.bytes(i)[4..needle_len] == needle[4..]); - matched != negated - }) + BitBuffer::collect_bool_in( + len, + |i| { + let view = &haystack.views[i]; + let matched = view.len() as usize >= needle_len + && (view_prefix(view) & prefix_mask) == needle_prefix + && (needle_len <= 4 || haystack.bytes(i)[4..needle_len] == needle[4..]); + matched != negated + }, + allocator, + ) } LikePattern::EndsWith(needle) => { // Inlined values compare their suffix inside the view struct without touching the // data buffers; reference views slice exactly the suffix out of their buffer. let needle_len = needle.len(); - BitBuffer::collect_bool(len, |i| { - // SAFETY: `i` is below the array length, and the suffix length is only read - // once the view is known to be at least `needle_len` long. - let matched = unsafe { - let view = haystack.views.get_unchecked(i); - view.len() as usize >= needle_len - && bytes_eq(haystack.suffix_bytes_unchecked(view, needle_len), needle) - }; - matched != negated - }) + BitBuffer::collect_bool_in( + len, + |i| { + // SAFETY: `i` is below the array length, and the suffix length is only read + // once the view is known to be at least `needle_len` long. + let matched = unsafe { + let view = haystack.views.get_unchecked(i); + view.len() as usize >= needle_len + && bytes_eq(haystack.suffix_bytes_unchecked(view, needle_len), needle) + }; + matched != negated + }, + allocator, + ) } - LikePattern::IEqAscii(needle) => BitBuffer::collect_bool(len, |i| { - let view = &haystack.views[i]; - let matched = view.len() as usize == needle.len() - && haystack.bytes(i).eq_ignore_ascii_case(needle); - matched != negated - }), - LikePattern::Contains(finder, needle_len) => BitBuffer::collect_bool(len, |i| { - let view = &haystack.views[i]; - let matched = - view.len() as usize >= *needle_len && finder.find(haystack.bytes(i)).is_some(); - matched != negated - }), - _ => BitBuffer::collect_bool(len, |i| pattern.matches(haystack.bytes(i)) != negated), + LikePattern::IEqAscii(needle) => BitBuffer::collect_bool_in( + len, + |i| { + let view = &haystack.views[i]; + let matched = view.len() as usize == needle.len() + && haystack.bytes(i).eq_ignore_ascii_case(needle); + matched != negated + }, + allocator, + ), + LikePattern::Contains(finder, needle_len) => BitBuffer::collect_bool_in( + len, + |i| { + let view = &haystack.views[i]; + let matched = + view.len() as usize >= *needle_len && finder.find(haystack.bytes(i)).is_some(); + matched != negated + }, + allocator, + ), + _ => BitBuffer::collect_bool_in( + len, + |i| pattern.matches(haystack.bytes(i)) != negated, + allocator, + ), } } diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index d2508014089..28b147ec0d5 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -242,7 +242,7 @@ fn list_contains_scalar( let elems = list_array.elements(); if elems.is_empty() { // Must return false when a list is empty (but valid), or null when the list itself is null. - return list_false_or_null(&list_array, nullability); + return list_false_or_null(&list_array, nullability, ctx); } let rhs = ConstantArray::new(value.clone(), elems.len()); @@ -263,7 +263,7 @@ fn list_contains_scalar( "Search value must not be null here" ); // False, unless the list itself is null in which case we return null. - list_false_or_null(&list_array, nullability) + list_false_or_null(&list_array, nullability, ctx) } // No elements match, and all comparisons are valid (result in `false`). Some(false) => { @@ -294,7 +294,7 @@ fn list_contains_scalar( // Process based on the offset and size types. let list_matches = match_each_unsigned_integer_ptype!(offsets.ptype(), |O| { match_each_unsigned_integer_ptype!(sizes.ptype(), |S| { - process_matches::(matches, list_array.len(), offsets, sizes) + process_matches::(matches, list_array.len(), offsets, sizes, ctx) }) }); @@ -312,6 +312,7 @@ fn process_matches( list_array_len: usize, offsets: PrimitiveArray, sizes: PrimitiveArray, + ctx: &mut ExecutionCtx, ) -> BitBuffer where O: IntegerPType, @@ -321,8 +322,9 @@ where let sizes_slice = sizes.as_slice::(); let bits = matches.bit_buffer_view(); - (0..list_array_len) - .map(|i| { + BitBuffer::collect_bool_in( + list_array_len, + |i| { let offset = offsets_slice[i].as_(); let size = sizes_slice[i].as_(); @@ -330,8 +332,9 @@ where // `Some(_)`, at least one element in this list's range matches. let mut set_bits = BitIndexIterator::new(bits.inner(), offset, size); set_bits.next().is_some() - }) - .collect::() + }, + ctx.allocator().clone(), + ) } /// Returns a `Bool` array with `false` for lists that are valid, @@ -339,6 +342,7 @@ where fn list_false_or_null( list_array: &ListViewArray, nullability: Nullability, + ctx: &mut ExecutionCtx, ) -> VortexResult { match list_array.validity()? { Validity::NonNullable => { @@ -362,7 +366,7 @@ fn list_false_or_null( } Validity::Array(validity_array) => { // Create a new bool array with false, and the provided nulls - let buffer = BitBuffer::new_unset(list_array.len()); + let buffer = BitBuffer::new_unset_in(list_array.len(), ctx.allocator().clone()); Ok(BoolArray::new(buffer, Validity::Array(validity_array)).into_array()) } } @@ -386,7 +390,11 @@ fn list_is_not_empty( let sizes = list_array.sizes().clone().execute::(ctx)?; let buffer = match_each_integer_ptype!(sizes.ptype(), |S| { - BitBuffer::from_iter(sizes.as_slice::().iter().map(|&size| size != S::zero())) + BitBuffer::collect_bool_in( + sizes.len(), + |idx| sizes.as_slice::()[idx] != S::zero(), + ctx.allocator().clone(), + ) }); // Copy over the validity mask from the input. diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 2a81b00d5c1..206f6b456db 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -15,6 +15,7 @@ use crate::Alignment; use crate::BitBufferMeta; use crate::BitBufferMut; use crate::Buffer; +use crate::BufferAllocatorRef; use crate::BufferMut; use crate::ByteBuffer; use crate::bit::BitChunks; @@ -30,7 +31,6 @@ use crate::bit::ops::bitwise_binary_op_lhs_owned; use crate::bit::ops::bitwise_unary_op; use crate::bit::ops::bitwise_unary_op_copy; use crate::bit::select::bit_select; -use crate::buffer; /// An immutable bitset stored as a packed byte buffer. #[derive(Debug, Clone, Eq)] @@ -147,8 +147,14 @@ impl BitBuffer { /// Create a new `BoolBuffer` of length `len` where all bits are set (true). #[inline] pub fn new_set(len: usize) -> Self { + Self::new_set_in(len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a set bit buffer with the provided allocator. + #[inline] + pub fn new_set_in(len: usize, allocator: BufferAllocatorRef) -> Self { let words = len.div_ceil(8); - let buffer = buffer![0xFF; words]; + let buffer = Buffer::full_in(0xFF, words, allocator); Self { buffer, @@ -160,8 +166,14 @@ impl BitBuffer { /// Create a new `BoolBuffer` of length `len` where all bits are unset (false). #[inline] pub fn new_unset(len: usize) -> Self { + Self::new_unset_in(len, BufferAllocatorRef::statically_allocated()) + } + + /// Create an unset bit buffer with the provided allocator. + #[inline] + pub fn new_unset_in(len: usize, allocator: BufferAllocatorRef) -> Self { let words = len.div_ceil(8); - let buffer = Buffer::zeroed(words); + let buffer = Buffer::zeroed_in(words, allocator); Self { buffer, @@ -184,10 +196,16 @@ impl BitBuffer { /// Create a new `BitBuffer` of length `len` where all bits are set to `value`. #[inline] pub fn full(value: bool, len: usize) -> Self { + Self::full_in(value, len, BufferAllocatorRef::statically_allocated()) + } + + /// Create a full bit buffer with the provided allocator. + #[inline] + pub fn full_in(value: bool, len: usize, allocator: BufferAllocatorRef) -> Self { if value { - Self::new_set(len) + Self::new_set_in(len, allocator) } else { - Self::new_unset(len) + Self::new_unset_in(len, allocator) } } @@ -213,6 +231,16 @@ impl BitBuffer { BitBufferMut::collect_bool(len, f).freeze() } + /// Collects predicate results with the provided allocator. + #[inline] + pub fn collect_bool_in bool>( + len: usize, + f: F, + allocator: BufferAllocatorRef, + ) -> Self { + BitBufferMut::collect_bool_in(len, f, allocator).freeze() + } + /// Like [`Self::collect_bool`], but compiles the packing loop — with `f` inside it — once /// per CPU feature level (AVX-512BW/AVX2/baseline) and selects a clone by runtime feature /// detection. @@ -226,7 +254,17 @@ impl BitBuffer { /// [`collect_bool_words_multiversioned`](crate::bit::collect_bool_words_multiversioned). #[inline] pub fn collect_bool_multiversioned bool>(len: usize, f: F) -> Self { - BitBufferMut::collect_bool_multiversioned(len, f).freeze() + Self::collect_bool_multiversioned_in(len, f, BufferAllocatorRef::statically_allocated()) + } + + /// Collects multiversioned predicate results with the provided allocator. + #[inline] + pub fn collect_bool_multiversioned_in bool>( + len: usize, + f: F, + allocator: BufferAllocatorRef, + ) -> Self { + BitBufferMut::collect_bool_multiversioned_in(len, f, allocator).freeze() } /// Maps over each bit in this buffer, calling `f(index, bit_value)` and collecting results. diff --git a/vortex-buffer/src/bit/buf_mut.rs b/vortex-buffer/src/bit/buf_mut.rs index c0c6338378d..2954c078ee9 100644 --- a/vortex-buffer/src/bit/buf_mut.rs +++ b/vortex-buffer/src/bit/buf_mut.rs @@ -231,7 +231,17 @@ impl BitBufferMut { /// ideally with a benchmark. #[inline] pub fn collect_bool bool>(len: usize, f: F) -> Self { - Self::collect_words(len, |words| collect_bool_words(words, len, f)) + Self::collect_bool_in(len, f, BufferAllocatorRef::statically_allocated()) + } + + /// Collects predicate results with the provided allocator. + #[inline] + pub fn collect_bool_in bool>( + len: usize, + f: F, + allocator: BufferAllocatorRef, + ) -> Self { + Self::collect_words_in(len, allocator, |words| collect_bool_words(words, len, f)) } /// Like [`Self::collect_bool`], but compiles the packing loop — with `f` inside it — once @@ -247,17 +257,29 @@ impl BitBufferMut { /// [`collect_bool_words_multiversioned`]. #[inline] pub fn collect_bool_multiversioned bool>(len: usize, f: F) -> Self { - Self::collect_words(len, |words| { + Self::collect_bool_multiversioned_in(len, f, BufferAllocatorRef::statically_allocated()) + } + + /// Collects multiversioned predicate results with the provided allocator. + #[inline] + pub fn collect_bool_multiversioned_in bool>( + len: usize, + f: F, + allocator: BufferAllocatorRef, + ) -> Self { + Self::collect_words_in(len, allocator, |words| { collect_bool_words_multiversioned(words, len, f) }) } - /// Allocate a zero-copy word buffer for `len` bits, let `fill` populate it, and wrap it as a - /// `BitBufferMut`. #[inline] - fn collect_words(len: usize, fill: impl FnOnce(&mut [u64])) -> Self { + fn collect_words_in( + len: usize, + allocator: BufferAllocatorRef, + fill: impl FnOnce(&mut [u64]), + ) -> Self { let num_words = len.div_ceil(64); - let mut buffer: BufferMut = BufferMut::with_capacity(num_words); + let mut buffer = BufferMut::::with_capacity_in(num_words, allocator); // SAFETY: `fill` (a `collect_bool_words` variant) writes every word in `0..num_words` // below before any read; `u64` has no invalid bit patterns and the assignments inside // `collect_bool_words` are pure writes. From 2b0d9f06af1d6e70f4f4debec2001eb68c305485 Mon Sep 17 00:00:00 2001 From: Nicholas Gates Date: Fri, 28 Aug 2026 17:00:04 -0400 Subject: [PATCH 2/2] fix(array): pass allocator to AVX2 take tests Signed-off-by: Nicholas Gates --- vortex-array/src/arrays/fixed_width/take/avx2/tests.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs b/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs index 2ec97bcc38c..57903fb1ddf 100644 --- a/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs @@ -7,6 +7,7 @@ use std::panic::RefUnwindSafe; use std::panic::catch_unwind; use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; use super::super::FixedWidthTakeValue; use super::take_avx2; @@ -22,7 +23,7 @@ fn take_avx2_if_supported( // SAFETY: AVX2 support was detected above, and `FixedWidthTakeValue` guarantees that every // byte in the values is initialized. - Some(unsafe { take_avx2(values, indices) }) + Some(unsafe { take_avx2(values, indices, BufferAllocatorRef::statically_allocated()) }) } fn assert_avx2_take_panics(values: &[V], indices: &[I], expected: &str) @@ -36,7 +37,9 @@ where // SAFETY: AVX2 support was detected above, and `FixedWidthTakeValue` guarantees that every // byte in the values is initialized. - let result = catch_unwind(|| unsafe { take_avx2(values, indices) }); + let result = catch_unwind(|| unsafe { + take_avx2(values, indices, BufferAllocatorRef::statically_allocated()) + }); let Err(payload) = result else { panic!("take should panic for an invalid index"); };