diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index 0684aea5e6a..e14cd112ec5 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -17,14 +17,13 @@ use vortex_array::match_each_integer_ptype; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::patches::Patches; use vortex_array::scalar::Scalar; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::FL_CHUNK_SIZE; -use crate::unpack_iter::BitPacked as BitPackedUnpack; -use crate::unpack_iter::BitUnpackedChunks; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedUnpack; +use crate::bitpacking::unpack_iter::BitUnpackedChunks; /// Unpacks a bit-packed array into a primitive array. pub fn unpack_array( @@ -124,8 +123,9 @@ where // SAFETY: `decode` writes a value to every slot in this range. let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, len) }; + let widths = array.chunk_layout(ctx)?; let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; - let mut chunks = array.unpacked_chunks::(&mut scratch)?; + let mut chunks = array.unpacked_chunks::(&widths, &mut scratch)?; decode(&mut chunks, uninit_slice, &map); if let Some(patches) = array.patches() { @@ -164,19 +164,28 @@ pub(crate) fn apply_patches_to_uninit_range, index: usize) -> Scalar { - let bit_width = array.bit_width() as usize; +pub fn unpack_single( + array: ArrayView<'_, BitPacked>, + index: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { let ptype = array.dtype().as_ptype(); - // let packed = array.packed().into_primitive()?; let index_in_encoded = index + array.offset() as usize; + let chunk = index_in_encoded / FL_CHUNK_SIZE; + let index_in_chunk = index_in_encoded % FL_CHUNK_SIZE; + let (range, bit_width) = array.chunk_range(chunk, ctx)?; let scalar: Scalar = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |P| { + let packed_chunk = + &array.packed_slice::

()[range.start / size_of::

()..range.end / size_of::

()]; + // SAFETY: `packed_chunk` is exactly one packed block at `bit_width`, and the index is + // within the chunk. unsafe { - unpack_single_primitive::

(array.packed_slice::

(), bit_width, index_in_encoded) - .into() + BitPacking::unchecked_unpack_single(bit_width as usize, packed_chunk, index_in_chunk) } + .into() }); // Cast to fix signedness and nullability - scalar.cast(array.dtype()).vortex_expect("cast failure") + scalar.cast(array.dtype()) } /// # Safety @@ -227,12 +236,13 @@ mod tests { use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_session::VortexSession; use super::*; use crate::BitPackedArray; use crate::BitPackedData; - use crate::bitpack_compress::bitpack_encode; + use crate::bitpacking::bitpack_compress::bitpack_encode; fn encode(array: &PrimitiveArray, bit_width: u8) -> BitPackedArray { bitpack_encode(array, bit_width, None, &mut SESSION.create_execution_ctx()).unwrap() @@ -259,7 +269,7 @@ mod tests { .iter() .enumerate() .for_each(|(i, v)| { - let scalar: u16 = (&unpack_single(compressed.as_view(), i)) + let scalar: u16 = (&unpack_single(compressed.as_view(), i, &mut ctx).unwrap()) .try_into() .unwrap(); assert_eq!(scalar, *v); diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 1b6f02715cc..ac4dce6daec 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -250,7 +250,6 @@ pub(crate) fn materialized_layout(offsets: &ArrayRef) -> VortexResult, @@ -271,7 +269,7 @@ pub struct BitPackedData { impl Display for BitPackedData { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "bit_width: {}, offset: {}", self.bit_width, self.offset) + write!(f, "offset: {}", self.offset) } } @@ -302,17 +300,16 @@ impl BitPackedData { /// * Any patches must have any `array_len` equal to `length` /// * The offsets child must hold `num_chunks + 1` non-nullable `u64` byte boundaries. /// - /// Offset differences must still imply the scalar `bit_width`, and the child must be - /// materialized until the kernels are migrated. The packed buffer must match its span. + /// Adjacent boundaries must differ by a multiple of 128 bytes, implying a width no greater + /// than `ptype`. The final boundary minus the origin must equal the packed buffer size. + /// Compressed offsets are checked at execution time, before unpacking. /// /// Any violation of these preconditions will result in an error. pub fn try_new( packed: BufferHandle, patches: Option, - bit_width: u8, offset: u16, ) -> VortexResult { - vortex_ensure!(bit_width <= 64, "Unsupported bit width {bit_width}"); vortex_ensure!( (offset as usize) < FL_CHUNK_SIZE, "Offset must be less than the full block i.e., {FL_CHUNK_SIZE}, got {offset}" @@ -320,7 +317,6 @@ impl BitPackedData { Ok(Self { offset, - bit_width, packed, patches_data: patches.as_ref().map(PatchesData::from_patches), }) @@ -359,15 +355,10 @@ impl BitPackedData { num_chunks + 1, offsets.len() ); - let widths = materialized_layout(offsets)?.ok_or_else(|| { - vortex_err!("BitPacked chunk layout must be materialized while kernels use bit_width") - })?; - vortex_ensure!( - widths.is_empty() || widths.uniform_width() == Some(self.bit_width), - "Chunk offsets must imply bit_width {}", - self.bit_width - ); - Self::validate_layout(&self.packed, ptype, &widths)?; + // Compressed offsets are checked once materialized, before any unchecked unpacking. + if let Some(widths) = materialized_layout(offsets)? { + Self::validate_layout(&self.packed, ptype, &widths)?; + } Ok(()) } @@ -432,25 +423,20 @@ impl BitPackedData { unsafe { std::slice::from_raw_parts(packed_ptr, packed_len) } } - /// Accessor for bit unpacked chunks - pub fn unpacked_chunks<'a, T: BitPackedIter>( - &'a self, - dtype: &DType, - len: usize, - scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], - ) -> VortexResult> { - assert_eq!( - T::PTYPE, - self.ptype(dtype), - "Requested type doesn't match the array ptype" - ); - BitUnpackedChunks::try_new(self, len, scratch) - } - - /// Bit-width of the packed values + /// The packed FastLanes block of `chunk` as `T` words, along with that chunk's bit width. #[inline] - pub fn bit_width(&self) -> u8 { - self.bit_width + pub(crate) fn packed_chunk( + &self, + widths: &ChunkLayout, + chunk: usize, + ) -> (&[T], usize) { + let bit_width = widths.width(chunk); + let start = widths.byte_offset(chunk) / size_of::(); + let len = chunk_packed_bytes(bit_width) / size_of::(); + ( + &self.packed_slice::()[start..][..len], + bit_width as usize, + ) } #[inline] @@ -480,14 +466,6 @@ impl BitPackedData { .map_err(|a| vortex_err!(InvalidArgument: "Bitpacking can only encode primitive arrays, got {}", a.encoding_id()))?; bitpack_encode(&parray, bit_width, None, ctx) } - - /// Calculate the maximum value that **can** be contained by this array, given its bit-width. - /// - /// Note that this value need not actually be present in the array. - #[inline] - pub fn max_packed_value(&self) -> usize { - (1 << self.bit_width()) - 1 - } } pub trait BitPackedArrayExt: BitPackedArraySlotsExt { @@ -496,11 +474,6 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { BitPackedData::packed(self) } - #[inline] - fn bit_width(&self) -> u8 { - BitPackedData::bit_width(self) - } - /// Materialize and validate offsets once for bulk access, without computing prefix sums. fn chunk_layout(&self, ctx: &mut ExecutionCtx) -> VortexResult { let layout = match materialized_layout(self.chunk_offsets())? { @@ -516,6 +489,71 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { Ok(layout) } + /// Read and validate offsets only when the child is already materialized. + fn materialized_chunk_layout(&self) -> VortexResult> { + let widths = materialized_layout(self.chunk_offsets())?; + if let Some(widths) = &widths { + BitPackedData::validate_layout( + self.packed(), + self.as_ref().dtype().as_ptype(), + widths, + )?; + } + Ok(widths) + } + + /// Read one byte boundary without executing the entire offsets child. + fn chunk_byte_offset(&self, boundary: usize, ctx: &mut ExecutionCtx) -> VortexResult { + vortex_ensure!( + boundary < self.chunk_offsets().len(), + "Chunk boundary out of bounds" + ); + if let Some(offsets) = self + .chunk_offsets() + .as_opt::() + .filter(|a| a.buffer_handle().is_on_host()) + { + Ok(offsets.as_slice::()[boundary]) + } else { + u64::try_from(&self.chunk_offsets().execute_scalar(boundary, ctx)?) + } + } + + /// Locate and validate one chunk using scalar child access, without materializing the offsets. + fn chunk_range( + &self, + chunk: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult<(Range, u8)> { + vortex_ensure!( + chunk < self.chunk_offsets().len() - 1, + "Chunk index out of bounds" + ); + let base = self.chunk_byte_offset(0, ctx)?; + let start = if chunk == 0 { + base + } else { + self.chunk_byte_offset(chunk, ctx)? + }; + let end = self.chunk_byte_offset(chunk + 1, ctx)?; + let width = width_from_offsets(start, end)?; + vortex_ensure!( + width as usize <= self.as_ref().dtype().as_ptype().bit_width(), + "Unsupported bit width {width}" + ); + let start = start + .checked_sub(base) + .ok_or_else(|| vortex_err!("Chunk offset precedes buffer origin"))?; + let end = end + .checked_sub(base) + .ok_or_else(|| vortex_err!("Chunk offset precedes buffer origin"))?; + vortex_ensure!( + start % (FL_CHUNK_SIZE / 8) as u64 == 0 && end <= self.packed().len() as u64, + "Chunk offsets are unaligned or exceed the packed buffer" + ); + Ok((usize::try_from(start)?..usize::try_from(end)?, width)) + } + #[inline] fn offset(&self) -> u16 { BitPackedData::offset(self) @@ -541,17 +579,17 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { BitPackedData::packed_slice::(self) } - #[inline] + /// Iterate packed chunks using boundaries materialized for this operation. fn unpacked_chunks<'a, T: BitPackedIter>( &'a self, + widths: &'a ChunkLayout, scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], ) -> VortexResult> { - BitPackedData::unpacked_chunks::( - self, - self.as_ref().dtype(), - self.as_ref().len(), - scratch, - ) + vortex_ensure!( + T::PTYPE == self.as_ref().dtype().as_ptype(), + "Requested unpack type does not match array dtype" + ); + BitUnpackedChunks::try_new(self, self.as_ref().len(), widths, scratch) } } diff --git a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index 4877fa9c57f..95d97983dec 100644 --- a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs @@ -16,6 +16,8 @@ use vortex_error::vortex_ensure; use crate::BitPackedData; use crate::FL_CHUNK_SIZE; +use crate::bitpacking::array::ChunkLayout; +use crate::bitpacking::array::chunk_packed_bytes; const CHUNK_SIZE: usize = FL_CHUNK_SIZE; @@ -48,6 +50,16 @@ impl> UnpackStrategy for BitPackingStr } } +/// The packed FastLanes block of `chunk` and its bit width. +#[allow(clippy::inline_always)] +#[inline(always)] +fn packed_chunk<'a, P>(packed: &'a [P], widths: &ChunkLayout, chunk: usize) -> (&'a [P], usize) { + let bit_width = widths.width(chunk); + let start = widths.byte_offset(chunk) / size_of::

(); + let len = chunk_packed_bytes(bit_width) / size_of::

(); + (&packed[start..][..len], bit_width as usize) +} + /// Accessor to unpacked chunks of bitpacked arrays /// /// The usual pattern of usage should follow @@ -63,23 +75,22 @@ impl> UnpackStrategy for BitPackingStr /// use vortex_buffer::buffer; /// use vortex_fastlanes::BitPackedData; /// use vortex_fastlanes::BitPackedArrayExt; +/// use vortex_fastlanes::FL_CHUNK_SIZE; /// use vortex_fastlanes::unpack_iter::BitUnpackedChunks; /// /// let mut ctx = vortex_array::array_session().create_execution_ctx(); /// let array = BitPackedData::encode(&buffer![2, 3, 4, 5].into_array(), 2, &mut ctx).unwrap(); -/// let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; -/// let mut unpacked_chunks: BitUnpackedChunks = -/// array.unpacked_chunks(&mut scratch).unwrap(); +/// let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; +/// let widths = array.chunk_layout(&mut ctx).unwrap(); +/// let mut unpacked_chunks: BitUnpackedChunks = array.unpacked_chunks(&widths, &mut scratch).unwrap(); /// /// if let Some(header) = unpacked_chunks.initial() { /// // handle partial initial chunk /// } /// -/// { -/// let mut chunks_iter = unpacked_chunks.full_chunks(); -/// while let Some(chunk) = chunks_iter.next() { -/// // handle full bitpacked chunks of 1024 elements -/// } +/// let mut chunks_iter = unpacked_chunks.full_chunks(); +/// while let Some(chunk) = chunks_iter.next() { +/// // handle full bitpacked chunks of 1024 elements /// } /// /// if let Some(trailer) = unpacked_chunks.trailer() { @@ -88,7 +99,7 @@ impl> UnpackStrategy for BitPackingStr /// ``` pub struct UnpackedChunks<'a, T: PhysicalPType, S: UnpackStrategy> { strategy: S, - bit_width: usize, + widths: &'a ChunkLayout, offset: usize, len: usize, num_chunks: usize, @@ -104,12 +115,13 @@ impl<'a, T: BitPacked> BitUnpackedChunks<'a, T> { pub fn try_new( array: &'a BitPackedData, len: usize, + widths: &'a ChunkLayout, scratch: &'a mut [MaybeUninit; CHUNK_SIZE], ) -> VortexResult { Self::try_new_with_strategy( BitPackingStrategy, array.packed_slice::(), - array.bit_width() as usize, + widths, array.offset() as usize, len, scratch, @@ -117,14 +129,12 @@ impl<'a, T: BitPacked> BitUnpackedChunks<'a, T> { } pub fn full_chunks(&mut self) -> BitUnpackIterator<'_, T> { - let elems_per_chunk = self.elems_per_chunk(); let last_chunk_is_sliced = self.last_chunk_is_sliced() as usize; let first_chunk_is_sliced = self.first_chunk_is_sliced(); BitUnpackIterator::new( self.packed, + self.widths, self.scratch, - self.bit_width, - elems_per_chunk, self.num_chunks - last_chunk_is_sliced, first_chunk_is_sliced, ) @@ -135,16 +145,16 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { pub fn try_new_with_strategy( strategy: S, packed: &'a [T::Physical], - bit_width: usize, + widths: &'a ChunkLayout, offset: usize, len: usize, scratch: &'a mut [MaybeUninit; CHUNK_SIZE], ) -> VortexResult { let (num_chunks, last_chunk_length) = - validate_packed::(packed.len(), bit_width, offset, len)?; + validate_packed::(packed.len(), widths, offset, len)?; Ok(Self { strategy, - bit_width, + widths, offset, len, num_chunks, @@ -156,14 +166,14 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { #[allow(clippy::inline_always)] #[inline(always)] - const fn elems_per_chunk(&self) -> usize { - 128 * self.bit_width / size_of::() + fn chunk(&self, chunk: usize) -> (&'a [T::Physical], usize) { + packed_chunk(self.packed, self.widths, chunk) } /// Access first chunk of the array if the last chunk has fewer than 1024 due to slicing pub fn initial(&mut self) -> Option<&mut [T]> { (self.first_chunk_is_sliced() || self.num_chunks == 1).then(|| { - let chunk: &[T::Physical] = &self.packed[..self.elems_per_chunk()]; + let (chunk, bit_width) = self.chunk(0); let dst: &mut [MaybeUninit] = self.scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; @@ -173,10 +183,10 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { CHUNK_SIZE - self.offset }; // SAFETY: - // 1. chunk is elems_per_chunk. + // 1. chunk holds exactly one packed block at bit_width. // 2. buffer is exactly CHUNK_SIZE. unsafe { - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width, chunk, dst); mem::transmute(&mut self.scratch[self.offset..][..header_end_slice]) } }) @@ -230,13 +240,18 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { } if self.num_chunks > 1 { - let packed_slice = self.packed; - let elems_per_chunk = self.elems_per_chunk(); - for i in self.full_chunks_range() { - let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; + let range = self.full_chunks_range(); + let packed = self.packed; + let widths: &'a ChunkLayout = self.widths; + let mut start = widths.byte_offset(range.start) / size_of::(); + for chunk_idx in range { + let bit_width = widths.width(chunk_idx); + let len = chunk_packed_bytes(bit_width) / size_of::(); + let chunk = &packed[start..start + len]; + start += len; unsafe { let dst: &mut [T::Physical] = mem::transmute(&mut self.scratch[..]); - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width as usize, chunk, dst); let unpacked: &mut [T] = mem::transmute(&mut self.scratch[..]); f(unpacked, local_idx..local_idx + CHUNK_SIZE); } @@ -265,16 +280,19 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { let mut local_idx = start_idx; - let packed_slice = self.packed; - let elems_per_chunk = self.elems_per_chunk(); - for i in self.full_chunks_range() { - let chunk = &packed_slice[i * elems_per_chunk..][..elems_per_chunk]; + let range = self.full_chunks_range(); + let mut start = self.widths.byte_offset(range.start) / size_of::(); + for chunk_idx in range { + let bit_width = self.widths.width(chunk_idx); + let len = chunk_packed_bytes(bit_width) / size_of::(); + let chunk = &self.packed[start..start + len]; + start += len; unsafe { let uninit_dst = &mut output[local_idx..local_idx + CHUNK_SIZE]; // SAFETY: &[T] and &[MaybeUninit] have the same layout. let dst: &mut [T::Physical] = mem::transmute(uninit_dst); - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width as usize, chunk, dst); } local_idx += CHUNK_SIZE; } @@ -289,15 +307,14 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { /// Access last chunk of the array if the last chunk has fewer than 1024 due to slicing pub fn trailer(&mut self) -> Option<&mut [T]> { (self.last_chunk_is_sliced() && self.num_chunks > 1).then(|| { - let chunk: &[T::Physical] = &self.packed - [(self.num_chunks - 1) * self.elems_per_chunk()..][..self.elems_per_chunk()]; + let (chunk, bit_width) = self.chunk(self.num_chunks - 1); let dst: &mut [MaybeUninit] = self.scratch; let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; // SAFETY: - // 1. chunk is elems_per_chunk. + // 1. chunk holds exactly one packed block at bit_width. // 2. buffer is exactly CHUNK_SIZE. unsafe { - self.strategy.unpack_chunk(self.bit_width, chunk, dst); + self.strategy.unpack_chunk(bit_width, chunk, dst); mem::transmute(&mut self.scratch[..self.last_chunk_length]) } }) @@ -312,33 +329,49 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { } } -/// Walk every packed chunk in array order without allocating an unpack scratch buffer. +/// Walk every *packed* chunk in array order, yielding the raw packed FastLanes block, its bit +/// width, and the padded bit range it covers, without unpacking it. +/// +/// Unlike [`UnpackedChunks::for_each_unpacked_chunk`], this does not fill a scratch buffer: it +/// hands the still-packed block to the callback so fused kernels (e.g. compare) can unpack and +/// consume it in a single pass. +/// +/// The yielded range is in *padded* coordinates: block `c` covers +/// `[c * 1024, min((c + 1) * 1024, offset + len))`, so it includes the leading `offset` rows +/// that slicing skips. Block starts are therefore always 1024-aligned regardless of `offset`. +/// Callers must account for the array's `offset` when mapping a block's rows back to logical +/// output positions (e.g. by viewing the output buffer at a bit offset of `offset`). pub(crate) fn for_each_packed_chunk( packed: &[T::Physical], - bit_width: usize, + widths: &ChunkLayout, offset: usize, len: usize, mut f: F, ) -> VortexResult<()> where T: PhysicalPType, - F: FnMut(&[T::Physical], Range), + F: FnMut(&[T::Physical], usize, Range), { - let (num_chunks, _) = validate_packed::(packed.len(), bit_width, offset, len)?; - let elems_per_chunk = 128 * bit_width / size_of::(); + validate_packed::(packed.len(), widths, offset, len)?; let padded_len = offset + len; - for chunk in 0..num_chunks { - let packed_chunk = &packed[chunk * elems_per_chunk..][..elems_per_chunk]; - let start = chunk * CHUNK_SIZE; - let end = (start + CHUNK_SIZE).min(padded_len); - f(packed_chunk, start..end); + let mut start = 0; + for chunk in 0..widths.len() { + let bit_width = widths.width(chunk); + let packed_len = chunk_packed_bytes(bit_width) / size_of::(); + let packed_chunk = &packed[start..start + packed_len]; + start += packed_len; + let row_start = chunk * CHUNK_SIZE; + let row_end = (row_start + CHUNK_SIZE).min(padded_len); + f(packed_chunk, bit_width as usize, row_start..row_end); } Ok(()) } +/// Check that `packed_len` words of `T::Physical` hold exactly the chunks described by `widths` +/// for `offset + len` padded elements, returning the chunk count and the trailing chunk's length. fn validate_packed( packed_len: usize, - bit_width: usize, + widths: &ChunkLayout, offset: usize, len: usize, ) -> VortexResult<(usize, usize)> { @@ -346,12 +379,20 @@ fn validate_packed( offset < CHUNK_SIZE, "Invalid bit-packed offset {offset}, expected < {CHUNK_SIZE}" ); - let elems_per_chunk = 128 * bit_width / size_of::(); let num_chunks = (offset + len).div_ceil(CHUNK_SIZE); vortex_ensure!( - packed_len == num_chunks * elems_per_chunk, - "Invalid packed length: got {packed_len}, expected {}", - num_chunks * elems_per_chunk + widths.len() == num_chunks, + "Invalid chunk widths: got {}, expected {num_chunks}", + widths.len() + ); + vortex_ensure!( + widths.max_width() as usize <= size_of::() * 8, + "Chunk width exceeds unpacked type width" + ); + let expected = widths.packed_bytes() / size_of::(); + vortex_ensure!( + packed_len == expected, + "Invalid packed length: got {packed_len}, expected {expected}" ); Ok((num_chunks, (offset + len) % CHUNK_SIZE)) } @@ -359,29 +400,30 @@ fn validate_packed( /// Iterator over full chunks of bitpacked array that yields unpacked chunks one at a time pub struct BitUnpackIterator<'a, T: BitPacked + 'a> { packed: &'a [T::Physical], + widths: &'a ChunkLayout, buffer: &'a mut [MaybeUninit; CHUNK_SIZE], - bit_width: usize, - elems_per_chunk: usize, num_chunks: usize, idx: usize, + /// Word offset of chunk `idx` within `packed`. + start: usize, } impl<'a, T: BitPacked> BitUnpackIterator<'a, T> { pub fn new( packed: &'a [T::Physical], + widths: &'a ChunkLayout, buffer: &'a mut [MaybeUninit; CHUNK_SIZE], - bit_width: usize, - elems_per_chunk: usize, num_chunks: usize, first_chunk_is_sliced: bool, ) -> Self { + let idx = if first_chunk_is_sliced { 1 } else { 0 }; Self { packed, + widths, buffer, - bit_width, - elems_per_chunk, num_chunks, - idx: if first_chunk_is_sliced { 1 } else { 0 }, + idx, + start: widths.byte_offset(idx) / size_of::(), } } } @@ -398,15 +440,18 @@ impl<'a, T: BitPacked + 'a> LendingIterator for BitUnpackIterator<'a, T> { return None; } - let chunk = &self.packed[self.idx * self.elems_per_chunk..][..self.elems_per_chunk]; + let bit_width = self.widths.width(self.idx); + let len = chunk_packed_bytes(bit_width) / size_of::(); + let chunk = &self.packed[self.start..self.start + len]; let dst: &mut [MaybeUninit] = self.buffer; unsafe { let dst: &mut [T::Physical] = mem::transmute(dst); - BitPacking::unchecked_unpack(self.bit_width, chunk, dst); + BitPacking::unchecked_unpack(bit_width as usize, chunk, dst); } self.idx += 1; + self.start += len; // SAFETY: The buffer has the appropriate lifetime, the iterator signature doesn't account for it Some(unsafe { mem::transmute::<&mut [MaybeUninit; 1024], &mut [T; 1024]>(self.buffer) }) } diff --git a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs index 20e445fb0de..1e3da8864bf 100644 --- a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs +++ b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Tests for the chunk layout children of uniformly bit-packed arrays. +//! Behavioural tests for bit-packed arrays whose chunks are packed at different widths. use std::sync::LazyLock; @@ -9,7 +9,9 @@ use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::slice::SliceKernel; use vortex_array::assert_arrays_eq; +use vortex_array::scalar::Scalar; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; @@ -64,7 +66,26 @@ fn invalid_offsets_rejected_before_unpacking(#[case] offsets: Buffer) -> Vo let values = PrimitiveArray::from_iter((0..3072u32).map(|i| i % 2)); let packed = bitpack_to_best_bit_width(&values, &mut ctx)?; let offsets = offsets.into_array(); - assert!(BitPacked::with_chunk_offsets(packed, offsets).is_err()); + assert!(BitPacked::with_chunk_offsets(packed.clone(), offsets.clone()).is_err()); + let offsets = offsets.execute::(&mut ctx)?; + let offsets = bitpack_to_best_bit_width(&offsets, &mut ctx)?.into_array(); + let packed = BitPacked::with_chunk_offsets(packed, offsets)?.into_array(); + // An isolated scalar checks only its own chunk; bulk unpacking validates the whole layout. + assert_eq!(packed.execute_scalar(1, &mut ctx)?, Scalar::from(1u32)); + assert!(packed.execute_scalar(2048, &mut ctx).is_err()); + assert!(packed.execute::(&mut ctx).is_err()); + Ok(()) +} + +#[test] +fn slice_rejects_unaligned_offsets() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0..3072u32).map(|i| i % 2)); + let packed = bitpack_to_best_bit_width(&values, &mut ctx)?; + let offsets = PrimitiveArray::from_iter([0u64, 127, 255, 383]); + let offsets = bitpack_to_best_bit_width(&offsets, &mut ctx)?.into_array(); + let packed = BitPacked::with_chunk_offsets(packed, offsets)?; + assert!(::slice(packed.as_view(), 1024..2048, &mut ctx).is_err()); Ok(()) } diff --git a/encodings/fastlanes/src/bitpacking/compute/between.rs b/encodings/fastlanes/src/bitpacking/compute/between.rs index 8010fa208d2..97d3444b18c 100644 --- a/encodings/fastlanes/src/bitpacking/compute/between.rs +++ b/encodings/fastlanes/src/bitpacking/compute/between.rs @@ -23,6 +23,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::bitpacking::compute::stream_predicate::stream_predicate; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; impl BetweenKernel for BitPacked { fn between( @@ -74,7 +75,7 @@ fn between_constant_typed( ctx: &mut ExecutionCtx, ) -> VortexResult where - T: NativePType + Copy + crate::unpack_iter::BitPacked, + T: NativePType + Copy + BitPackedIter, { // Branch on strictness once at the top so each call into `between_impl` monomorphises // a single tight predicate — same shape as `Primitive::between` in `vortex-array`. @@ -128,7 +129,7 @@ fn between_impl( ctx: &mut ExecutionCtx, ) -> VortexResult where - T: NativePType + Copy + crate::unpack_iter::BitPacked, + T: NativePType + Copy + BitPackedIter, Lo: Fn(T, T) -> bool, Up: Fn(T, T) -> bool, { diff --git a/encodings/fastlanes/src/bitpacking/compute/compare.rs b/encodings/fastlanes/src/bitpacking/compute/compare.rs index c9d6b815b0d..807b6f21b8f 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare.rs @@ -28,7 +28,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::bitpacking::compute::compare_fused::stream_compare_fused; -use crate::unpack_iter::BitPacked as BitPackedIter; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; impl CompareKernel for BitPacked { fn compare( diff --git a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs index 1259ed815fe..7a9497fc090 100644 --- a/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs +++ b/encodings/fastlanes/src/bitpacking/compute/compare_fused.rs @@ -11,7 +11,7 @@ //! `[bool; 1024]` or `[T; 1024]` scratch. A single SIMD [`transpose_bits`] per block then rotates //! that mask into logical row order. //! -//! The packed blocks are walked through [`crate::unpack_iter::for_each_packed_chunk`], so chunk +//! The packed blocks are walked through [`crate::bitpacking::unpack_iter::for_each_packed_chunk`], so chunk //! sizing and bounds live in one place without allocating an unpack scratch buffer. //! //! Slicing is handled by working in *padded* coordinates: bit `offset + i` holds element `i`. The @@ -48,8 +48,8 @@ use vortex_error::VortexResult; use super::stream_predicate::stream_predicate; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::unpack_iter::BitPacked as BitPackedIter; -use crate::unpack_iter::for_each_packed_chunk; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; +use crate::bitpacking::unpack_iter::for_each_packed_chunk; const CHUNK_SIZE: usize = 1024; const U64_BITS: usize = u64::BITS as usize; @@ -78,12 +78,12 @@ where F: Fn(T, T) -> bool + Copy, { let len = array.len(); - let bit_width = array.bit_width() as usize; + let widths = array.chunk_layout(ctx)?; let offset = array.offset() as usize; // A degenerate width has no packed payload for the fused kernel to consume; defer to the scalar // streaming predicate, which handles every layout (including the empty array). - if len == 0 || bit_width == 0 { + if len == 0 || widths.max_width() == 0 { return stream_predicate::(array, nullability, move |v| cmp(v, rhs), ctx); } @@ -97,14 +97,19 @@ where let mut lane_major = [0u64; WORDS_PER_CHUNK]; for_each_packed_chunk::( array.packed_slice::<::Physical>(), - bit_width, + &widths, offset, len, - |packed_chunk, range| { + |packed_chunk, bit_width, range| { // Block starts are always 1024-aligned (padded coords), so the slot is a full block. let out = words[range.start / U64_BITS..] .first_chunk_mut::() .vortex_expect("over-allocated buffer holds a full block per chunk"); + // A zero-width chunk has no packed payload: every value in it is zero. + if bit_width == 0 { + out.fill(if cmp(T::default(), rhs) { u64::MAX } else { 0 }); + return; + } // SAFETY: `packed_chunk` holds exactly `128 * bit_width / size_of::()` packed // elements and `bit_width <= U::T`, satisfying `unchecked_unpack_cmp`'s contract. The // kernel assigns every word in `transposed`, so its previous contents are irrelevant. diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 0b1b9422f86..bdd551f2dd2 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -26,6 +26,7 @@ use super::take::UNPACK_CHUNK_THRESHOLD; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::BitPackedData; +use crate::ChunkLayout; /// The threshold over which it is faster to fully unpack the entire [`BitPackedArray`](crate::BitPackedArray) and then /// filter the result than to unpack only specific bitpacked values into the output buffer. @@ -65,7 +66,7 @@ impl FilterKernel for BitPacked { // Filter and patch using the correct unsigned type for FastLanes, then cast to signed if needed. let primitive = match_each_unsigned_integer_ptype!(array.dtype().as_ptype().to_unsigned(), |U| { - let (buffer, validity) = filter_primitive_without_patches::(array, values)?; + let (buffer, validity) = filter_primitive_without_patches::(array, values, ctx)?; // reinterpret_cast for signed types. let primitive = PrimitiveArray::new(buffer, validity); if array.dtype().as_ptype().is_signed_int() { @@ -109,8 +110,10 @@ impl FilterKernel for BitPacked { fn filter_primitive_without_patches( array: ArrayView<'_, BitPacked>, selection: &MaskValuesRef, + ctx: &mut ExecutionCtx, ) -> VortexResult<(Buffer, Validity)> { - let values = filter_with_indices(array.data(), selection.indices()); + let widths = array.chunk_layout(ctx)?; + let values = filter_with_indices(array.data(), &widths, selection.indices()); let validity = array .validity()? .filter(&Mask::Values(MaskValuesRef::clone(selection)))?; @@ -120,24 +123,21 @@ fn filter_primitive_without_patches( fn filter_with_indices( array: &BitPackedData, + widths: &ChunkLayout, indices: &[usize], ) -> BufferMut { let offset = array.offset() as usize; - let bit_width = array.bit_width() as usize; let mut values = BufferMut::with_capacity(indices.len()); // Some re-usable memory to store per-chunk indices. let mut unpacked = [const { MaybeUninit::::uninit() }; 1024]; - let packed_bytes = array.packed_slice::(); // Group the indices by the FastLanes chunk they belong to. - let chunk_size = 128 * bit_width / size_of::(); - chunked_indices( indices.iter().copied(), offset, |chunk_idx, indices_within_chunk| { - let packed = &packed_bytes[chunk_idx * chunk_size..][..chunk_size]; + let (packed, bit_width) = array.packed_chunk::(widths, chunk_idx); if indices_within_chunk.len() == 1024 { // Unpack the entire chunk. diff --git a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs index 0ab01a635ba..b249090abdf 100644 --- a/encodings/fastlanes/src/bitpacking/compute/is_constant.rs +++ b/encodings/fastlanes/src/bitpacking/compute/is_constant.rs @@ -23,7 +23,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::unpack_iter::BitPacked as BitPackedUnpack; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedUnpack; /// BitPacked-specific is_constant kernel with SIMD support. #[derive(Debug)] @@ -56,8 +56,9 @@ fn bitpacked_is_constant( array: ArrayView<'_, BitPacked>, ctx: &mut ExecutionCtx, ) -> VortexResult { + let widths = array.chunk_layout(ctx)?; let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; - let mut bit_unpack_iterator = array.unpacked_chunks::(&mut scratch)?; + let mut bit_unpack_iterator = array.unpacked_chunks::(&widths, &mut scratch)?; let patches = array .patches() .map(|p| -> VortexResult<_> { diff --git a/encodings/fastlanes/src/bitpacking/compute/mod.rs b/encodings/fastlanes/src/bitpacking/compute/mod.rs index 38f86f781bb..4a06d67b9d2 100644 --- a/encodings/fastlanes/src/bitpacking/compute/mod.rs +++ b/encodings/fastlanes/src/bitpacking/compute/mod.rs @@ -53,7 +53,7 @@ mod tests { use vortex_array::compute::conformance::consistency::test_array_consistency; use crate::BitPackedArray; - use crate::bitpack_compress::bitpack_encode; + use crate::bitpacking::bitpack_compress::bitpack_encode; use crate::bitpacking::compute::chunked_indices; fn bp(array: &PrimitiveArray, bit_width: u8) -> BitPackedArray { diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs index 30679b1914f..db8193ff991 100644 --- a/encodings/fastlanes/src/bitpacking/compute/slice.rs +++ b/encodings/fastlanes/src/bitpacking/compute/slice.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::cmp::max; use std::ops::Range; use vortex_array::ArrayRef; @@ -12,6 +11,7 @@ use vortex_array::arrays::slice::SliceKernel; use vortex_array::arrays::slice::SliceReduce; use vortex_array::patches::Patches; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use crate::BitPacked; use crate::BitPackedArraySlotsExt; @@ -24,7 +24,12 @@ impl SliceReduce for BitPacked { return Ok(None); } - Ok(Some(slice_bitpacked(array, range, None)?)) + let Some(widths) = array.materialized_chunk_layout()? else { + return Ok(None); + }; + let (chunks, _) = slice_chunks(array.offset(), &range); + let encoded = widths.byte_offset(chunks.start)..widths.byte_offset(chunks.end); + Ok(Some(slice_bitpacked(array, encoded, range, None)?)) } } @@ -32,7 +37,7 @@ impl SliceKernel for BitPacked { fn slice( array: ArrayView<'_, Self>, range: Range, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let patches = array .patches() @@ -40,38 +45,51 @@ impl SliceKernel for BitPacked { .transpose()? .flatten(); - Ok(Some(slice_bitpacked(array, range, patches)?)) + let (chunks, _) = slice_chunks(array.offset(), &range); + let base = array.chunk_byte_offset(0, ctx)?; + let start = array.chunk_byte_offset(chunks.start, ctx)?; + let end = array.chunk_byte_offset(chunks.end, ctx)?; + vortex_ensure!( + base <= start + && start <= end + && end - base <= array.packed().len() as u64 + && (start - base).is_multiple_of(128) + && (end - base).is_multiple_of(128), + "Slice chunk offsets exceed the packed buffer" + ); + let encoded = usize::try_from(start - base)?..usize::try_from(end - base)?; + Ok(Some(slice_bitpacked(array, encoded, range, patches)?)) } } fn slice_bitpacked( array: ArrayView<'_, BitPacked>, + encoded: Range, range: Range, patches: Option, ) -> VortexResult { - let offset_start = range.start + array.offset() as usize; - let offset_stop = range.end + array.offset() as usize; - let offset = offset_start % 1024; - let block_start = max(0, offset_start - offset); - let block_stop = offset_stop.div_ceil(1024) * 1024; - - let encoded_start = (block_start / 8) * array.bit_width() as usize; - let encoded_stop = (block_stop / 8) * array.bit_width() as usize; + let (chunks, offset) = slice_chunks(array.offset(), &range); + let chunk_start = chunks.start; + let chunk_stop = chunks.end; Ok(BitPacked::try_new( - array.packed().slice(encoded_start..encoded_stop), + array.packed().slice(encoded), array.dtype().as_ptype(), array.validity()?.slice(range.clone())?, patches, - array - .chunk_offsets() - .slice(block_start / 1024..block_stop / 1024 + 1)?, + array.chunk_offsets().slice(chunk_start..chunk_stop + 1)?, range.len(), offset as u16, )? .into_array()) } +fn slice_chunks(offset: u16, range: &Range) -> (Range, usize) { + let start = range.start + offset as usize; + let stop = range.end + offset as usize; + (start / 1024..stop.div_ceil(1024), start % 1024) +} + #[cfg(test)] mod tests { use vortex_array::IntoArray; @@ -82,7 +100,7 @@ mod tests { use vortex_error::VortexResult; use crate::BitPacked; - use crate::bitpack_compress::bitpack_encode; + use crate::bitpacking::bitpack_compress::bitpack_encode; #[test] fn test_reduce_parent_returns_bitpacked_slice() -> VortexResult<()> { diff --git a/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs b/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs index 9154ca736c1..fe0724df31e 100644 --- a/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs +++ b/encodings/fastlanes/src/bitpacking/compute/stream_predicate.rs @@ -35,7 +35,7 @@ use vortex_error::VortexResult; use crate::BitPacked; use crate::BitPackedArrayExt; use crate::FL_CHUNK_SIZE; -use crate::unpack_iter::BitPacked as BitPackedIter; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; /// Stream `predicate` over the unpacked values of a [`BitPackedArray`](crate::BitPackedArray), one FastLanes /// block at a time, producing a [`BoolArray`]. @@ -53,8 +53,9 @@ where let mut words: BufferMut = BufferMut::zeroed(len.div_ceil(u64::BITS as usize)); if len > 0 { + let widths = array.chunk_layout(ctx)?; let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; - let mut chunks = array.unpacked_chunks::(&mut scratch)?; + let mut chunks = array.unpacked_chunks::(&widths, &mut scratch)?; let words = words.as_mut_slice(); if let Some(p) = array.patches() { diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 86e97623cf6..ab570c072db 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/encodings/fastlanes/src/bitpacking/compute/take.rs @@ -25,7 +25,7 @@ use vortex_error::VortexResult; use super::chunked_indices; use crate::BitPacked; use crate::BitPackedArrayExt; -use crate::bitpack_decompress; +use crate::bitpacking::bitpack_decompress; // TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. /// assuming the buffer is already allocated (which will happen at most once) then unpacking @@ -80,10 +80,8 @@ fn take_primitive( return Ok(PrimitiveArray::new(Buffer::::empty(), taken_validity)); } + let widths = array.chunk_layout(ctx)?; let offset = array.offset() as usize; - let bit_width = array.bit_width() as usize; - - let packed = array.packed_slice::(); // Group indices by 1024-element chunk, *without* allocating on the heap let indices_iter = indices.as_slice::().iter().map(|i| { @@ -93,10 +91,9 @@ fn take_primitive( let mut output = BufferMut::::with_capacity(indices.len()); let mut unpacked = [const { MaybeUninit::uninit() }; 1024]; - let chunk_len = 128 * bit_width / size_of::(); chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| { - let packed = &packed[chunk_idx * chunk_len..][..chunk_len]; + let (packed, bit_width) = array.data().packed_chunk::(&widths, chunk_idx); let mut have_unpacked = false; let (offset_chunks, remainder) = indices_within_chunk.as_chunks::(); diff --git a/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs b/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs index 30ea0ac70d9..7ecb9893c35 100644 --- a/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs +++ b/encodings/fastlanes/src/bitpacking/plugin/bitpacked.rs @@ -181,7 +181,7 @@ impl ArrayPlugin for BitPackedPlugin { s.push(Some(offsets)); s }; - let data = BitPackedData::try_new(packed, patches, bit_width, offset)?; + let data = BitPackedData::try_new(packed, patches, offset)?; Ok(Array::::try_from_parts( ArrayParts::new(BitPacked, dtype.clone(), len, data).with_slots(slots), )? diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 317d0978f42..1a0384462d2 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -34,7 +34,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -45,9 +44,7 @@ use crate::BitPackedData; use crate::BitPackedDataParts; use crate::bitpacking::array::BitPackedSlots; use crate::bitpacking::array::BitPackedSlotsView; -use crate::bitpacking::array::CHUNK_OFFSETS_DTYPE; use crate::bitpacking::array::PATCH_SLOTS; -use crate::bitpacking::array::materialized_layout; use crate::bitpacking::bitpack_decompress::unpack_array; use crate::bitpacking::bitpack_decompress::unpack_into_primitive_builder; use crate::bitpacking::vtable::rules::RULES; @@ -66,7 +63,6 @@ pub(crate) fn initialize(session: &VortexSession) { impl ArrayHash for BitPackedData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.offset.hash(state); - self.bit_width.hash(state); self.packed.array_hash(state, accuracy); self.patches_data.hash(state); } @@ -75,7 +71,6 @@ impl ArrayHash for BitPackedData { impl ArrayEq for BitPackedData { fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { self.offset == other.offset - && self.bit_width == other.bit_width && self.packed.array_eq(&other.packed, accuracy) && self.patches_data == other.patches_data } @@ -237,14 +232,6 @@ impl BitPacked { len: usize, offset: u16, ) -> VortexResult { - vortex_ensure!( - chunk_offsets.dtype() == &CHUNK_OFFSETS_DTYPE, - "Expected non-nullable u64 offsets" - ); - let layout = materialized_layout(&chunk_offsets)?.ok_or_else(|| { - vortex_err!("Chunk offsets must be materialized while kernels use bit_width") - })?; - let bit_width = layout.uniform_width().unwrap_or(0); let dtype = DType::Primitive(ptype, validity.nullability()); let slots = { let mut s = ArraySlots::with_capacity(BitPackedSlots::COUNT); @@ -253,12 +240,12 @@ impl BitPacked { s.push(Some(chunk_offsets)); s }; - let data = BitPackedData::try_new(packed, patches, bit_width, offset)?; + let data = BitPackedData::try_new(packed, patches, offset)?; Array::try_from_parts(ArrayParts::new(BitPacked, dtype, len, data).with_slots(slots)) } /// Replace the non-nullable `u64` chunk boundaries, including the trailing boundary. - /// Boundaries must remain materialized and imply the scalar `bit_width`. + /// Compressed values are validated at execution time, before unpacking. pub fn with_chunk_offsets( array: BitPackedArray, offsets: ArrayRef, @@ -280,7 +267,6 @@ impl BitPacked { let data = array.into_data(); BitPackedDataParts { offset: data.offset, - bit_width: data.bit_width, chunk_offsets, len, packed: data.packed, diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index ba3914c1a14..bc319f53196 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -8,15 +8,15 @@ use vortex_array::vtable::OperationsVTable; use vortex_error::VortexResult; use crate::BitPacked; -use crate::bitpack_decompress; use crate::bitpacking::array::BitPackedArrayExt; +use crate::bitpacking::bitpack_decompress; impl OperationsVTable for BitPacked { type ProbeState = (); fn scalar_at( array: ArrayView<'_, BitPacked>, index: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { Ok( if let Some(patches) = array.patches() @@ -24,7 +24,7 @@ impl OperationsVTable for BitPacked { { patch } else { - bitpack_decompress::unpack_single(array, index) + bitpack_decompress::unpack_single(array, index, ctx)? }, ) } diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index e41eb27a5be..f0162301d34 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -96,6 +96,7 @@ pub(crate) fn fused_decompress< .as_::() .vortex_expect("cannot be null"); + let widths = bp.chunk_layout(ctx)?; let strategy = FoRStrategy { reference: ref_ }; let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; @@ -103,7 +104,7 @@ pub(crate) fn fused_decompress< let mut unpacked = UnpackedChunks::try_new_with_strategy( strategy, bp.packed_slice::(), - bp.bit_width() as usize, + &widths, bp.offset() as usize, bp.len(), &mut scratch, diff --git a/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap index c7bf7d84c3d..98368621a48 100644 --- a/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap @@ -8,7 +8,7 @@ root: vortex.onpair(utf8, len=16384) nbytes=140258 dict_offsets: vortex.primitive(u16, len=1626) nbytes=3252 metadata: ptype: u16 codes: fastlanes.bitpacked(u16, len=63845) nbytes=89216 - metadata: bit_width: 11, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=64) nbytes=512 metadata: ptype: u64 codes_offsets: vortex.primitive(u16, len=16385) nbytes=32770 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap index 84ec2952483..adbc898b4a7 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap @@ -6,7 +6,7 @@ input: binary, len=16384, nbytes=315856 root: vortex.dict(binary, len=16384) nbytes=6332 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=6280 - metadata: bit_width: 3, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 values: vortex.varbin(binary, len=5) nbytes=52 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap index 3c027e5421f..26a9b2b46a7 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap @@ -6,6 +6,6 @@ input: decimal(12,2), len=16384, nbytes=131072 root: vortex.decimal_byte_parts.v2(decimal(12,2), len=16384) nbytes=49288 metadata: msp: fastlanes.bitpacked(i32, len=16384) nbytes=49288 - metadata: bit_width: 24, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap b/vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap index e0e11486a3d..4bc54dd8303 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap @@ -6,6 +6,6 @@ input: f64, len=16384, nbytes=131072 root: vortex.alp(f64, len=16384) nbytes=49288 metadata: exponents: e: 14, f: 12 encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49288 - metadata: bit_width: 24, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap b/vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap index 15588bd7c38..eb394c3205d 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap @@ -6,11 +6,11 @@ input: f64, len=16384, nbytes=131072 root: vortex.alprd(f64, len=16384) nbytes=113152 metadata: right_bit_width: 52, patch_offset: 0 left_parts: fastlanes.bitpacked(u16, len=16384) nbytes=6280 - metadata: bit_width: 3, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 right_parts: fastlanes.bitpacked(u64, len=16384) nbytes=106632 - metadata: bit_width: 52, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 patch_indices: vortex.primitive(u16, len=60) nbytes=120 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap index a212976976e..a4a91d272fc 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap @@ -8,7 +8,7 @@ root: vortex.alp(f64, len=16384) nbytes=6344 encoded: vortex.dict(i64, len=16384) nbytes=6344 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=6280 - metadata: bit_width: 3, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 values: vortex.primitive(i64, len=8) nbytes=64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap index b4c82316d67..61314dd9d82 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap @@ -6,7 +6,7 @@ input: i64, len=16384, nbytes=131072 root: vortex.dict(i64, len=16384) nbytes=6328 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=6280 - metadata: bit_width: 3, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 values: vortex.primitive(i64, len=6) nbytes=48 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap index 91b7c28c9c8..41f71eb7c44 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap @@ -6,6 +6,6 @@ input: u64, len=16384, nbytes=131072 root: fastlanes.for(u64, len=16384) nbytes=49288 metadata: reference: 1700000001036u64 encoded: fastlanes.bitpacked(u64, len=16384) nbytes=49288 - metadata: bit_width: 24, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap index 765c40faf1c..8867e80a3fd 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap @@ -8,7 +8,7 @@ root: vortex.sparse(i32?, len=16384) nbytes=3047 patch_indices: vortex.primitive(u16, len=823) nbytes=1646 metadata: ptype: u16 patch_values: fastlanes.bitpacked(i32?, len=823) nbytes=1399 - metadata: bit_width: 10, offset: 0 + metadata: offset: 0 validity_child: vortex.bool(bool, len=823) nbytes=103 metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=2) nbytes=16 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap index 0e9544544cc..e2bcbdf797f 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap @@ -6,6 +6,6 @@ input: i64, len=16384, nbytes=131072 root: fastlanes.for(i64, len=16384) nbytes=16520 metadata: reference: -128i64 encoded: fastlanes.bitpacked(i64, len=16384) nbytes=16520 - metadata: bit_width: 8, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap index b1cba7c8b98..ee9737ba574 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap @@ -8,12 +8,12 @@ root: vortex.runend(i32, len=16384) nbytes=4000 ends: fastlanes.for(u16, len=1020) nbytes=1808 metadata: reference: 13u16 encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1808 - metadata: bit_width: 14, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=2) nbytes=16 metadata: ptype: u64 values: fastlanes.for(i32, len=1020) nbytes=2192 metadata: reference: -49931i32 encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2192 - metadata: bit_width: 17, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=2) nbytes=16 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap index fd5d06d865d..dcd742436bb 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap @@ -10,6 +10,6 @@ root: vortex.sparse(i64, len=16384) nbytes=5556 patch_values: fastlanes.for(i64, len=848) nbytes=3856 metadata: reference: 1000830099i64 encoded: fastlanes.bitpacked(i64, len=848) nbytes=3856 - metadata: bit_width: 30, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=2) nbytes=16 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__list_of_int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__regular__list_of_int_runs.snap index 45a1cabc688..26029e2a0b8 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__list_of_int_runs.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__list_of_int_runs.snap @@ -10,17 +10,17 @@ root: vortex.list(list(i32), len=4066) nbytes=11218 ends: fastlanes.for(u16, len=1020) nbytes=1808 metadata: reference: 13u16 encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1808 - metadata: bit_width: 14, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=2) nbytes=16 metadata: ptype: u64 values: fastlanes.for(i32, len=1020) nbytes=2192 metadata: reference: -49931i32 encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2192 - metadata: bit_width: 17, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=2) nbytes=16 metadata: ptype: u64 offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7218 - metadata: bit_width: 14, offset: 0 + metadata: offset: 0 patch_indices: vortex.primitive(u16, len=1) nbytes=2 metadata: ptype: u16 patch_values: vortex.constant(u16, len=1) nbytes=4 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap b/vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap index 9f1a8748081..1925bedfec8 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap @@ -12,6 +12,6 @@ root: vortex.fsst(utf8, len=16384) nbytes=151526 patch_values: vortex.constant(u8, len=1575) nbytes=2 metadata: scalar: 23u8 codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=37136 - metadata: bit_width: 17, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=18) nbytes=144 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap b/vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap index 7a194041570..b5ffec9ec8b 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap @@ -6,7 +6,7 @@ input: utf8, len=16384, nbytes=262144 root: vortex.dict(utf8, len=16384) nbytes=8509 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=8328 - metadata: bit_width: 4, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 values: vortex.fsst(utf8, len=12) nbytes=181 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap index 7eab0810314..ecf96abcb3c 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap @@ -10,7 +10,7 @@ root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=57797 category: vortex.dict(utf8, len=16384) nbytes=8509 metadata: all_values_referenced: true codes: fastlanes.bitpacked(u8, len=16384) nbytes=8328 - metadata: bit_width: 4, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 values: vortex.fsst(utf8, len=12) nbytes=181 @@ -22,6 +22,6 @@ root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=57797 value: vortex.alp(f64, len=16384) nbytes=49288 metadata: exponents: e: 14, f: 12 encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49288 - metadata: bit_width: 24, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap b/vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap index 0184162d4d3..a85cb90068d 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap @@ -8,6 +8,6 @@ root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=67720 storage: fastlanes.for(i64, len=16384) nbytes=67720 metadata: reference: 1700000000891673i64 encoded: fastlanes.bitpacked(i64, len=16384) nbytes=67720 - metadata: bit_width: 33, offset: 0 + metadata: offset: 0 chunk_offsets: vortex.primitive(u64, len=17) nbytes=136 metadata: ptype: u64 diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index 43f6a5a3712..7fb3daec17f 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -559,7 +559,14 @@ impl FusedPlan { let bp = child.as_::(); let offset = slice_arr.data().slice_range().start; let len = array.len(); - let (packed, bitpacked_offset, patch_range) = bitpacked_slice_view(bp, offset, len)?; + let Some(widths) = bp.materialized_chunk_layout()? else { + vortex_bail!("Fused bit-unpack requires materialized chunk widths"); + }; + let (packed, widths, bitpacked_offset, patch_range) = + bitpacked_slice_view(bp, &widths, offset, len)?; + let Some(bit_width) = widths.uniform_width() else { + vortex_bail!("CUDA bit-unpack requires every chunk to share one bit width"); + }; let source_ptype = ptype_to_tag(PType::try_from(bp.dtype()).map_err(|_| { vortex_err!("BitPacked must have primitive dtype, got {:?}", bp.dtype()) @@ -567,7 +574,7 @@ impl FusedPlan { let buf_index = self.source_buffers.len(); self.source_buffers.push(Some(packed)); return Ok(Stage::new( - SourceOp::bitunpack(bp.bit_width(), bitpacked_offset), + SourceOp::bitunpack(bit_width, bitpacked_offset), Some(buf_index), source_ptype, ) @@ -615,6 +622,12 @@ impl FusedPlan { fn walk_bitpacked(&mut self, array: ArrayRef) -> VortexResult { let bp = array.as_::(); + let Some(bit_width) = bp + .materialized_chunk_layout()? + .and_then(|widths| widths.uniform_width()) + else { + vortex_bail!("CUDA bit-unpack requires every chunk to share one bit width"); + }; let source_ptype = ptype_to_tag(PType::try_from(bp.dtype()).map_err(|_| { vortex_err!("BitPacked must have primitive dtype, got {:?}", bp.dtype()) @@ -622,7 +635,7 @@ impl FusedPlan { let buf_index = self.source_buffers.len(); self.source_buffers.push(Some(bp.packed().clone())); Ok(Stage::new( - SourceOp::bitunpack(bp.bit_width(), bp.offset()), + SourceOp::bitunpack(bit_width, bp.offset()), Some(buf_index), source_ptype, ) diff --git a/vortex-cuda/src/kernel/encodings/bitpacked.rs b/vortex-cuda/src/kernel/encodings/bitpacked.rs index d7b8b089470..bb5b13ae3de 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -14,6 +14,8 @@ use vortex::array::ArrayRef; use vortex::array::ArrayVTable; use vortex::array::ArrayView; use vortex::array::Canonical; +use vortex::array::ExecutionCtx; +use vortex::array::IntoArray; use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::Slice; use vortex::array::arrays::slice::SliceArraySlotsExt; @@ -25,10 +27,11 @@ use vortex::dtype::NativePType; use vortex::encodings::fastlanes::BitPacked; use vortex::encodings::fastlanes::BitPackedArray; use vortex::encodings::fastlanes::BitPackedArrayExt; -use vortex::encodings::fastlanes::BitPackedArraySlotsExt; use vortex::encodings::fastlanes::BitPackedDataParts; +use vortex::encodings::fastlanes::ChunkLayout; use vortex::encodings::fastlanes::unpack_iter::BitPacked as BitPackedUnpack; use vortex::error::VortexResult; +use vortex::error::vortex_bail; use vortex::error::vortex_ensure; use vortex::error::vortex_err; @@ -52,9 +55,10 @@ pub(crate) struct BitPackedExecutor; /// materialization so exception metadata is sliced consistently. pub(crate) fn bitpacked_slice_view( bp: ArrayView<'_, BitPacked>, + widths: &ChunkLayout, offset: usize, len: usize, -) -> VortexResult<(BufferHandle, u16, Range)> { +) -> VortexResult<(BufferHandle, ChunkLayout, u16, Range)> { let patch_range = offset..offset + len; let offset_start = patch_range.start + bp.offset() as usize; let offset_stop = offset_start + len; @@ -62,11 +66,14 @@ pub(crate) fn bitpacked_slice_view( let block_start = offset_start - bitpacked_offset; let block_stop = offset_stop.div_ceil(PATCH_CHUNK_SIZE) * PATCH_CHUNK_SIZE; - let encoded_start = (block_start / 8) * bp.bit_width() as usize; - let encoded_stop = (block_stop / 8) * bp.bit_width() as usize; + let chunk_start = block_start / PATCH_CHUNK_SIZE; + let chunk_stop = block_stop / PATCH_CHUNK_SIZE; + let encoded_start = widths.byte_offset(chunk_start); + let encoded_stop = widths.byte_offset(chunk_stop); Ok(( bp.packed().slice(encoded_start..encoded_stop), + widths.slice(chunk_start..chunk_stop), u16::try_from(bitpacked_offset)?, patch_range, )) @@ -75,6 +82,7 @@ pub(crate) fn bitpacked_slice_view( impl BitPackedExecutor { fn try_specialize( array: ArrayRef, + ctx: &mut ExecutionCtx, ) -> VortexResult>)>> { if let Ok(array) = array.clone().try_downcast::() { return Ok(Some((array, None))); @@ -91,15 +99,16 @@ impl BitPackedExecutor { let bp = child.as_::(); let offset = slice.data().slice_range().start; let len = array.len(); - let (packed, bitpacked_offset, patch_range) = bitpacked_slice_view(bp, offset, len)?; - let chunk_start = (offset + bp.offset() as usize) / PATCH_CHUNK_SIZE; - let chunk_stop = chunk_start + (len + bitpacked_offset as usize).div_ceil(PATCH_CHUNK_SIZE); + let widths = bp.chunk_layout(ctx)?; + let (packed, widths, bitpacked_offset, patch_range) = + bitpacked_slice_view(bp, &widths, offset, len)?; + let offsets = widths.offsets_array(); let sliced = BitPacked::try_new( packed, bp.ptype(bp.dtype()), child.validity()?.slice(patch_range.clone())?, bp.patches(), - bp.chunk_offsets().slice(chunk_start..chunk_stop + 1)?, + offsets, len, bitpacked_offset, )?; @@ -116,8 +125,8 @@ impl CudaExecute for BitPackedExecutor { array: ArrayRef, ctx: &mut CudaExecutionCtx, ) -> VortexResult { - let (array, patch_range) = - Self::try_specialize(array)?.ok_or_else(|| vortex_err!("Expected BitPackedArray"))?; + let (array, patch_range) = Self::try_specialize(array, ctx.execution_ctx())? + .ok_or_else(|| vortex_err!("Expected BitPackedArray"))?; let ptype = array.ptype(array.dtype()); match_each_integer_ptype!(ptype, |A| { @@ -163,9 +172,9 @@ where A: BitPackedUnpack + NativePType + DeviceRepr + Send + Sync + 'static, A::Physical: DeviceRepr + Send + Sync + 'static, { + let widths = array.chunk_layout(ctx.execution_ctx())?; let BitPackedDataParts { offset, - bit_width, chunk_offsets: _, len, packed, @@ -174,6 +183,9 @@ where } = BitPacked::into_parts(array); vortex_ensure!(len > 0, "Non empty array"); + let Some(bit_width) = widths.uniform_width() else { + vortex_bail!("CUDA bit-unpack requires every chunk to share one bit width"); + }; let offset = offset as usize; let device_input = ctx.ensure_on_device(packed).await?; @@ -628,8 +640,8 @@ mod tests { bitpacked.into_array() }; - let (specialized, patch_range) = - BitPackedExecutor::try_specialize(array)?.vortex_expect("expected BitPacked input"); + let (specialized, patch_range) = BitPackedExecutor::try_specialize(array, &mut ctx)? + .vortex_expect("expected BitPacked input"); assert_eq!(specialized.len(), expected_len); assert_eq!(specialized.offset(), expected_offset);