diff --git a/encodings/fastlanes/benches/bitpack_compare.rs b/encodings/fastlanes/benches/bitpack_compare.rs index 739fcd72dc6..6dde3e13a56 100644 --- a/encodings/fastlanes/benches/bitpack_compare.rs +++ b/encodings/fastlanes/benches/bitpack_compare.rs @@ -59,7 +59,7 @@ fn page_aligned(array: BitPackedArray) -> BitPackedArray { ptype, parts.validity, parts.patches, - parts.bit_width, + parts.widths, parts.len, parts.offset, ) diff --git a/encodings/fastlanes/benches/bitpack_compare_sweep.rs b/encodings/fastlanes/benches/bitpack_compare_sweep.rs index ec7cf9b6892..6bac192a754 100644 --- a/encodings/fastlanes/benches/bitpack_compare_sweep.rs +++ b/encodings/fastlanes/benches/bitpack_compare_sweep.rs @@ -85,7 +85,7 @@ fn page_aligned(array: BitPackedArray) -> BitPackedArray { ptype, parts.validity, parts.patches, - parts.bit_width, + parts.widths, parts.len, parts.offset, ) diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index 4fe91aad4ff..d30eacd301c 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -29,6 +29,8 @@ use vortex_mask::Mask; use crate::BitPacked; use crate::BitPackedArray; +use crate::ChunkWidths; +use crate::FL_CHUNK_SIZE; use crate::bitpack_decompress; pub fn bitpack_to_best_bit_width( @@ -84,7 +86,7 @@ pub fn bitpack_encode( array.ptype(), array.validity()?, patches, - bit_width, + ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)).into_array(), array.len(), 0, )?; @@ -113,7 +115,7 @@ pub unsafe fn bitpack_encode_unchecked( array.ptype(), array.validity()?, None, - bit_width, + ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)).into_array(), array.len(), 0, ) diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 8002320d434..e88fe2b1770 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -14,11 +14,14 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::array_slots; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::patches::PatchSlotIndices; use vortex_array::patches::Patches; @@ -37,9 +40,9 @@ pub mod unpack_iter; use crate::BitPackedArray; use crate::FL_CHUNK_SIZE; -use crate::bitpack_compress::bitpack_encode; -use crate::unpack_iter::BitPacked as BitPackedIter; -use crate::unpack_iter::BitUnpackedChunks; +use crate::bitpacking::bitpack_compress::bitpack_encode; +use crate::bitpacking::unpack_iter::BitPacked as BitPackedIter; +use crate::bitpacking::unpack_iter::BitUnpackedChunks; /// Bytes occupied by one packed FastLanes chunk of `bit_width`-bit values. #[inline] @@ -48,7 +51,7 @@ pub const fn chunk_packed_bytes(bit_width: u8) -> usize { } /// Chunk widths and byte offsets used while encoding or executing bit-packed data. -/// Execution borrows the materialized children; only encoding computes prefix sums. +/// Operations use this view to address each packed chunk. #[derive(Clone, Debug)] pub struct ChunkWidths { widths: Widths, @@ -232,6 +235,9 @@ pub struct BitPackedSlots { /// The validity bitmap indicating which elements are non-null. #[slot(3)] pub validity_child: Option, + /// One non-nullable `u8` width per 1024-element chunk. Uniform widths use a constant array. + #[slot(4)] + pub width_table: ArrayRef, } pub(crate) const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { @@ -240,9 +246,36 @@ pub(crate) const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { chunk_offsets: BitPackedSlots::PATCH_CHUNK_OFFSETS, }; +/// The dtype of the width table child: one byte per chunk. +pub(crate) const WIDTH_TABLE_DTYPE: DType = DType::Primitive(PType::U8, Nullability::NonNullable); + +impl IntoArray for ChunkWidths { + fn into_array(self) -> ArrayRef { + if self.is_uniform() { + ConstantArray::new(self.max_width, self.len()).into_array() + } else { + self.as_buffer().into_array() + } + } +} + +/// Read the width child without executing it during reduction. +pub(crate) fn materialized_widths(table: &ArrayRef) -> VortexResult> { + if let Some(constant) = table.as_opt::() { + return Ok(Some(ChunkWidths::uniform( + u8::try_from(constant.scalar())?, + table.len(), + ))); + } + Ok(table + .as_opt::() + .filter(|a| a.buffer_handle().is_on_host()) + .map(|a| ChunkWidths::new(a.to_buffer::()))) +} + pub struct BitPackedDataParts { pub offset: u16, - pub bit_width: u8, + pub widths: ArrayRef, pub len: usize, pub packed: BufferHandle, pub patches: Option, @@ -254,7 +287,6 @@ pub struct BitPackedData { /// The offset within the first block (created with a slice). /// 0 <= offset < 1024 pub(super) offset: u16, - pub(super) bit_width: u8, pub(super) packed: BufferHandle, /// Patch metadata for reconstructing Patches from slots. pub(super) patches_data: Option, @@ -262,84 +294,67 @@ 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) } } impl BitPackedData { /// Create a new bitpacked array using a buffer of packed data. /// - /// The packed data should be interpreted as a sequence of values with size `bit_width`. - /// - /// # Errors - /// - /// This method returns errors if any of the metadata is inconsistent, for example the packed - /// buffer provided does not have the right size according to the supplied length and target - /// PType. + /// The packed data holds one FastLanes block per 1024-element chunk, each packed at that + /// chunk's width from the width-table child and concatenated in chunk order. The buffer is padded with + /// zeros to the next multiple of 1024 elements if the length is not divisible by 1024. /// /// # Safety /// /// For signed arrays, it is the caller's responsibility to ensure that there are no values - /// that can be interpreted once unpacked to the provided PType. + /// that can be interpreted as negative once unpacked to the provided PType. /// /// This invariant is upheld by the compressor, but callers must ensure this if they wish to /// construct a new `BitPackedArray` from parts. /// /// See also the [`encode`][Self::encode] method on this type for a safe path to create a new /// bit-packed array. - /// A safe constructor for a `BitPackedArray` from its components: - /// - /// * `packed` is ByteBuffer holding the compressed data that was packed with FastLanes - /// bit-packing to a `bit_width` bits per value. `length` is the length of the original - /// vector. Note that the packed is padded with zeros to the next multiple of 1024 elements - /// if `length` is not divisible by 1024. - /// * `ptype` of the original data - /// * `validity` to track any nulls - /// * `patches` optionally provided for values that did not pack - /// - /// Any failure in validation will result in an error. /// /// # Validation /// + /// Performed when the array is built from its parts: + /// /// * The `ptype` must be an integer /// * `validity` must have `length` len /// * Any patches must have any `array_len` equal to `length` - /// * The `packed` buffer must be exactly sized to hold `length` values of `bit_width` rounded - /// up to the next multiple of 1024. + /// * The width-table child must hold one non-nullable `u8` per chunk. + /// + /// Once the widths are materialized, they must be no wider than `ptype`, and the packed + /// buffer must be exactly the sum of the chunks' packed sizes. Compressed children 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 < 1024, - "Offset must be less than the full block i.e., 1024, got {offset}" + (offset as usize) < FL_CHUNK_SIZE, + "Offset must be less than the full block i.e., {FL_CHUNK_SIZE}, got {offset}" ); Ok(Self { offset, - bit_width, packed, patches_data: patches.as_ref().map(PatchesData::from_patches), }) } pub(crate) fn validate( - packed: &BufferHandle, + &self, ptype: PType, validity: &Validity, patches: Option<&Patches>, - bit_width: u8, + table: &ArrayRef, length: usize, - offset: u16, ) -> VortexResult<()> { vortex_ensure!(ptype.is_int(), MismatchedTypes: "integer", ptype); - vortex_ensure!(bit_width <= 64, "Unsupported bit width {bit_width}"); - if let Some(validity_len) = validity.maybe_len() { vortex_ensure!( validity_len == length, @@ -352,16 +367,47 @@ impl BitPackedData { Self::validate_patches(patches, ptype, length)?; } - // Validate packed buffer - let expected_packed_len = - (length + offset as usize).div_ceil(1024) * (128 * bit_width as usize); + let num_chunks = (length + self.offset as usize).div_ceil(FL_CHUNK_SIZE); vortex_ensure!( - packed.len() == expected_packed_len, - "Expected {} packed bytes, got {}", - expected_packed_len, - packed.len() + table.dtype() == &WIDTH_TABLE_DTYPE, + "BitPacked width table must be {WIDTH_TABLE_DTYPE}, got {}", + table.dtype() + ); + vortex_ensure!( + table.len() == num_chunks, + "Expected {num_chunks} chunk widths, got {}", + table.len() ); + // Compressed children are checked once materialized, before any unchecked unpacking. + if let Some(widths) = materialized_widths(table)? { + Self::validate_widths(&self.packed, ptype, &widths)?; + } + Ok(()) + } + pub(crate) fn validate_widths( + packed: &BufferHandle, + ptype: PType, + widths: &ChunkWidths, + ) -> VortexResult<()> { + vortex_ensure!( + widths.max_width() as usize <= ptype.bit_width(), + "Unsupported bit width {} for {ptype}", + widths.max_width() + ); + for chunk in 0..widths.len() { + vortex_ensure!( + widths.byte_offsets[chunk + 1].checked_sub(widths.byte_offsets[chunk]) + == Some(chunk_packed_bytes(widths.width(chunk)) as u64), + "Chunk {chunk} offsets do not match its bit width" + ); + } + let expected = widths.byte_offsets[widths.len()] - widths.byte_offsets[0]; + vortex_ensure!( + packed.len() as u64 == expected, + "Expected {expected} packed bytes, got {}", + packed.len() + ); Ok(()) } @@ -407,13 +453,7 @@ impl BitPackedData { unsafe { std::slice::from_raw_parts(packed_ptr, packed_len) } } - /// Bit-width of the packed values - #[inline] - pub fn bit_width(&self) -> u8 { - self.bit_width - } - - /// Access a chunk using the layout prepared for this operation. + /// The packed FastLanes block of `chunk` as `T` words, along with that chunk's bit width. #[inline] pub(crate) fn packed_chunk( &self, @@ -456,43 +496,54 @@ 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. +pub trait BitPackedArrayExt: BitPackedArraySlotsExt { #[inline] - pub fn max_packed_value(&self) -> usize { - (1 << self.bit_width()) - 1 + fn packed(&self) -> &BufferHandle { + BitPackedData::packed(self) } -} -pub trait BitPackedArrayExt: BitPackedArraySlotsExt { - /// Prepare a chunk layout for a bulk operation. - fn chunk_widths(&self, _ctx: &mut ExecutionCtx) -> VortexResult { - Ok(ChunkWidths::uniform( - self.bit_width(), - (self.as_ref().len() + self.offset() as usize).div_ceil(FL_CHUNK_SIZE), - )) + /// Prepare and validate the width child once for a bulk operation. + fn chunk_widths(&self, ctx: &mut ExecutionCtx) -> VortexResult { + let widths = match materialized_widths(self.width_table())? { + Some(widths) => widths, + None => ChunkWidths::new( + self.width_table() + .clone() + .execute::(ctx)? + .to_buffer::(), + ), + }; + BitPackedData::validate_widths(self.packed(), self.as_ref().dtype().as_ptype(), &widths)?; + Ok(widths) + } + + /// Read and validate widths only when the child is already materialized. + fn materialized_chunk_widths(&self) -> VortexResult> { + let widths = materialized_widths(self.width_table())?; + if let Some(widths) = &widths { + BitPackedData::validate_widths( + self.packed(), + self.as_ref().dtype().as_ptype(), + widths, + )?; + } + Ok(widths) } - /// Locate one chunk without allocating a layout for the entire array. + /// Locate and validate one chunk using the width table. fn chunk_range( &self, chunk: usize, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult<(Range, u8)> { - let len = chunk_packed_bytes(self.bit_width()); - Ok((chunk * len..(chunk + 1) * len, self.bit_width())) - } - - #[inline] - fn packed(&self) -> &BufferHandle { - BitPackedData::packed(self) - } - - #[inline] - fn bit_width(&self) -> u8 { - BitPackedData::bit_width(self) + let widths = self.chunk_widths(ctx)?; + vortex_ensure!(chunk < widths.len(), "Chunk index out of bounds"); + Ok(( + widths.byte_offset(chunk)..widths.byte_offset(chunk + 1), + widths.width(chunk), + )) } #[inline] @@ -520,7 +571,7 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { BitPackedData::packed_slice::(self) } - #[inline] + /// Iterate packed chunks using widths materialized for this operation. fn unpacked_chunks<'a, T: BitPackedIter>( &'a self, widths: &'a ChunkWidths, @@ -611,6 +662,7 @@ mod test { &mut ctx ); } + #[test] fn chunk_widths_offsets() { assert_eq!(ChunkWidths::uniform(3, 3).uniform_width(), Some(3)); diff --git a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs new file mode 100644 index 00000000000..0d425ce8005 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs @@ -0,0 +1,78 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Behavioural tests for bit-packed arrays whose chunks are packed at different widths. + +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Constant; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::BitPacked; +use crate::BitPackedArray; +use crate::BitPackedArrayExt; +use crate::BitPackedArraySlotsExt; +use crate::FL_CHUNK_SIZE; +use crate::bitpacking::bitpack_compress::bitpack_to_best_bit_width; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session +}); + +/// Four full chunks plus a partial trailer, each chunk with a distinctly different magnitude: +/// 3-bit values, 12-bit values with a few 20-bit outliers, all zeros, 20-bit values, and a +/// 5-bit tail. +fn varied(len_tail: usize) -> Vec { + (0..4 * FL_CHUNK_SIZE + len_tail) + .map(|i| { + let chunk = i / FL_CHUNK_SIZE; + let pos = (i % FL_CHUNK_SIZE) as u32; + match chunk { + 0 => pos % 8, + 1 if pos % 300 == 7 => 1 << 20 | pos, + 1 => pos % 4096, + 2 => 0, + 3 => (pos * 977) % (1 << 20), + _ => pos % 32, + } + }) + .collect() +} + +fn encode(values: &[u32]) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + bitpack_to_best_bit_width(&PrimitiveArray::from_iter(values.iter().copied()), &mut ctx) +} + +/// Every array carries a non-nullable `u8` width per chunk, including uniform arrays. +#[test] +fn width_table_is_validated() -> VortexResult<()> { + let packed = encode(&varied(100))?; + let num_chunks = packed + .chunk_widths(&mut SESSION.create_execution_ctx())? + .len(); + let short = PrimitiveArray::from_iter(vec![3u8; num_chunks - 1]).into_array(); + assert!(BitPacked::with_width_table(packed.clone(), short).is_err()); + let wide = PrimitiveArray::from_iter(vec![3u16; num_chunks]).into_array(); + assert!(BitPacked::with_width_table(packed, wide).is_err()); + + let uniform = encode(&(0..3000u32).map(|i| i % 128).collect::>())?; + assert!(uniform.width_table().is::()); + let table = PrimitiveArray::from_iter(vec![ + 7u8; + uniform + .chunk_widths(&mut SESSION.create_execution_ctx())? + .len() + ]) + .into_array(); + let replaced = BitPacked::with_width_table(uniform.clone(), table)?; + assert_arrays_eq!(uniform, replaced, &mut SESSION.create_execution_ctx()); + Ok(()) +} diff --git a/encodings/fastlanes/src/bitpacking/compute/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index cdbb8141cea..8765f415f2f 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -16,6 +16,7 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_array::validity::Validity; use vortex_error::VortexResult; +use crate::BitPackedArraySlotsExt; use crate::bitpacking::BitPacked; use crate::bitpacking::array::BitPackedArrayExt; use crate::bitpacking::array::bitpack_decompress::unpack_map_into_builder; @@ -44,7 +45,7 @@ fn build_with_validity( .patches() .map(|patches| patches.map_values(|values| values.cast(dtype.clone()))) .transpose()?, - array.bit_width(), + array.width_table().clone(), array.len(), array.offset(), )? diff --git a/encodings/fastlanes/src/bitpacking/compute/slice.rs b/encodings/fastlanes/src/bitpacking/compute/slice.rs index 996565a2672..3f7d6314565 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; @@ -14,6 +13,7 @@ use vortex_array::patches::Patches; use vortex_error::VortexResult; use crate::BitPacked; +use crate::BitPackedArraySlotsExt; use crate::bitpacking::array::BitPackedArrayExt; impl SliceReduce for BitPacked { @@ -23,7 +23,12 @@ impl SliceReduce for BitPacked { return Ok(None); } - Ok(Some(slice_bitpacked(array, range, None)?)) + let Some(widths) = array.materialized_chunk_widths()? 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)?)) } } @@ -31,7 +36,7 @@ impl SliceKernel for BitPacked { fn slice( array: ArrayView<'_, Self>, range: Range, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult> { let patches = array .patches() @@ -39,36 +44,41 @@ impl SliceKernel for BitPacked { .transpose()? .flatten(); - Ok(Some(slice_bitpacked(array, range, patches)?)) + let (chunks, _) = slice_chunks(array.offset(), &range); + let widths = array.chunk_widths(ctx)?; + let encoded = widths.byte_offset(chunks.start)..widths.byte_offset(chunks.end); + 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.bit_width(), + array.width_table().slice(chunk_start..chunk_stop)?, 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; @@ -79,7 +89,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/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index ba87ae92f1b..8ce20e5f160 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -29,3 +29,6 @@ pub(crate) fn initialize(session: &vortex_session::VortexSession) { #[cfg(test)] mod serde_tests; + +#[cfg(test)] +mod chunk_widths_tests; diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin.rs index dd66efe44d1..b434ca77442 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin.rs @@ -32,7 +32,10 @@ use vortex_session::VortexSession; use crate::BitPacked; use crate::BitPackedArrayExt; +use crate::BitPackedArraySlotsExt; use crate::BitPackedData; +use crate::ChunkWidths; +use crate::FL_CHUNK_SIZE; use crate::bitpacking::array::BitPackedSlots; /// Metadata of the frozen `fastlanes.bitpacked` wire format. @@ -66,15 +69,18 @@ fn deserialize_v1(parts: ArrayDeserialization<'_>) -> VortexResult { metadata.bit_width ) })?; + let num_chunks = (len + offset as usize).div_ceil(FL_CHUNK_SIZE); let (patches, validity, _) = deserialize_children(children, metadata.patches, dtype, len, 0)?; let slots = { let mut s = ArraySlots::with_capacity(BitPackedSlots::COUNT); PatchesData::push_slots(&mut s, patches.as_ref()); s.push(validity_to_child(&validity, len)); + let widths = ChunkWidths::uniform(bit_width, num_chunks); + s.push(Some(widths.into_array())); 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), )? @@ -158,23 +164,36 @@ impl ArrayPlugin for BitPackedPlugin { fn serialize( &self, array: &ArrayRef, - _session: &VortexSession, + session: &VortexSession, ) -> VortexResult> { let view = array.as_::(); - let metadata = BitPackedMetadata { - bit_width: view.bit_width() as u32, - offset: view.offset() as u32, - patches: view - .patches() - .map(|p| p.to_metadata(view.len(), view.dtype())) - .transpose()?, + let widths = view.chunk_widths(&mut session.create_execution_ctx())?; + vortex_ensure!( + widths.is_uniform(), + "Nonuniform widths require the v2 wire format" + ); + { + let metadata = BitPackedMetadata { + bit_width: widths.max_width() as u32, + offset: view.offset() as u32, + patches: view + .patches() + .map(|p| p.to_metadata(view.len(), view.dtype())) + .transpose()?, + } + .encode_to_vec(); + let children = array.slots()[..BitPackedSlots::WIDTH_TABLE] + .iter() + .flatten() + .cloned() + .collect(); + Ok(Some(ArraySerialization::new( + self.id(), + metadata, + array.buffers(), + children, + ))) } - .encode_to_vec(); - Ok(Some(ArraySerialization::from_array( - self.id(), - array, - metadata, - ))) } fn deserialize( @@ -229,7 +248,7 @@ impl ArrayPlugin for BitPackedPatchedPlugin { let packed = bitpacked.packed().clone(); let ptype = bitpacked.dtype().as_ptype(); let validity = bitpacked.validity()?; - let bw = bitpacked.bit_width; + let bw = bitpacked.width_table().clone(); let len = bitpacked.len(); let offset = bitpacked.offset(); @@ -300,7 +319,7 @@ mod tests { let array = bitpacked.as_array(); let serialization = SESSION.array_serialize(array)?.unwrap(); - let children = array.children(); + let children = serialization.children.clone(); let buffers = array .buffers() .into_iter() @@ -353,7 +372,7 @@ mod tests { let array = bitpacked.as_array(); let serialization = SESSION.array_serialize(array)?.unwrap(); - let children = array.children(); + let children = serialization.children.clone(); let buffers = array .buffers() .into_iter() @@ -386,7 +405,7 @@ mod tests { let array = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); let serialization = SESSION.array_serialize(&array)?.unwrap(); - let children = array.children(); + let children = serialization.children.clone(); let buffers = array .buffers() .into_iter() diff --git a/encodings/fastlanes/src/bitpacking/serde_tests.rs b/encodings/fastlanes/src/bitpacking/serde_tests.rs index 7c8a8d3d952..90940d6d073 100644 --- a/encodings/fastlanes/src/bitpacking/serde_tests.rs +++ b/encodings/fastlanes/src/bitpacking/serde_tests.rs @@ -57,7 +57,7 @@ fn uniform_widths_serialize_as_original_format() -> VortexResult<()> { let values: Vec = (0..3000).map(|i| i % 128).collect(); let packed = bitpack_to_best_bit_width(&PrimitiveArray::from_iter(values.iter().copied()), &mut ctx)?; - assert!(packed.as_array().children().is_empty()); + assert_eq!(packed.as_array().children().len(), 1); assert!( SESSION .array_serialize(packed.as_array())? diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index aa18bd289c6..b879e6a4d40 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -39,13 +39,14 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::BitPackedArrayExt; +use crate::BitPackedArraySlotsExt; use crate::BitPackedData; use crate::BitPackedDataParts; -use crate::bitpack_decompress::unpack_array; -use crate::bitpack_decompress::unpack_into_primitive_builder; use crate::bitpacking::array::BitPackedSlots; use crate::bitpacking::array::BitPackedSlotsView; use crate::bitpacking::array::PATCH_SLOTS; +use crate::bitpacking::bitpack_decompress::unpack_array; +use crate::bitpacking::bitpack_decompress::unpack_into_primitive_builder; use crate::bitpacking::vtable::rules::RULES; mod kernels; mod operations; @@ -62,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); } @@ -71,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 } @@ -95,19 +94,27 @@ impl VTable for BitPacked { len: usize, slots: &[Option], ) -> VortexResult<()> { + vortex_ensure!( + slots.len() == BitPackedSlots::COUNT, + "Expected {} slots, got {}", + BitPackedSlots::COUNT, + slots.len() + ); + vortex_ensure!( + slots[BitPackedSlots::WIDTH_TABLE].is_some(), + "Missing width table or chunk offsets" + ); let bp_slots = BitPackedSlotsView::from_slots(slots); let validity = child_to_validity(bp_slots.validity_child, dtype.nullability()); let patches = PatchesData::patches_from_slots(data.patches_data.as_ref(), len, slots, PATCH_SLOTS); - BitPackedData::validate( - &data.packed, + data.validate( dtype.as_ptype(), &validity, patches.as_ref(), - data.bit_width, + bp_slots.width_table, len, - data.offset, ) } @@ -214,34 +221,51 @@ impl VTable for BitPacked { pub struct BitPacked; impl BitPacked { + /// Build a bit-packed array with one width per chunk. pub fn try_new( packed: BufferHandle, ptype: PType, validity: Validity, patches: Option, - bit_width: u8, + widths: ArrayRef, len: usize, offset: u16, ) -> VortexResult { let dtype = DType::Primitive(ptype, validity.nullability()); let slots = { - let mut s = ArraySlots::with_capacity(4); + let mut s = ArraySlots::with_capacity(BitPackedSlots::COUNT); PatchesData::push_slots(&mut s, patches.as_ref()); s.push(validity_to_child(&validity, len)); + s.push(Some(widths)); 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 width child, dropping statistics that may no longer describe the values. + pub fn with_width_table( + array: BitPackedArray, + table: ArrayRef, + ) -> VortexResult { + let mut slots: ArraySlots = array.slots().iter().cloned().collect(); + slots[BitPackedSlots::WIDTH_TABLE] = Some(table); + let dtype = array.dtype().clone(); + let len = array.len(); + Array::try_from_parts( + ArrayParts::new(BitPacked, dtype, len, array.into_data()).with_slots(slots), + ) + } + pub fn into_parts(array: BitPackedArray) -> BitPackedDataParts { let len = array.len(); let patches = array.patches(); let validity = array.validity().vortex_expect("BitPacked validity"); + let widths = array.width_table().clone(); let data = array.into_data(); BitPackedDataParts { offset: data.offset, - bit_width: data.bit_width, + widths, len, packed: data.packed, patches, diff --git a/encodings/fastlanes/src/bitpacking/vtable/operations.rs b/encodings/fastlanes/src/bitpacking/vtable/operations.rs index 0fe41c2ee3a..00a4f4fab4c 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -56,6 +56,7 @@ mod test { use crate::BitPacked; use crate::BitPackedArray; use crate::BitPackedData; + use crate::ChunkWidths; use crate::bitpacking::array::BitPackedArrayExt; use crate::test::SESSION; @@ -255,7 +256,7 @@ mod test { ) .unwrap(), ), - 1, + ChunkWidths::uniform(1, 1).into_array(), 8, 0, ) diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 5ac7d0e4078..78e63fb6dee 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -97,7 +97,7 @@ impl Scheme for BitPackingScheme { ptype, parts.validity, None, - parts.bit_width, + parts.widths, parts.len, parts.offset, )? @@ -122,7 +122,7 @@ impl Scheme for BitPackingScheme { ptype, parts.validity, parts.patches, - parts.bit_width, + parts.widths, parts.len, parts.offset, )? 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 028dbe81dc9..2dcf421a04e 100644 --- a/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__onpair__string_fsst_structured.snap @@ -3,12 +3,14 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=653785 -root: vortex.onpair(utf8, len=16384) nbytes=139746 +root: vortex.onpair(utf8, len=16384) nbytes=139748 metadata: dict_bytes_len: 11866 dict_offsets: vortex.primitive(u16, len=1626) nbytes=3252 metadata: ptype: u16 - codes: fastlanes.bitpacked(u16, len=63845) nbytes=88704 - metadata: bit_width: 11, offset: 0 + codes: fastlanes.bitpacked(u16, len=63845) nbytes=88706 + metadata: offset: 0 + width_table: vortex.constant(u8, len=63) nbytes=2 + metadata: scalar: 11u8 codes_offsets: vortex.primitive(u16, len=16385) nbytes=32770 metadata: ptype: u16 uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154 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 788cf54ecd8..6d1b6de8c18 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__binary_low_cardinality.snap @@ -3,10 +3,12 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: binary, len=16384, nbytes=315856 -root: vortex.dict(binary, len=16384) nbytes=6196 +root: vortex.dict(binary, len=16384) nbytes=6198 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6146 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 3u8 values: vortex.varbin(binary, len=5) nbytes=52 metadata: offsets: vortex.primitive(u8, len=6) nbytes=6 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap index 6eb4ccfed8d..98b05b7e4e1 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap @@ -3,7 +3,9 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: decimal(12,2), len=16384, nbytes=131072 -root: vortex.decimal_byte_parts.v2(decimal(12,2), len=16384) nbytes=49152 +root: vortex.decimal_byte_parts.v2(decimal(12,2), len=16384) nbytes=49154 metadata: - msp: fastlanes.bitpacked(i32, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 + msp: fastlanes.bitpacked(i32, len=16384) nbytes=49154 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 24u8 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 f70d90a301d..881812695ba 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__float_alp_prices.snap @@ -3,7 +3,9 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: f64, len=16384, nbytes=131072 -root: vortex.alp(f64, len=16384) nbytes=49152 +root: vortex.alp(f64, len=16384) nbytes=49154 metadata: exponents: e: 14, f: 12 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49154 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 24u8 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 487e7c7900b..7dee3501126 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__float_full_precision.snap @@ -3,12 +3,16 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: f64, len=16384, nbytes=131072 -root: vortex.alprd(f64, len=16384) nbytes=112880 +root: vortex.alprd(f64, len=16384) nbytes=112884 metadata: right_bit_width: 52, patch_offset: 0 - left_parts: fastlanes.bitpacked(u16, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 - right_parts: fastlanes.bitpacked(u64, len=16384) nbytes=106496 - metadata: bit_width: 52, offset: 0 + left_parts: fastlanes.bitpacked(u16, len=16384) nbytes=6146 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 3u8 + right_parts: fastlanes.bitpacked(u64, len=16384) nbytes=106498 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 52u8 patch_indices: vortex.primitive(u16, len=60) nbytes=120 metadata: ptype: u16 patch_values: 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 9c2192c70d7..dfe78f90377 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__float_low_cardinality.snap @@ -3,11 +3,13 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: f64, len=16384, nbytes=131072 -root: vortex.alp(f64, len=16384) nbytes=6208 +root: vortex.alp(f64, len=16384) nbytes=6210 metadata: exponents: e: 16, f: 11 - encoded: vortex.dict(i64, len=16384) nbytes=6208 + encoded: vortex.dict(i64, len=16384) nbytes=6210 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6146 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 3u8 values: vortex.primitive(i64, len=8) nbytes=64 metadata: ptype: i64 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 8648200a138..58df4406720 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_low_cardinality.snap @@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: i64, len=16384, nbytes=131072 -root: vortex.dict(i64, len=16384) nbytes=6192 +root: vortex.dict(i64, len=16384) nbytes=6194 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=6144 - metadata: bit_width: 3, offset: 0 + codes: fastlanes.bitpacked(u8, len=16384) nbytes=6146 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 3u8 values: vortex.primitive(i64, len=6) nbytes=48 metadata: ptype: i64 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 4800aae289d..531b01ee424 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_monotone_jitter.snap @@ -3,7 +3,9 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: u64, len=16384, nbytes=131072 -root: fastlanes.for(u64, len=16384) nbytes=49152 +root: fastlanes.for(u64, len=16384) nbytes=49154 metadata: reference: 1700000001036u64 - encoded: fastlanes.bitpacked(u64, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 + encoded: fastlanes.bitpacked(u64, len=16384) nbytes=49154 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 24u8 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 73566369279..c0ba7ca5240 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_mostly_null.snap @@ -3,11 +3,13 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: i32?, len=16384, nbytes=67584 -root: vortex.sparse(i32?, len=16384) nbytes=3031 +root: vortex.sparse(i32?, len=16384) nbytes=3033 metadata: fill_value: null patch_indices: vortex.primitive(u16, len=823) nbytes=1646 metadata: ptype: u16 - patch_values: fastlanes.bitpacked(i32?, len=823) nbytes=1383 - metadata: bit_width: 10, offset: 0 + patch_values: fastlanes.bitpacked(i32?, len=823) nbytes=1385 + metadata: offset: 0 validity_child: vortex.bool(bool, len=823) nbytes=103 metadata: offset: 0 + width_table: vortex.constant(u8, len=1) nbytes=2 + metadata: scalar: 10u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap index 24949893190..761958a68a6 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_negatives.snap @@ -3,7 +3,9 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: i64, len=16384, nbytes=131072 -root: fastlanes.for(i64, len=16384) nbytes=16384 +root: fastlanes.for(i64, len=16384) nbytes=16386 metadata: reference: -128i64 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=16384 - metadata: bit_width: 8, offset: 0 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=16386 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 8u8 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap b/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap index 271de8faf93..2f3e237bc26 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_runs.snap @@ -3,13 +3,17 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: i32, len=16384, nbytes=65536 -root: vortex.runend(i32, len=16384) nbytes=3968 +root: vortex.runend(i32, len=16384) nbytes=3972 metadata: offset: 0 - ends: fastlanes.for(u16, len=1020) nbytes=1792 + ends: fastlanes.for(u16, len=1020) nbytes=1794 metadata: reference: 13u16 - encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 - metadata: bit_width: 14, offset: 0 - values: fastlanes.for(i32, len=1020) nbytes=2176 + encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1794 + metadata: offset: 0 + width_table: vortex.constant(u8, len=1) nbytes=2 + metadata: scalar: 14u8 + values: fastlanes.for(i32, len=1020) nbytes=2178 metadata: reference: -49931i32 - encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 - metadata: bit_width: 17, offset: 0 + encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2178 + metadata: offset: 0 + width_table: vortex.constant(u8, len=1) nbytes=2 + metadata: scalar: 17u8 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 53a5d1e023c..f2bb58d7016 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__int_sparse_outliers.snap @@ -3,11 +3,13 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: i64, len=16384, nbytes=131072 -root: vortex.sparse(i64, len=16384) nbytes=5540 +root: vortex.sparse(i64, len=16384) nbytes=5542 metadata: fill_value: 1000000i64 patch_indices: vortex.primitive(u16, len=848) nbytes=1696 metadata: ptype: u16 - patch_values: fastlanes.for(i64, len=848) nbytes=3840 + patch_values: fastlanes.for(i64, len=848) nbytes=3842 metadata: reference: 1000830099i64 - encoded: fastlanes.bitpacked(i64, len=848) nbytes=3840 - metadata: bit_width: 30, offset: 0 + encoded: fastlanes.bitpacked(i64, len=848) nbytes=3842 + metadata: offset: 0 + width_table: vortex.constant(u8, len=1) nbytes=2 + metadata: scalar: 30u8 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 c9554add05a..bd235627afc 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 @@ -3,23 +3,29 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: list(i32), len=4066, nbytes=81804 -root: vortex.list(list(i32), len=4066) nbytes=11146 +root: vortex.list(list(i32), len=4066) nbytes=11152 metadata: - elements: vortex.runend(i32, len=16384) nbytes=3968 + elements: vortex.runend(i32, len=16384) nbytes=3972 metadata: offset: 0 - ends: fastlanes.for(u16, len=1020) nbytes=1792 + ends: fastlanes.for(u16, len=1020) nbytes=1794 metadata: reference: 13u16 - encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1792 - metadata: bit_width: 14, offset: 0 - values: fastlanes.for(i32, len=1020) nbytes=2176 + encoded: fastlanes.bitpacked(u16, len=1020) nbytes=1794 + metadata: offset: 0 + width_table: vortex.constant(u8, len=1) nbytes=2 + metadata: scalar: 14u8 + values: fastlanes.for(i32, len=1020) nbytes=2178 metadata: reference: -49931i32 - encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2176 - metadata: bit_width: 17, offset: 0 - offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7178 - metadata: bit_width: 14, offset: 0 + encoded: fastlanes.bitpacked(i32, len=1020) nbytes=2178 + metadata: offset: 0 + width_table: vortex.constant(u8, len=1) nbytes=2 + metadata: scalar: 17u8 + offsets: fastlanes.bitpacked(u16, len=4067) nbytes=7180 + metadata: offset: 0 patch_indices: vortex.primitive(u16, len=1) nbytes=2 metadata: ptype: u16 patch_values: vortex.constant(u16, len=1) nbytes=4 metadata: scalar: 16384u16 patch_chunk_offsets: vortex.primitive(u8, len=4) nbytes=4 metadata: ptype: u8 + width_table: vortex.constant(u8, len=4) nbytes=2 + metadata: scalar: 14u8 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 327f050b0d7..1fb4bb8092b 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__string_fsst_structured.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=653785 -root: vortex.fsst(utf8, len=16384) nbytes=151382 +root: vortex.fsst(utf8, len=16384) nbytes=151384 metadata: len: 16384, nsymbols: 223 uncompressed_lengths: vortex.sparse(u8, len=16384) nbytes=3154 metadata: fill_value: 24u8 @@ -11,5 +11,7 @@ root: vortex.fsst(utf8, len=16384) nbytes=151382 metadata: ptype: u16 patch_values: vortex.constant(u8, len=1575) nbytes=2 metadata: scalar: 23u8 - codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36992 - metadata: bit_width: 17, offset: 0 + codes_offsets: fastlanes.bitpacked(u32, len=16385) nbytes=36994 + metadata: offset: 0 + width_table: vortex.constant(u8, len=17) nbytes=2 + metadata: scalar: 17u8 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 ec79989d46a..fa867d2d4b4 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__string_low_cardinality.snap @@ -3,10 +3,12 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: utf8, len=16384, nbytes=262144 -root: vortex.dict(utf8, len=16384) nbytes=8373 +root: vortex.dict(utf8, len=16384) nbytes=8375 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 - metadata: bit_width: 4, offset: 0 + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8194 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 4u8 values: vortex.fsst(utf8, len=12) nbytes=181 metadata: len: 12, nsymbols: 9 uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap b/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap index be8f48b779b..68302be477d 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__struct_mixed.snap @@ -3,21 +3,25 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: {id=i64, category=utf8, value=f64}, len=16384, nbytes=524288 -root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=57525 +root: vortex.struct({id=i64, category=utf8, value=f64}, len=16384) nbytes=57529 metadata: id: vortex.sequence(i64, len=16384) nbytes=0 metadata: base: 10000i64, multiplier: 7i64 - category: vortex.dict(utf8, len=16384) nbytes=8373 + category: vortex.dict(utf8, len=16384) nbytes=8375 metadata: all_values_referenced: true - codes: fastlanes.bitpacked(u8, len=16384) nbytes=8192 - metadata: bit_width: 4, offset: 0 + codes: fastlanes.bitpacked(u8, len=16384) nbytes=8194 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 4u8 values: vortex.fsst(utf8, len=12) nbytes=181 metadata: len: 12, nsymbols: 9 uncompressed_lengths: vortex.primitive(u8, len=12) nbytes=12 metadata: ptype: u8 codes_offsets: vortex.primitive(u8, len=13) nbytes=13 metadata: ptype: u8 - value: vortex.alp(f64, len=16384) nbytes=49152 + value: vortex.alp(f64, len=16384) nbytes=49154 metadata: exponents: e: 14, f: 12 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49152 - metadata: bit_width: 24, offset: 0 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=49154 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 24u8 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 6f38e8e6f5d..a9dc41403d7 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__temporal_timestamp_micros.snap @@ -3,9 +3,11 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: vortex.timestamp[µs, tz=UTC](i64), len=16384, nbytes=131072 -root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=67584 +root: vortex.ext(vortex.timestamp[µs, tz=UTC](i64), len=16384) nbytes=67586 metadata: - storage: fastlanes.for(i64, len=16384) nbytes=67584 + storage: fastlanes.for(i64, len=16384) nbytes=67586 metadata: reference: 1700000000891673i64 - encoded: fastlanes.bitpacked(i64, len=16384) nbytes=67584 - metadata: bit_width: 33, offset: 0 + encoded: fastlanes.bitpacked(i64, len=16384) nbytes=67586 + metadata: offset: 0 + width_table: vortex.constant(u8, len=16) nbytes=2 + metadata: scalar: 33u8 diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index 43f6a5a3712..bb491b22d4d 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_widths()? 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_widths()? + .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 86b7a88b276..f72702fd16e 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; @@ -26,8 +28,10 @@ use vortex::encodings::fastlanes::BitPacked; use vortex::encodings::fastlanes::BitPackedArray; use vortex::encodings::fastlanes::BitPackedArrayExt; use vortex::encodings::fastlanes::BitPackedDataParts; +use vortex::encodings::fastlanes::ChunkWidths; 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; @@ -51,9 +55,10 @@ pub(crate) struct BitPackedExecutor; /// materialization so exception metadata is sliced consistently. pub(crate) fn bitpacked_slice_view( bp: ArrayView<'_, BitPacked>, + widths: &ChunkWidths, offset: usize, len: usize, -) -> VortexResult<(BufferHandle, u16, Range)> { +) -> VortexResult<(BufferHandle, ChunkWidths, u16, Range)> { let patch_range = offset..offset + len; let offset_start = patch_range.start + bp.offset() as usize; let offset_stop = offset_start + len; @@ -61,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, )) @@ -74,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))); @@ -90,13 +99,15 @@ 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 widths = bp.chunk_widths(ctx)?; + let (packed, widths, bitpacked_offset, patch_range) = + bitpacked_slice_view(bp, &widths, offset, len)?; let sliced = BitPacked::try_new( packed, bp.ptype(bp.dtype()), child.validity()?.slice(patch_range.clone())?, bp.patches(), - bp.bit_width(), + widths.into_array(), len, bitpacked_offset, )?; @@ -113,8 +124,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| { @@ -160,9 +171,10 @@ where A: BitPackedUnpack + NativePType + DeviceRepr + Send + Sync + 'static, A::Physical: DeviceRepr + Send + Sync + 'static, { + let widths = array.chunk_widths(ctx.execution_ctx())?; let BitPackedDataParts { offset, - bit_width, + widths: _, len, packed, patches, @@ -170,6 +182,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?; @@ -624,8 +639,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);