diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index bbcf0b9ba56..0321ede2462 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -2,14 +2,21 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors 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; @@ -20,6 +27,9 @@ 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 super::unpack_chunk_threshold; @@ -27,6 +37,7 @@ use super::unpack_indices_into; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitWidthsView; +use crate::FL_CHUNK_SIZE; const FULL_ARRAY_DECODE_RATIO: usize = 8; @@ -45,6 +56,25 @@ impl TakeExecute for BitPacked { return prim.into_array().take(indices.clone()).map(Some); } + 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| { + 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())?; @@ -70,6 +100,130 @@ 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); + }; + 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 <= sparse_threshold + { + 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 <= sparse_threshold) + { + 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; + // 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. + 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::() { + 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 + // 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, @@ -148,15 +302,25 @@ mod tests { 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; @@ -368,6 +532,205 @@ mod tests { 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 = 65523; + 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 - 257) as u64 + ] + .into_array(), + buffer![0u16, 256, 257, 511, 768, 1024, 1536, 256, 257].into_array(), + ConstantArray::new(1u8, 9).into_array(), + 4865, + )? + .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), + (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, 128].into_array(), usize::from(length) + 128), + ( + 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] + #[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, + ) -> 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(32700, 128, 128)] + #[case(1, 128, 127)] + #[case(u64::MAX, 128, 128)] + 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() }