From bf3aea006bd8e81ec84eea734a2767e9d3102ffa Mon Sep 17 00:00:00 2001 From: blaginin Date: Thu, 15 Oct 2026 14:45:09 +0100 Subject: [PATCH 1/3] Preserve contiguous index ranges in BitPacked take Signed-off-by: blaginin --- .../fastlanes/benches/bitpacking_take.rs | 133 ++++++++ .../fastlanes/src/bitpacking/compute/take.rs | 321 ++++++++++++++++++ 2 files changed, 454 insertions(+) diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index 9292143e406..ab15fed0730 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -17,7 +17,12 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray as _; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -406,3 +411,131 @@ fn patched_take_10k_adversarial(bencher: Bencher) { .unwrap() }) } + +// These are the element runs emitted by a sparse FSL gather. The baseline materializes those +// runs inside the measured operation, reproducing the previous BitPacked take path. +fn embedding_runs_fixture(dimension: usize, ptype: PType, bit_width: u8) -> (ArrayRef, ArrayRef) { + const ROWS: usize = 4096; + const SELECTED: usize = 16; + let mut ctx = SESSION.create_execution_ctx(); + let values = fixture(ROWS * dimension, usize::from(bit_width)) + .into_array() + .cast(ptype.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx) + .unwrap() + .into_array(); + let mut rng = StdRng::seed_from_u64(0); + let starts = PrimitiveArray::from_iter( + (0..SELECTED).map(|_| (rng.random_range(0..ROWS) * dimension) as u64), + ) + .into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(dimension as u64, SELECTED).into_array(), + ConstantArray::new(1u64, SELECTED).into_array(), + SELECTED * dimension, + ) + .unwrap() + .into_array(); + (packed, indices) +} + +#[divan::bench(args = [768, 1024, 1536])] +fn embedding_runs(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U32, 16); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [768, 1024, 1536])] +fn embedding_runs_expanded_baseline(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U32, 16); + bench_piecewise_take(bencher, packed, indices, true); +} + +#[divan::bench(args = [512, 1024, 2048, 4096])] +fn embedding_code_runs(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U8, 4); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [512, 1024, 2048, 4096])] +fn embedding_code_runs_expanded_baseline(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U8, 4); + bench_piecewise_take(bencher, packed, indices, true); +} + +// Short runs share chunks; long runs visit distant chunks. Neither crosses the full-decode +// threshold, including the largest case (16 * 1024 selected out of 128 * 1024 values). +fn fragmented_runs_fixture(run_length: usize) -> (ArrayRef, ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + let values = fixture(128 * CHUNK_SIZE, 16).into_array(); + let packed = BitPackedData::encode(&values, 16, &mut ctx) + .unwrap() + .into_array(); + let (num_runs, stride) = if run_length < 64 { + (128, run_length + 1) + } else { + (16, 4 * CHUNK_SIZE) + }; + let starts = + PrimitiveArray::from_iter((0..num_runs).map(|run| (run * stride) as u64)).into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(run_length as u64, num_runs).into_array(), + ConstantArray::new(1u64, num_runs).into_array(), + num_runs * run_length, + ) + .unwrap() + .into_array(); + (packed, indices) +} + +#[divan::bench(args = [1, 2, 7, 8, 64, 512, 1024])] +fn fragmented_runs(bencher: Bencher, run_length: usize) { + let (packed, indices) = fragmented_runs_fixture(run_length); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [1, 2, 7, 8, 64, 512, 1024])] +fn fragmented_runs_expanded_baseline(bencher: Bencher, run_length: usize) { + let (packed, indices) = fragmented_runs_fixture(run_length); + bench_piecewise_take(bencher, packed, indices, true); +} + +fn bench_piecewise_take( + bencher: Bencher, + packed: ArrayRef, + indices: ArrayRef, + expanded_baseline: bool, +) { + let mut ctx = SESSION.create_execution_ctx(); + let canonical = packed.clone().execute::(&mut ctx).unwrap(); + assert_arrays_eq!( + packed.take(indices.clone()).unwrap(), + canonical.take(indices.clone()).unwrap(), + &mut ctx + ); + bencher + .counter(ItemsCount::new(indices.len())) + .with_inputs(|| (indices.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(indices, ctx)| { + let indices = if expanded_baseline { + indices + .clone() + .execute::(ctx) + .unwrap() + .into_array() + } else { + indices.clone() + }; + packed + .take(indices) + .unwrap() + .execute::(ctx) + .unwrap() + }); +} diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 1bb778212c3..fcbdd90f615 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -3,14 +3,21 @@ use std::mem; use std::mem::MaybeUninit; +use std::ops::Range; use fastlanes::BitPacking; +use itertools::Itertools; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::PiecewiseSequence; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::TakeExecute; +use vortex_array::arrays::piecewise_sequence::array::PiecewiseSequenceArraySlotsExt; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::IntegerPType; use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; @@ -21,11 +28,15 @@ use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; use super::chunked_indices; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitWidthsView; +use crate::FL_CHUNK_SIZE; use crate::bitpack_decompress; // TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. @@ -49,6 +60,23 @@ impl TakeExecute for BitPacked { return prim.into_array().take(indices.clone()).map(Some); } + if let Some(ranges) = contiguous_ranges(indices, array.len(), ctx)? { + let taken_validity = array.validity()?.take(indices)?; + let ptype = array.dtype().as_ptype(); + let taken = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| { + take_ranges::(array, bit_width, &ranges, indices.len(), taken_validity) + }); + let taken = taken.reinterpret_cast(ptype); + let taken = if let Some(patches) = array.patches() + && let Some(patches) = patches.take(indices, ctx)? + { + taken.patch(&patches, ctx)? + } else { + taken + }; + return Ok(Some(taken.into_array())); + } + // NOTE: we use the unsigned PType because all values in the BitPackedArray must // be non-negative (pre-condition of creating the BitPackedArray). let ptype: PType = PType::try_from(array.dtype())?; @@ -74,6 +102,129 @@ impl TakeExecute for BitPacked { } } +fn contiguous_ranges( + indices: &ArrayRef, + array_len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult>>> { + let Some(indices) = indices.as_opt::() else { + return Ok(None); + }; + // The primitive path batches short runs that share a packed chunk. + if let Some(length) = indices + .lengths() + .as_opt::() + .and_then(|length| length.scalar().as_primitive().as_::()) + && length != 0 + && length < UNPACK_CHUNK_THRESHOLD as u64 + { + return Ok(None); + } + + if indices + .multipliers() + .as_opt::() + .and_then(|multiplier| multiplier.scalar().as_primitive().as_::()) + != Some(1) + { + return Ok(None); + } + + let lengths = indices + .lengths() + .cast(PType::U64.into())? + .execute::(ctx)?; + if lengths + .as_slice::() + .iter() + .any(|&length| length != 0 && length < UNPACK_CHUNK_THRESHOLD as u64) + { + return Ok(None); + } + let starts = indices + .starts() + .cast(PType::U64.into())? + .execute::(ctx)?; + let mut ranges = Vec::with_capacity(starts.len()); + let mut output_len = 0usize; + for (&start, &length) in starts + .as_slice::() + .iter() + .zip_eq(lengths.as_slice::()) + { + if length == 0 { + continue; + } + let start = usize::try_from(start)?; + let length = usize::try_from(length)?; + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray range overflows usize"))?; + vortex_ensure!( + end <= array_len, + "PiecewiseSequenceArray range exceeds array length" + ); + output_len = output_len + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + ranges.push(start..end); + } + vortex_ensure_eq!(output_len, indices.len()); + Ok(Some(ranges)) +} + +fn take_ranges( + array: ArrayView<'_, BitPacked>, + bit_width: u8, + ranges: &[Range], + output_len: usize, + validity: Validity, +) -> PrimitiveArray { + let bit_width = usize::from(bit_width); + let packed = array.packed_slice::(); + let chunk_len = 128 * bit_width / size_of::(); + let offset = usize::from(array.offset()); + let mut output = BufferMut::::with_capacity(output_len); + let mut unpacked = [T::zero(); FL_CHUNK_SIZE]; + let mut unpacked_chunk = None; + + // Vector gathers select runs of components. Keep those runs intact so neither the index + // buffer nor the chunk-grouping work scales with the number of selected components. + for range in ranges { + let mut start = range.start + offset; + let end = range.end + offset; + while start < end { + let chunk_idx = start / FL_CHUNK_SIZE; + let within_chunk = start % FL_CHUNK_SIZE; + let length = (end - start).min(FL_CHUNK_SIZE - within_chunk); + let packed_chunk = &packed[chunk_idx * chunk_len..][..chunk_len]; + if unpacked_chunk != Some(chunk_idx) && length < UNPACK_CHUNK_THRESHOLD { + for index in within_chunk..within_chunk + length { + // SAFETY: the validated range lies within this array; the packed chunk has + // exactly the words required by its bit width, and index is below 1024. + output.push(unsafe { + bitpack_decompress::unpack_single_primitive::( + packed_chunk, + bit_width, + index, + ) + }); + } + } else { + if unpacked_chunk != Some(chunk_idx) { + // SAFETY: BitPacked validates its bit width and buffer size. The input holds + // one packed chunk and the initialized output has exactly 1024 elements. + unsafe { BitPacking::unchecked_unpack(bit_width, packed_chunk, &mut unpacked) }; + unpacked_chunk = Some(chunk_idx); + } + output.extend_from_slice(&unpacked[within_chunk..within_chunk + length]); + } + start += length; + } + } + PrimitiveArray::new(output, validity) +} + fn take_primitive( array: ArrayView<'_, BitPacked>, bit_width: u8, @@ -171,14 +322,25 @@ mod test { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::FixedSizeListArray; + use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::dict::TakeExecute; + use vortex_array::arrays::slice::SliceKernel; use vortex_array::assert_arrays_eq; + use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_session::VortexSession; + use crate::BitPacked; use crate::BitPackedArray; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; @@ -318,6 +480,165 @@ mod test { assert_eq!(taken_primitive_prim.invalid_count(&mut ctx).unwrap(), 1); } + #[rstest] + #[case(0, 0, false, false, u64::MAX, PType::U64)] + #[case(7, 0, false, false, u64::MAX, PType::U8)] + #[case(7, 137, true, true, 0, PType::I32)] + #[case(15, 1023, false, true, 0, PType::U16)] + #[case(4, 137, false, true, 0, PType::U8)] + #[case(4, 1023, true, true, 0, PType::U8)] + #[case(63, 1023, false, true, 0, PType::U64)] + #[case(62, 137, true, true, 0, PType::I64)] + fn test_take_piecewise_ranges( + #[case] bit_width: u8, + #[case] slice_offset: usize, + #[case] patched: bool, + #[case] nullable: bool, + #[case] empty_start: u64, + #[case] ptype: PType, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source_len: usize = 32755; + let values = PrimitiveArray::from_option_iter((0..source_len).map(|i| { + if nullable && i.is_multiple_of(29) { + None + } else if patched && i.is_multiple_of(31) { + Some(1u64 << bit_width) + } else { + Some((i as u64).wrapping_mul(0x0123_4567_89ab_cdef) & ((1u64 << bit_width) - 1)) + } + })) + .into_array() + .cast(DType::Primitive(ptype, Nullability::Nullable))? + .execute::(&mut ctx)? + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx)?; + let packed = + BitPacked::slice(packed.as_view(), slice_offset..source_len, &mut ctx)?.unwrap(); + let values = values.slice(slice_offset..source_len)?; + let indices = PiecewiseSequenceArray::try_new( + buffer![ + empty_start, + 17, + 991, + 1000, + 3072, + 12000, + 17000, + 17, + (source_len - slice_offset - 27) as u64 + ] + .into_array(), + buffer![0u16, 8, 9, 27, 768, 1024, 1536, 27, 27].into_array(), + ConstantArray::new(1u8, 9).into_array(), + 3426, + )? + .into_array(); + let expanded = indices + .clone() + .execute::(&mut ctx)? + .into_array(); + let actual = BitPacked::take(packed.as_::(), &indices, &mut ctx)?.unwrap(); + assert_arrays_eq!(actual, values.take(expanded)?, &mut ctx); + Ok(()) + } + + #[test] + fn test_take_piecewise_fallbacks() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0u32..32768).map(|i| i % 127)).into_array(); + let packed = BitPackedData::encode(&values, 7, &mut ctx)?; + for (length, multiplier, constant) in + [(7u8, 1u8, true), (8, 0, true), (8, 2, true), (8, 1, false)] + { + let multipliers = if constant { + ConstantArray::new(multiplier, 2).into_array() + } else { + buffer![multiplier, multiplier].into_array() + }; + for (lengths, len) in [ + (buffer![length, 16].into_array(), usize::from(length) + 16), + ( + ConstantArray::new(length, 2).into_array(), + usize::from(length) * 2, + ), + ] { + let indices = PiecewiseSequenceArray::try_new( + buffer![1017u16, 1023].into_array(), + lengths, + multipliers.clone(), + len, + )? + .into_array(); + let expanded = indices + .clone() + .execute::(&mut ctx)? + .into_array(); + assert_arrays_eq!(packed.take(indices)?, values.take(expanded)?, &mut ctx); + } + } + Ok(()) + } + + #[rstest] + fn test_take_piecewise_fsl_integration( + #[values(false, true)] blocked: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_option_iter( + (0usize..65536).map(|i| (!i.is_multiple_of(29)).then_some((i % 128) as u8)), + ) + .into_array(); + let packed = if blocked { + let widths = (0..64) + .map(|i| if i % 2 == 0 { 4 } else { 7 }) + .collect::>(); + BitPackedData::encode_blocked(&values, &widths, &mut ctx)? + } else { + BitPackedData::encode(&values, 4, &mut ctx)? + }; + let validity = Validity::from_iter((0usize..1024).map(|i| !i.is_multiple_of(11))); + let expected = FixedSizeListArray::new(values, 64, validity.clone(), 1024); + let array = FixedSizeListArray::new(packed.into_array(), 64, validity, 1024); + let indices = + PrimitiveArray::from_option_iter([Some(1000u32), None, Some(0), Some(16), Some(16)]) + .into_array(); + assert_arrays_eq!( + array.take(indices.clone())?, + expected.take(indices)?, + &mut ctx + ); + Ok(()) + } + + #[rstest] + #[case(32760, 16, 16)] + #[case(1, 16, 15)] + #[case(u64::MAX, 8, 8)] + fn test_take_piecewise_invalid_range( + #[case] start: u64, + #[case] length: u64, + #[case] declared_len: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0..32768).map(|i| i as u32 % 127)).into_array(); + let packed = BitPackedData::encode(&values, 7, &mut ctx)?; + let indices = PiecewiseSequenceArray::try_new( + buffer![start].into_array(), + buffer![length].into_array(), + ConstantArray::new(1u64, 1).into_array(), + declared_len, + )? + .into_array(); + assert!( + packed + .take(indices)? + .execute::(&mut ctx) + .is_err() + ); + Ok(()) + } + fn bp(array: vortex_array::ArrayRef, bit_width: u8) -> BitPackedArray { BitPackedData::encode(&array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() } From 73192344ef8645b4dba38a8894776b68fb6d99cb Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 16 Oct 2026 14:37:38 +0100 Subject: [PATCH 2/3] Use sparse decode thresholds for BitPacked ranges Signed-off-by: blaginin --- .../fastlanes/benches/bitpacking_take.rs | 47 +++++++++ .../fastlanes/src/bitpacking/compute/take.rs | 97 +++++++++++++------ 2 files changed, 116 insertions(+), 28 deletions(-) diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index ab15fed0730..ba2cac32499 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -22,6 +22,7 @@ use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; @@ -506,6 +507,52 @@ fn fragmented_runs_expanded_baseline(bencher: Bencher, run_length: usize) { bench_piecewise_take(bencher, packed, indices, true); } +// Isolated runs cannot amortize a whole-chunk decode across neighboring runs. Offsets near +// a chunk boundary also exercise sparse fragments of runs above the full-decode threshold. +#[divan::bench( + types = [u8, u16, u32, u64], + args = [ + (8, 0), (16, 0), (32, 0), (64, 0), (128, 0), (160, 0), (161, 0), + (256, 0), (512, 0), (161, 944), (256, 1020), (256, 896), + ], +)] +fn isolated_runs(bencher: Bencher, (run_length, offset): (usize, usize)) { + let bit_width = match T::PTYPE { + PType::U8 => 4, + PType::U16 => 9, + PType::U32 => 16, + PType::U64 => 33, + _ => unreachable!(), + }; + let mut ctx = SESSION.create_execution_ctx(); + let mut rng = StdRng::seed_from_u64(42); + let values = PrimitiveArray::from_iter( + (0..128 * CHUNK_SIZE).map(|_| rng.random_range(0..(1u64 << bit_width))), + ) + .into_array() + .cast(T::PTYPE.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx) + .unwrap() + .into_array(); + let num_runs = 16; + let starts = + PrimitiveArray::from_iter((0..num_runs).map(|run| (run * 4 * CHUNK_SIZE + offset) as u64)) + .into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(run_length as u64, num_runs).into_array(), + ConstantArray::new(1u64, num_runs).into_array(), + num_runs * run_length, + ) + .unwrap() + .into_array(); + bench_piecewise_take(bencher, packed, indices, false); +} + fn bench_piecewise_take( bencher: Bencher, packed: ArrayRef, diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 0ebbff2af34..0321ede2462 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -38,10 +38,8 @@ use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitWidthsView; use crate::FL_CHUNK_SIZE; -use crate::bitpack_decompress; const FULL_ARRAY_DECODE_RATIO: usize = 8; -const MIN_CONTIGUOUS_RANGE_LEN: usize = 8; impl TakeExecute for BitPacked { fn take( @@ -58,7 +56,9 @@ impl TakeExecute for BitPacked { return prim.into_array().take(indices.clone()).map(Some); } - if let Some(ranges) = contiguous_ranges(indices, array.len(), ctx)? { + if let Some(ranges) = + contiguous_ranges(indices, array.len(), array.dtype().as_ptype(), ctx)? + { let taken_validity = array.validity()?.take(indices)?; let ptype = array.dtype().as_ptype(); let taken = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| { @@ -103,18 +103,23 @@ impl TakeExecute for BitPacked { fn contiguous_ranges( indices: &ArrayRef, array_len: usize, + ptype: PType, ctx: &mut ExecutionCtx, ) -> VortexResult>>> { let Some(indices) = indices.as_opt::() else { return Ok(None); }; - // The primitive path batches short runs that share a packed chunk. + let sparse_threshold = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| { + unpack_chunk_threshold::() as u64 + }); + // The primitive path batches short runs that share a packed chunk and selectively decodes + // isolated short runs. Reserve the range path for runs long enough to unpack a whole chunk. if let Some(length) = indices .lengths() .as_opt::() .and_then(|length| length.scalar().as_primitive().as_::()) && length != 0 - && length < MIN_CONTIGUOUS_RANGE_LEN as u64 + && length <= sparse_threshold { return Ok(None); } @@ -135,7 +140,7 @@ fn contiguous_ranges( if lengths .as_slice::() .iter() - .any(|&length| length != 0 && length < MIN_CONTIGUOUS_RANGE_LEN as u64) + .any(|&length| length != 0 && length <= sparse_threshold) { return Ok(None); } @@ -185,6 +190,8 @@ fn take_ranges( let mut output = BufferMut::::with_capacity(output_len); let mut unpacked = [T::zero(); FL_CHUNK_SIZE]; let mut unpacked_chunk = None; + // u64 has the largest sparse decode threshold. + let mut sparse_indices = [0usize; unpack_chunk_threshold::()]; // Vector gathers select runs of components. Keep those runs intact so neither the index // buffer nor the chunk-grouping work scales with the number of selected components. @@ -196,18 +203,12 @@ fn take_ranges( let within_chunk = start % FL_CHUNK_SIZE; let length = (end - start).min(FL_CHUNK_SIZE - within_chunk); let packed_chunk = &packed[chunk_idx * chunk_len..][..chunk_len]; - if unpacked_chunk != Some(chunk_idx) && length < MIN_CONTIGUOUS_RANGE_LEN { - for index in within_chunk..within_chunk + length { - // SAFETY: the validated range lies within this array; the packed chunk has - // exactly the words required by its bit width, and index is below 1024. - output.push(unsafe { - bitpack_decompress::unpack_single_primitive::( - packed_chunk, - bit_width, - index, - ) - }); + if unpacked_chunk != Some(chunk_idx) && length <= unpack_chunk_threshold::() { + let indices = &mut sparse_indices[..length]; + for (index, position) in indices.iter_mut().zip(within_chunk..) { + *index = position; } + unpack_indices_into(&mut output, bit_width, packed_chunk, indices); } else { if unpacked_chunk != Some(chunk_idx) { // SAFETY: BitPacked validates its bit width and buffer size. The input holds @@ -549,7 +550,7 @@ mod tests { #[case] ptype: PType, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let source_len: usize = 32755; + let source_len: usize = 65523; let values = PrimitiveArray::from_option_iter((0..source_len).map(|i| { if nullable && i.is_multiple_of(29) { None @@ -577,12 +578,12 @@ mod tests { 12000, 17000, 17, - (source_len - slice_offset - 27) as u64 + (source_len - slice_offset - 257) as u64 ] .into_array(), - buffer![0u16, 8, 9, 27, 768, 1024, 1536, 27, 27].into_array(), + buffer![0u16, 256, 257, 511, 768, 1024, 1536, 256, 257].into_array(), ConstantArray::new(1u8, 9).into_array(), - 3426, + 4865, )? .into_array(); let expanded = indices @@ -599,16 +600,19 @@ mod tests { let mut ctx = SESSION.create_execution_ctx(); let values = PrimitiveArray::from_iter((0u32..32768).map(|i| i % 127)).into_array(); let packed = BitPackedData::encode(&values, 7, &mut ctx)?; - for (length, multiplier, constant) in - [(7u8, 1u8, true), (8, 0, true), (8, 2, true), (8, 1, false)] - { + for (length, multiplier, constant) in [ + (7u8, 1u8, true), + (65, 0, true), + (65, 2, true), + (65, 1, false), + ] { let multipliers = if constant { ConstantArray::new(multiplier, 2).into_array() } else { buffer![multiplier, multiplier].into_array() }; for (lengths, len) in [ - (buffer![length, 16].into_array(), usize::from(length) + 16), + (buffer![length, 128].into_array(), usize::from(length) + 128), ( ConstantArray::new(length, 2).into_array(), usize::from(length) * 2, @@ -631,6 +635,43 @@ mod tests { Ok(()) } + #[rstest] + #[case(PType::U8, 4, 16)] + #[case(PType::U16, 9, 32)] + #[case(PType::U32, 16, 64)] + #[case(PType::U64, 33, 160)] + fn test_take_piecewise_sparse_threshold( + #[case] ptype: PType, + #[case] bit_width: u8, + #[case] threshold: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter( + (0u64..8192).map(|i| i.wrapping_mul(0x0123_4567_89ab_cdef) & ((1u64 << bit_width) - 1)), + ) + .into_array() + .cast(ptype.into())? + .execute::(&mut ctx)? + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx)?; + for length in [threshold, threshold + 1, 2 * threshold + 1] { + // A run split across two chunks can require sparse extraction on both sides. + let indices = PiecewiseSequenceArray::try_new( + buffer![0u64, (1024 - threshold / 2) as u64, 4096].into_array(), + ConstantArray::new(length as u64, 3).into_array(), + ConstantArray::new(1u64, 3).into_array(), + length * 3, + )? + .into_array(); + let expanded = indices + .clone() + .execute::(&mut ctx)? + .into_array(); + assert_arrays_eq!(packed.take(indices)?, values.take(expanded)?, &mut ctx); + } + Ok(()) + } + #[rstest] fn test_take_piecewise_fsl_integration( #[values(false, true)] blocked: bool, @@ -663,9 +704,9 @@ mod tests { } #[rstest] - #[case(32760, 16, 16)] - #[case(1, 16, 15)] - #[case(u64::MAX, 8, 8)] + #[case(32700, 128, 128)] + #[case(1, 128, 127)] + #[case(u64::MAX, 128, 128)] fn test_take_piecewise_invalid_range( #[case] start: u64, #[case] length: u64, From c799ef9ccc55cc07cb20b21792350ecde1ceb4cf Mon Sep 17 00:00:00 2001 From: blaginin Date: Fri, 16 Oct 2026 15:03:13 +0100 Subject: [PATCH 3/3] bench(fastlanes): cover contiguous and isolated BitPacked takes Signed-off-by: blaginin --- .../fastlanes/benches/bitpacking_take.rs | 180 ++++++++++++++++++ 1 file changed, 180 insertions(+) diff --git a/encodings/fastlanes/benches/bitpacking_take.rs b/encodings/fastlanes/benches/bitpacking_take.rs index 9292143e406..ba2cac32499 100644 --- a/encodings/fastlanes/benches/bitpacking_take.rs +++ b/encodings/fastlanes/benches/bitpacking_take.rs @@ -17,7 +17,13 @@ use vortex_array::ArrayRef; use vortex_array::IntoArray as _; use vortex_array::RecursiveCanonical; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PiecewiseSequenceArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -406,3 +412,177 @@ fn patched_take_10k_adversarial(bencher: Bencher) { .unwrap() }) } + +// These are the element runs emitted by a sparse FSL gather. The baseline materializes those +// runs inside the measured operation, reproducing the previous BitPacked take path. +fn embedding_runs_fixture(dimension: usize, ptype: PType, bit_width: u8) -> (ArrayRef, ArrayRef) { + const ROWS: usize = 4096; + const SELECTED: usize = 16; + let mut ctx = SESSION.create_execution_ctx(); + let values = fixture(ROWS * dimension, usize::from(bit_width)) + .into_array() + .cast(ptype.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx) + .unwrap() + .into_array(); + let mut rng = StdRng::seed_from_u64(0); + let starts = PrimitiveArray::from_iter( + (0..SELECTED).map(|_| (rng.random_range(0..ROWS) * dimension) as u64), + ) + .into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(dimension as u64, SELECTED).into_array(), + ConstantArray::new(1u64, SELECTED).into_array(), + SELECTED * dimension, + ) + .unwrap() + .into_array(); + (packed, indices) +} + +#[divan::bench(args = [768, 1024, 1536])] +fn embedding_runs(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U32, 16); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [768, 1024, 1536])] +fn embedding_runs_expanded_baseline(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U32, 16); + bench_piecewise_take(bencher, packed, indices, true); +} + +#[divan::bench(args = [512, 1024, 2048, 4096])] +fn embedding_code_runs(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U8, 4); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [512, 1024, 2048, 4096])] +fn embedding_code_runs_expanded_baseline(bencher: Bencher, dimension: usize) { + let (packed, indices) = embedding_runs_fixture(dimension, PType::U8, 4); + bench_piecewise_take(bencher, packed, indices, true); +} + +// Short runs share chunks; long runs visit distant chunks. Neither crosses the full-decode +// threshold, including the largest case (16 * 1024 selected out of 128 * 1024 values). +fn fragmented_runs_fixture(run_length: usize) -> (ArrayRef, ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + let values = fixture(128 * CHUNK_SIZE, 16).into_array(); + let packed = BitPackedData::encode(&values, 16, &mut ctx) + .unwrap() + .into_array(); + let (num_runs, stride) = if run_length < 64 { + (128, run_length + 1) + } else { + (16, 4 * CHUNK_SIZE) + }; + let starts = + PrimitiveArray::from_iter((0..num_runs).map(|run| (run * stride) as u64)).into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(run_length as u64, num_runs).into_array(), + ConstantArray::new(1u64, num_runs).into_array(), + num_runs * run_length, + ) + .unwrap() + .into_array(); + (packed, indices) +} + +#[divan::bench(args = [1, 2, 7, 8, 64, 512, 1024])] +fn fragmented_runs(bencher: Bencher, run_length: usize) { + let (packed, indices) = fragmented_runs_fixture(run_length); + bench_piecewise_take(bencher, packed, indices, false); +} + +#[divan::bench(args = [1, 2, 7, 8, 64, 512, 1024])] +fn fragmented_runs_expanded_baseline(bencher: Bencher, run_length: usize) { + let (packed, indices) = fragmented_runs_fixture(run_length); + bench_piecewise_take(bencher, packed, indices, true); +} + +// Isolated runs cannot amortize a whole-chunk decode across neighboring runs. Offsets near +// a chunk boundary also exercise sparse fragments of runs above the full-decode threshold. +#[divan::bench( + types = [u8, u16, u32, u64], + args = [ + (8, 0), (16, 0), (32, 0), (64, 0), (128, 0), (160, 0), (161, 0), + (256, 0), (512, 0), (161, 944), (256, 1020), (256, 896), + ], +)] +fn isolated_runs(bencher: Bencher, (run_length, offset): (usize, usize)) { + let bit_width = match T::PTYPE { + PType::U8 => 4, + PType::U16 => 9, + PType::U32 => 16, + PType::U64 => 33, + _ => unreachable!(), + }; + let mut ctx = SESSION.create_execution_ctx(); + let mut rng = StdRng::seed_from_u64(42); + let values = PrimitiveArray::from_iter( + (0..128 * CHUNK_SIZE).map(|_| rng.random_range(0..(1u64 << bit_width))), + ) + .into_array() + .cast(T::PTYPE.into()) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array(); + let packed = BitPackedData::encode(&values, bit_width, &mut ctx) + .unwrap() + .into_array(); + let num_runs = 16; + let starts = + PrimitiveArray::from_iter((0..num_runs).map(|run| (run * 4 * CHUNK_SIZE + offset) as u64)) + .into_array(); + let indices = PiecewiseSequenceArray::try_new( + starts, + ConstantArray::new(run_length as u64, num_runs).into_array(), + ConstantArray::new(1u64, num_runs).into_array(), + num_runs * run_length, + ) + .unwrap() + .into_array(); + bench_piecewise_take(bencher, packed, indices, false); +} + +fn bench_piecewise_take( + bencher: Bencher, + packed: ArrayRef, + indices: ArrayRef, + expanded_baseline: bool, +) { + let mut ctx = SESSION.create_execution_ctx(); + let canonical = packed.clone().execute::(&mut ctx).unwrap(); + assert_arrays_eq!( + packed.take(indices.clone()).unwrap(), + canonical.take(indices.clone()).unwrap(), + &mut ctx + ); + bencher + .counter(ItemsCount::new(indices.len())) + .with_inputs(|| (indices.clone(), SESSION.create_execution_ctx())) + .bench_refs(|(indices, ctx)| { + let indices = if expanded_baseline { + indices + .clone() + .execute::(ctx) + .unwrap() + .into_array() + } else { + indices.clone() + }; + packed + .take(indices) + .unwrap() + .execute::(ctx) + .unwrap() + }); +}