From 2756c3406a5498ae453e45a532f4e7b1afe8c341 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 3 Sep 2026 15:15:43 -0400 Subject: [PATCH 1/2] feat(fastlanes): store a bit width per FastLanes chunk in BitPacked Each 1024-element chunk of a BitPacked array is packed at its own width, chosen from that chunk's histogram by charging the exact packed block plus its exceptions. Uniform widths still serialize as fastlanes.bitpacked, byte for byte. Differing widths serialize as fastlanes.bitpacked_v2, carrying the widths in a u8 child so metadata stays bounded and a compressor can re-encode the table; that format's metadata is offset and patches only. Following the editions model, BitPackedPlugin owns both wire formats. The encoding's own serializer is labelled with the original ID, so it returns None for differing widths rather than emitting bytes that ID's readers would misread. The compressor still picks one global width, so its output is unchanged. CUDA decodes uniform arrays and rejects differing widths until it has a per-chunk kernel. Signed-off-by: Matt Katz --- .../fastlanes/benches/bitpack_compare.rs | 2 +- .../benches/bitpack_compare_sweep.rs | 2 +- .../src/bitpacking/array/bitpack_compress.rs | 764 ++++++++++++------ .../bitpacking/array/bitpack_decompress.rs | 19 +- .../fastlanes/src/bitpacking/array/mod.rs | 328 +++++++- .../src/bitpacking/array/unpack_iter.rs | 154 ++-- .../src/bitpacking/chunk_widths_tests.rs | 627 ++++++++++++++ .../src/bitpacking/compute/between.rs | 5 +- .../fastlanes/src/bitpacking/compute/cast.rs | 2 +- .../src/bitpacking/compute/compare.rs | 2 +- .../src/bitpacking/compute/compare_fused.rs | 19 +- .../src/bitpacking/compute/filter.rs | 6 +- .../src/bitpacking/compute/is_constant.rs | 2 +- .../fastlanes/src/bitpacking/compute/mod.rs | 2 +- .../fastlanes/src/bitpacking/compute/slice.rs | 11 +- .../bitpacking/compute/stream_predicate.rs | 2 +- .../fastlanes/src/bitpacking/compute/take.rs | 8 +- encodings/fastlanes/src/bitpacking/mod.rs | 6 + encodings/fastlanes/src/bitpacking/plugin.rs | 185 ++++- .../fastlanes/src/bitpacking/vtable/mod.rs | 194 +++-- .../src/bitpacking/vtable/operations.rs | 5 +- .../fastlanes/src/for/array/for_decompress.rs | 2 +- encodings/fastlanes/src/lib.rs | 2 +- .../src/schemes/integer/bitpacking.rs | 4 +- .../src/dynamic_dispatch/plan_builder.rs | 13 +- vortex-cuda/src/kernel/encodings/bitpacked.rs | 22 +- 26 files changed, 1914 insertions(+), 474 deletions(-) create mode 100644 encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs 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 a393db6ecc8..bc4b8aae80e 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -4,6 +4,7 @@ use fastlanes::BitPacking; use itertools::Itertools; use num_traits::PrimInt; +use num_traits::Zero; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -14,82 +15,311 @@ use vortex_array::buffer::BufferHandle; use vortex_array::dtype::IntegerPType; use vortex_array::dtype::NativePType; use vortex_array::dtype::PType; +use vortex_array::dtype::PhysicalPType; use vortex_array::match_each_integer_ptype; use vortex_array::match_each_unsigned_integer_ptype; use vortex_array::patches::Patches; use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_mask::AllOr; use vortex_mask::Mask; use crate::BitPacked; use crate::BitPackedArray; -use crate::bitpack_decompress; +use crate::FL_CHUNK_SIZE; +use crate::bitpack_decompress::count_exceptions; +use crate::bitpacking::array::ChunkWidths; +use crate::bitpacking::array::chunk_packed_bytes; -pub fn bitpack_to_best_bit_width( +/// Bit-pack an array choosing the cost-model-optimal width for every 1024-element chunk. +/// +/// Each chunk is charged for its packed block plus the exceptions left behind, so a chunk of small +/// values stays narrow no matter how wide its neighbours are. +/// +/// Every chunk is processed in one go while it sits in L1: histogram, width choice, exception +/// gathering and packing, so the values are streamed from memory once. +pub fn bitpack_to_best_chunk_widths( array: &PrimitiveArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let bit_width_freq = bit_width_histogram(array.as_view(), ctx)?; - let best_bit_width = find_best_bit_width(array.ptype(), &bit_width_freq)?; - bitpack_encode(array, best_bit_width, Some(&bit_width_freq), ctx) + ensure_non_negative(array, ctx)?; + let validity = array.validity()?; + let mask = validity.execute_mask(array.len(), ctx)?; + let patch_validity = match validity { + Validity::NonNullable => Validity::NonNullable, + _ => Validity::AllValid, + }; + + let len = array.len(); + let (widths, packed, patches) = match_each_integer_ptype!(array.ptype(), |T| { + encode_chunks::(array.as_slice::(), &mask, patch_validity)? + }); + + let bitpacked = BitPacked::try_new( + BufferHandle::new_host(packed), + array.ptype(), + validity, + patches, + widths, + len, + 0, + )?; + bitpacked.statistics().inherit_from(array.statistics()); + Ok(bitpacked) } -#[expect(unused_comparisons, clippy::absurd_extreme_comparisons)] -pub fn bitpack_encode( +/// Multi-pass reference for [`bitpack_to_best_chunk_widths`]: one pass to choose widths, one to +/// pack, and one to gather exceptions. Kept to check the fused encoder against. +#[cfg(test)] +pub(crate) fn bitpack_to_best_chunk_widths_multipass( array: &PrimitiveArray, - bit_width: u8, - bit_width_freq: Option<&[usize]>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let bit_width_freq = match bit_width_freq { - Some(freq) => freq, - None => &bit_width_histogram(array.as_view(), ctx)?, + let plan = chunk_width_plan(array.as_view(), ctx)?; + bitpack_encode_planned(array, plan, ctx) +} + +/// Histogram, choose a width, gather exceptions and pack, one 1024-element chunk at a time. +/// +/// Patch indices use the narrowest unsigned type that can address every value. +fn encode_chunks( + values: &[T], + mask: &Mask, + patch_validity: Validity, +) -> VortexResult<(ChunkWidths, ByteBuffer, Option)> +where + T: NativePType + PrimInt + PhysicalPType, + T::Physical: BitPacking + NativePType, +{ + let len = values.len(); + if len < u8::MAX as usize { + encode_chunks_indexed::(values, mask, patch_validity) + } else if len < u16::MAX as usize { + encode_chunks_indexed::(values, mask, patch_validity) + } else if len < u32::MAX as usize { + encode_chunks_indexed::(values, mask, patch_validity) + } else { + encode_chunks_indexed::(values, mask, patch_validity) + } +} + +fn encode_chunks_indexed( + values: &[T], + mask: &Mask, + patch_validity: Validity, +) -> VortexResult<(ChunkWidths, ByteBuffer, Option)> +where + T: NativePType + PrimInt + PhysicalPType, + T::Physical: BitPacking + NativePType, + P: IntegerPType, +{ + let bits = T::PTYPE.bit_width(); + let bytes_per_exception = bytes_per_exception(T::PTYPE); + let num_chunks = values.len().div_ceil(FL_CHUNK_SIZE); + + let mut widths = BufferMut::::with_capacity(num_chunks); + let mut chunk_offsets = BufferMut::::with_capacity(num_chunks); + let mut indices = BufferMut::

::empty(); + let mut patch_values = BufferMut::::empty(); + + let validity = match mask.bit_buffer() { + AllOr::All => None, + AllOr::Some(bits) => Some(bits), + // Every value is null: every chunk is zero-width and there is nothing to pack. + AllOr::None => { + widths.extend_trusted(std::iter::repeat_n(0u8, num_chunks)); + chunk_offsets.extend_trusted(std::iter::repeat_n(0u64, num_chunks)); + return Ok((ChunkWidths::new(widths.freeze()), ByteBuffer::empty(), None)); + } }; - // Check array contains no negative values. - if array.ptype().is_signed_int() { - let has_negative_values = match_each_integer_ptype!(array.ptype(), |P| { - array.statistics().compute_min::

(ctx).unwrap_or_default() < 0 + // Every chunk packs into a whole block, so the padded size bounds the output; a short trailing + // chunk can pack to more than its raw size. The buffer is shrunk to its exact size at the end. + let mut packed = BufferMut::::with_capacity(num_chunks * FL_CHUNK_SIZE); + let mut histogram = vec![0usize; bits + 1]; + // Zero-padded copy of the trailing partial chunk. + let mut padded = [T::Physical::zero(); FL_CHUNK_SIZE]; + + for (chunk_idx, chunk) in values.chunks(FL_CHUNK_SIZE).enumerate() { + let base = chunk_idx * FL_CHUNK_SIZE; + let chunk_validity = validity.map(|v| v.slice(base..base + chunk.len())); + + histogram.fill(0); + for_each_valid_width(chunk, chunk_validity.as_ref(), |_, _, width| { + histogram[width] += 1; }); - if has_negative_values { - vortex_bail!(InvalidArgument: "cannot bitpack_encode array containing negative integers") + + let bit_width = best_chunk_width(&histogram, bytes_per_exception); + widths.push(bit_width); + chunk_offsets.push(patch_values.len() as u64); + + // The chunk is still in L1, so a second walk over it is cheaper than remembering widths. + if count_exceptions(bit_width, &histogram) > 0 { + for_each_valid_width(chunk, chunk_validity.as_ref(), |i, value, width| { + if width > bit_width as usize { + indices.push(P::from(base + i).vortex_expect("cast index from usize")); + patch_values.push(value); + } + }); + } + + if bit_width > 0 { + let input: &[T::Physical] = if chunk.len() == FL_CHUNK_SIZE { + as_physical(chunk) + } else { + padded[..chunk.len()].copy_from_slice(as_physical(chunk)); + &padded + }; + let packed_len = chunk_packed_bytes(bit_width) / size_of::(); + let start = packed.len(); + // SAFETY: `input` holds exactly 1024 values and the output window is exactly one + // packed block at `bit_width`, within the raw-size capacity reserved above. + unsafe { + packed.set_len(start + packed_len); + BitPacking::unchecked_pack(bit_width as usize, input, &mut packed[start..]); + } + } + } + + let packed = if packed.len() < packed.capacity() { + let mut exact = BufferMut::::with_capacity(packed.len()); + exact.extend_from_slice(&packed); + exact.freeze() + } else { + packed.freeze() + }; + + let patches = if indices.is_empty() { + None + } else { + Some(Patches::new( + values.len(), + 0, + indices.into_array(), + PrimitiveArray::new(patch_values, patch_validity).into_array(), + Some(chunk_offsets.into_array()), + )?) + }; + + Ok(( + ChunkWidths::new(widths.freeze()), + packed.into_byte_buffer(), + patches, + )) +} + +/// Call `f(index, value, bit_width)` for every value of `chunk`; nulls report a width of zero. +#[inline] +fn for_each_valid_width( + chunk: &[T], + validity: Option<&BitBuffer>, + mut f: impl FnMut(usize, T, usize), +) { + let bits = T::PTYPE.bit_width(); + match validity { + None => { + for (i, &v) in chunk.iter().enumerate() { + f(i, v, bits - PrimInt::leading_zeros(v) as usize); + } + } + Some(validity) => { + for ((i, &v), valid) in chunk.iter().enumerate().zip(validity.iter()) { + let width = if valid { + bits - PrimInt::leading_zeros(v) as usize + } else { + 0 + }; + f(i, v, width); + } } } +} - let num_exceptions = bitpack_decompress::count_exceptions(bit_width, bit_width_freq); +/// View signed or unsigned values as their unsigned physical twin, which FastLanes packs. +fn as_physical(values: &[T]) -> &[T::Physical] { + const { + assert!(size_of::() == size_of::()); + assert!(align_of::() == align_of::()); + } + // SAFETY: `Physical` is the same-width unsigned integer, so the layouts match exactly. + unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) } +} - if bit_width >= array.ptype().bit_width() as u8 { - // Nothing we can do +/// The cost-model-optimal bit width of every 1024-element chunk of `array`. +pub fn best_chunk_widths( + array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + Ok(chunk_width_plan(array, ctx)?.widths) +} + +/// Bit-pack `array` at the given per-chunk widths, gathering values that do not fit their chunk's +/// width into patches. +pub fn bitpack_encode_with_widths( + array: &PrimitiveArray, + widths: ChunkWidths, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let num_chunks = array.len().div_ceil(FL_CHUNK_SIZE); + vortex_ensure!( + widths.len() == num_chunks, + "Expected {num_chunks} chunk widths for {} values, got {}", + array.len(), + widths.len() + ); + let plan = ChunkWidthPlan { + widths, + num_exceptions: None, + }; + bitpack_encode_planned(array, plan, ctx) +} + +/// Bit-pack `array` at the single best global width chosen by [`find_best_bit_width`]. +/// +/// Every chunk shares that width, so the result serializes under the original +/// `fastlanes.bitpacked` format. See [`bitpack_to_best_chunk_widths`] for per-chunk widths. +pub fn bitpack_to_best_bit_width( + array: &PrimitiveArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let bit_width_freq = bit_width_histogram(array.as_view(), ctx)?; + let best_bit_width = find_best_bit_width(array.ptype(), &bit_width_freq)?; + bitpack_encode(array, best_bit_width, Some(&bit_width_freq), ctx) +} + +/// Bit-pack every chunk of `array` at the same `bit_width`, which must be narrower than the type. +/// +/// `bit_width_freq` is the array's bit-width histogram if already known; it saves recomputing it +/// to count exceptions. +pub fn bitpack_encode( + array: &PrimitiveArray, + bit_width: u8, + bit_width_freq: Option<&[usize]>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if bit_width as usize >= array.ptype().bit_width() { vortex_bail!( InvalidArgument: "Cannot pack - specified bit width {bit_width} >= {}", array.ptype().bit_width() ) } - - // SAFETY: we check that array only contains non-negative values. - let packed = unsafe { bitpack_unchecked(array, bit_width) }; - let patches = (num_exceptions > 0) - .then(|| gather_patches(array, bit_width, num_exceptions, ctx)) - .transpose()? - .flatten(); - - let bitpacked = BitPacked::try_new( - BufferHandle::new_host(packed), - array.ptype(), - array.validity()?, - patches, - bit_width, - array.len(), - 0, - )?; - bitpacked.statistics().inherit_from(array.statistics()); - Ok(bitpacked) + let num_exceptions = bit_width_freq.map(|freq| count_exceptions(bit_width, freq)); + let widths = ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)); + bitpack_encode_planned( + array, + ChunkWidthPlan { + widths, + num_exceptions, + }, + ctx, + ) } /// Bitpack an array into the specified bit-width without checking statistics. @@ -104,8 +334,9 @@ pub unsafe fn bitpack_encode_unchecked( array: PrimitiveArray, bit_width: u8, ) -> VortexResult { + let widths = ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)); // SAFETY: non-negativity of input checked by caller. - let packed = unsafe { bitpack_unchecked(&array, bit_width) }; + let packed = unsafe { bitpack_unchecked_with_widths(&array, &widths) }; let arr_ref = array.clone().into_array(); let bitpacked = BitPacked::try_new( @@ -113,7 +344,7 @@ pub unsafe fn bitpack_encode_unchecked( array.ptype(), array.validity()?, None, - bit_width, + widths, array.len(), 0, ) @@ -134,69 +365,211 @@ pub unsafe fn bitpack_encode_unchecked( /// It is the caller's responsibility to ensure that `parray` is non-negative before calling /// this function. pub unsafe fn bitpack_unchecked(parray: &PrimitiveArray, bit_width: u8) -> ByteBuffer { - let parray = parray.reinterpret_cast(parray.ptype().to_unsigned()); - match_each_unsigned_integer_ptype!(parray.ptype(), |P| { - bitpack_primitive(parray.as_slice::

(), bit_width).into_byte_buffer() - }) + let widths = ChunkWidths::uniform(bit_width, parray.len().div_ceil(FL_CHUNK_SIZE)); + // SAFETY: forwarded to the caller. + unsafe { bitpack_unchecked_with_widths(parray, &widths) } } /// Bitpack a slice of primitives down to the given width. /// /// See `bitpack` for more caller information. pub fn bitpack_primitive(array: &[T], bit_width: u8) -> Buffer { - if bit_width == 0 { - return Buffer::::empty(); + let widths = ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)); + bitpack_primitive_chunked(array, &widths) +} + +/// Gather the values that do not fit `bit_width` into patches. +pub fn gather_patches( + parray: &PrimitiveArray, + bit_width: u8, + num_exceptions_hint: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let widths = ChunkWidths::uniform(bit_width, parray.len().div_ceil(FL_CHUNK_SIZE)); + gather_patches_with_widths(parray, &widths, num_exceptions_hint, ctx) +} + +/// Chosen chunk widths plus, when known, how many values do not fit them. +struct ChunkWidthPlan { + widths: ChunkWidths, + num_exceptions: Option, +} + +fn chunk_width_plan( + array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match_each_integer_ptype!(array.ptype(), |P| { + chunk_width_plan_typed::

(array, ctx) + }) +} + +fn chunk_width_plan_typed( + array: ArrayView<'_, Primitive>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let bytes_per_exception = bytes_per_exception(T::PTYPE); + let values = array.as_slice::(); + let num_chunks = values.len().div_ceil(FL_CHUNK_SIZE); + let bit_width: fn(T) -> usize = + |v: T| (8 * size_of::()) - (PrimInt::leading_zeros(v) as usize); + + let mut widths = BufferMut::::with_capacity(num_chunks); + let mut num_exceptions = 0usize; + let mut histogram = vec![0usize; size_of::() * 8 + 1]; + + // Score one chunk's histogram and reset it for the next chunk. + let mut finish_chunk = |histogram: &mut [usize]| -> u8 { + let best = best_chunk_width(histogram, bytes_per_exception); + num_exceptions += count_exceptions(best, histogram); + histogram.fill(0); + best + }; + + match array + .validity()? + .execute_mask(array.as_ref().len(), ctx)? + .bit_buffer() + { + AllOr::All => { + for chunk in values.chunks(FL_CHUNK_SIZE) { + for v in chunk { + histogram[bit_width(*v)] += 1; + } + widths.push(finish_chunk(&mut histogram)); + } + } + AllOr::None => { + for _ in 0..num_chunks { + widths.push(0); + } + } + AllOr::Some(buffer) => { + let mut valid = buffer.iter(); + for chunk in values.chunks(FL_CHUNK_SIZE) { + for v in chunk { + if valid.next().unwrap_or(true) { + histogram[bit_width(*v)] += 1; + } else { + histogram[0] += 1; + } + } + widths.push(finish_chunk(&mut histogram)); + } + } } - let bit_width = bit_width as usize; - - // How many fastlanes vectors we will process. - let num_chunks = array.len().div_ceil(1024); - let num_full_chunks = array.len() / 1024; - let packed_len = 128 * bit_width / size_of::(); - // packed_len says how many values of size T we're going to include. - // 1024 * bit_width / 8 == the number of bytes we're going to get. - // then we divide by the size of T to get the number of elements. - - // Allocate a result byte array. - let mut output = BufferMut::::with_capacity(num_chunks * packed_len); - - // Loop over all but the last chunk. - (0..num_full_chunks).for_each(|i| { - let start_elem = i * 1024; - let output_len = output.len(); - unsafe { - output.set_len(output_len + packed_len); - BitPacking::unchecked_pack( - bit_width, - &array[start_elem..][..1024], - &mut output[output_len..][..packed_len], - ); - }; - }); - // Pad the last chunk with zeros to a full 1024 elements. - if num_chunks != num_full_chunks { - let last_chunk_size = array.len() % 1024; - let mut last_chunk: [T; 1024] = [T::zero(); 1024]; - last_chunk[..last_chunk_size].copy_from_slice(&array[array.len() - last_chunk_size..]); + Ok(ChunkWidthPlan { + widths: ChunkWidths::new(widths.freeze()), + num_exceptions: Some(num_exceptions), + }) +} + +fn bitpack_encode_planned( + array: &PrimitiveArray, + plan: ChunkWidthPlan, + ctx: &mut ExecutionCtx, +) -> VortexResult { + ensure_non_negative(array, ctx)?; + let ChunkWidthPlan { + widths, + num_exceptions, + } = plan; + + // SAFETY: we check that array only contains non-negative values. + let packed = unsafe { bitpack_unchecked_with_widths(array, &widths) }; + let patches = if num_exceptions == Some(0) { + None + } else { + gather_patches_with_widths(array, &widths, num_exceptions.unwrap_or(0), ctx)? + }; + + let bitpacked = BitPacked::try_new( + BufferHandle::new_host(packed), + array.ptype(), + array.validity()?, + patches, + widths, + array.len(), + 0, + )?; + bitpacked.statistics().inherit_from(array.statistics()); + Ok(bitpacked) +} + +#[expect(unused_comparisons, clippy::absurd_extreme_comparisons)] +fn ensure_non_negative(array: &PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult<()> { + if array.ptype().is_signed_int() { + let has_negative_values = match_each_integer_ptype!(array.ptype(), |P| { + array.statistics().compute_min::

(ctx).unwrap_or_default() < 0 + }); + if has_negative_values { + vortex_bail!(InvalidArgument: "cannot bitpack_encode array containing negative integers") + } + } + Ok(()) +} + +/// Bitpack a [PrimitiveArray] with one width per 1024-element chunk. +/// +/// # Safety +/// +/// Internally this function will promote the provided array to its unsigned equivalent. This will +/// violate ordering guarantees if the array contains any negative values, so the caller must +/// ensure that `parray` is non-negative. +pub unsafe fn bitpack_unchecked_with_widths( + parray: &PrimitiveArray, + widths: &ChunkWidths, +) -> ByteBuffer { + let parray = parray.reinterpret_cast(parray.ptype().to_unsigned()); + match_each_unsigned_integer_ptype!(parray.ptype(), |P| { + bitpack_primitive_chunked(parray.as_slice::

(), widths).into_byte_buffer() + }) +} + +/// Bitpack a slice of primitives, packing each 1024-element chunk at its own width. +/// +/// Chunks of width zero contribute no packed bytes; the trailing partial chunk is zero-padded. +pub fn bitpack_primitive_chunked( + array: &[T], + widths: &ChunkWidths, +) -> Buffer { + let mut output = BufferMut::::with_capacity(widths.packed_bytes() / size_of::()); + let mut last_chunk = [T::zero(); FL_CHUNK_SIZE]; + + for (chunk_idx, chunk) in array.chunks(FL_CHUNK_SIZE).enumerate() { + let bit_width = widths.width(chunk_idx); + if bit_width == 0 { + continue; + } + let packed_len = chunk_packed_bytes(bit_width) / size_of::(); + let input: &[T] = if chunk.len() == FL_CHUNK_SIZE { + chunk + } else { + last_chunk[..chunk.len()].copy_from_slice(chunk); + &last_chunk + }; let output_len = output.len(); + // SAFETY: `input` holds exactly 1024 values and the output window is exactly one packed + // block at `bit_width`, which the capacity reserved above accounts for. unsafe { output.set_len(output_len + packed_len); BitPacking::unchecked_pack( - bit_width, - &last_chunk, + bit_width as usize, + input, &mut output[output_len..][..packed_len], ); - }; + } } output.freeze() } -pub fn gather_patches( +/// Gather the values that do not fit their chunk's bit width into patches. +pub fn gather_patches_with_widths( parray: &PrimitiveArray, - bit_width: u8, + widths: &ChunkWidths, num_exceptions_hint: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { @@ -215,7 +588,7 @@ pub fn gather_patches( match_each_integer_ptype!(parray.ptype(), |T| { gather_patches_impl::( parray.as_slice::(), - bit_width, + widths, num_exceptions_hint, patch_validity, validity_mask, @@ -225,7 +598,7 @@ pub fn gather_patches( match_each_integer_ptype!(parray.ptype(), |T| { gather_patches_impl::( parray.as_slice::(), - bit_width, + widths, num_exceptions_hint, patch_validity, validity_mask, @@ -235,7 +608,7 @@ pub fn gather_patches( match_each_integer_ptype!(parray.ptype(), |T| { gather_patches_impl::( parray.as_slice::(), - bit_width, + widths, num_exceptions_hint, patch_validity, validity_mask, @@ -245,7 +618,7 @@ pub fn gather_patches( match_each_integer_ptype!(parray.ptype(), |T| { gather_patches_impl::( parray.as_slice::(), - bit_width, + widths, num_exceptions_hint, patch_validity, validity_mask, @@ -258,7 +631,7 @@ pub fn gather_patches( fn gather_patches_impl( data: &[T], - bit_width: u8, + widths: &ChunkWidths, num_exceptions_hint: usize, patch_validity: Validity, validity_mask: Mask, @@ -270,16 +643,20 @@ where let mut indices: BufferMut

= BufferMut::with_capacity(num_exceptions_hint); let mut values: BufferMut = BufferMut::with_capacity(num_exceptions_hint); - let total_chunks = data.len().div_ceil(1024); + let total_chunks = data.len().div_ceil(FL_CHUNK_SIZE); let mut chunk_offsets: BufferMut = BufferMut::with_capacity(total_chunks); + // A value overflows its chunk's width when it has fewer leading zeros than this. + let mut overflow_leading_zeros = 0usize; for ((idx, value), valid) in data.iter().enumerate().zip(validity_mask.iter()) { - if (idx % 1024) == 0 { + if idx.is_multiple_of(FL_CHUNK_SIZE) { // Record the patch index offset for each chunk. chunk_offsets.push(values.len() as u64); + overflow_leading_zeros = + T::PTYPE.bit_width() - widths.width(idx / FL_CHUNK_SIZE) as usize; } - if (value.leading_zeros() as usize) < T::PTYPE.bit_width() - bit_width as usize && valid { + if (value.leading_zeros() as usize) < overflow_leading_zeros && valid { indices.push(P::from(idx).vortex_expect("cast index from usize")); values.push(*value); } @@ -298,6 +675,26 @@ where } } +/// The width minimising one chunk's cost: its packed block plus the exceptions left behind. +/// +/// A chunk always occupies a whole `128 * width` byte block, so a partial trailing chunk is +/// charged for its padding. +fn best_chunk_width(bit_width_freq: &[usize], bytes_per_exception: usize) -> u8 { + let len: usize = bit_width_freq.iter().sum(); + let mut num_packed = 0; + let mut best_cost = usize::MAX; + let mut best_width = 0; + for (bit_width, freq) in bit_width_freq.iter().enumerate() { + num_packed += *freq; + let cost = chunk_packed_bytes(bit_width as u8) + (len - num_packed) * bytes_per_exception; + if cost < best_cost { + best_cost = cost; + best_width = bit_width; + } + } + best_width as u8 +} + pub fn bit_width_histogram( array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, @@ -380,6 +777,7 @@ fn best_bit_width(bit_width_freq: &[usize], bytes_per_exception: usize) -> Vorte Ok(best_width as u8) } +/// Exceptions cost their value plus a u32 index; we cannot predict how patches compress. fn bytes_per_exception(ptype: PType) -> usize { ptype.byte_width() + 4 } @@ -427,11 +825,9 @@ pub mod test_harness { } #[cfg(test)] -mod test { +mod tests { use std::sync::LazyLock; - use rand::SeedableRng; - use rand::rngs::StdRng; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::assert_arrays_eq; @@ -443,9 +839,8 @@ mod test { use vortex_session::VortexSession; use super::*; + use crate::BitPackedArrayExt; use crate::BitPackedData; - use crate::bitpack_compress::test_harness::make_array; - use crate::bitpacking::array::BitPackedArrayExt; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -454,14 +849,16 @@ mod test { }); #[test] - fn test_best_bit_width() { - // 10 1-bit values, 20 2-bit, etc. - let freq = vec![0, 10, 20, 15, 1, 0, 0, 0]; - // 3-bits => (46 * 3) + (8 * 1 * 5) => 178 bits => 23 bytes and zero exceptions - assert_eq!( - best_bit_width(&freq, bytes_per_exception(PType::U8)).unwrap(), - 3 - ); + fn test_best_chunk_width() { + // 1000 3-bit values and 24 10-bit values in a u16 chunk: 3 bits plus 24 exceptions + // (384 + 24 * 6 bytes) beats 10 bits for everything (1280 bytes). + let mut freq = vec![0usize; 17]; + freq[3] = 1000; + freq[10] = 24; + assert_eq!(best_chunk_width(&freq, bytes_per_exception(PType::U16)), 3); + // Make the exceptions expensive enough and the wide width wins. + freq[10] = 200; + assert_eq!(best_chunk_width(&freq, bytes_per_exception(PType::U16)), 10); } #[test] @@ -501,14 +898,29 @@ mod test { assert!(matches!(err, VortexError::InvalidArgument(_, _))); } + /// Values below 100 with every 40th value pushed above 12 bits, and every 5th null. + fn patchy_nullable(len: usize, seed: u32) -> PrimitiveArray { + let values = (0..len as u32) + .map(|i| { + let v = (i * 7919 + seed) % 100; + if i % 40 == 0 { v + (1 << 13) } else { v } + }) + .map(|v| v as i32) + .collect::>(); + let validity = Validity::from_iter((0..len).map(|i| i % 5 != 0)); + PrimitiveArray::new(values, validity) + } + #[test] fn canonicalize_chunked_of_bitpacked() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let mut rng = StdRng::seed_from_u64(0); let chunks = (0..10) - .map(|_| make_array(&mut rng, 100, 0.25, 0.25, &mut ctx).unwrap()) - .collect::>(); + .map(|seed| { + bitpack_encode(&patchy_nullable(100, seed), 12, None, &mut ctx) + .map(|a| a.into_array()) + }) + .collect::>>()?; let chunked = ChunkedArray::from_iter(chunks).into_array(); let into_ca = chunked.clone().execute::(&mut ctx)?; @@ -518,148 +930,50 @@ mod test { let ca_into = primitive_builder.finish(); assert_arrays_eq!(into_ca, ca_into, &mut ctx); - - let mut primitive_builder = - PrimitiveBuilder::::with_capacity(chunked.dtype().nullability(), 10 * 100); - chunked.append_to_builder(&mut primitive_builder, &mut ctx)?; - let ca_into = primitive_builder.finish(); - - assert_arrays_eq!(into_ca, ca_into, &mut ctx); - Ok(()) } - #[test] - fn test_chunk_offsets() -> VortexResult<()> { + fn chunk_offsets_of(values: Vec) -> VortexResult { let mut ctx = SESSION.create_execution_ctx(); - let patch_value = 1u32 << 20; - let patch_indices = [100usize, 200, 3000, 3100]; - let mut values = vec![0u32; 4096usize]; - - patch_indices - .iter() - .for_each(|&idx| values[idx] = patch_value); - let array = PrimitiveArray::from_iter(values); let bitpacked = bitpack_encode(&array, 4, None, &mut ctx)?; - let patches = bitpacked .patches() .ok_or_else(|| vortex_err!("expected patches"))?; - let chunk_offsets = patches + patches .chunk_offsets() .as_ref() .ok_or_else(|| vortex_err!("expected chunk offsets"))? .clone() - .execute::(&mut ctx)?; + .execute::(&mut ctx) + } - // chunk 0 (0-1023): patches at 100, 200 -> starts at patch index 0 - // chunk 1 (1024-2047): no patches -> points to patch index 2 - // chunk 2 (2048-3071): patch at 3000 -> starts at patch index 2 - // chunk 3 (3072-4095): patch at 3100 -> starts at patch index 3 - assert_arrays_eq!( - chunk_offsets, - PrimitiveArray::from_iter([0u64, 2, 2, 3]), - &mut ctx - ); - Ok(()) + fn with_patches(len: usize, patch_indices: &[usize]) -> Vec { + let mut values = vec![0u32; len]; + patch_indices.iter().for_each(|&idx| values[idx] = 1 << 20); + values } #[test] - fn test_chunk_offsets_no_patches_in_middle() -> VortexResult<()> { + fn test_chunk_offsets() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let patch_value = 1u32 << 20; - let patch_indices = [100usize, 200, 2500]; - let mut values = vec![0u32; 3072usize]; - - patch_indices - .iter() - .for_each(|&idx| values[idx] = patch_value); - - let array = PrimitiveArray::from_iter(values); - let bitpacked = bitpack_encode(&array, 4, None, &mut ctx)?; - - let patches = bitpacked - .patches() - .ok_or_else(|| vortex_err!("expected patches"))?; - let chunk_offsets = patches - .chunk_offsets() - .as_ref() - .ok_or_else(|| vortex_err!("expected chunk offsets"))? - .clone() - .execute::(&mut ctx)?; - + // chunk 0: patches at 100, 200; chunk 1: none; chunk 2: 3000; chunk 3: 3100 assert_arrays_eq!( - chunk_offsets, - PrimitiveArray::from_iter([0u64, 2, 2]), + chunk_offsets_of(with_patches(4096, &[100, 200, 3000, 3100]))?, + PrimitiveArray::from_iter([0u64, 2, 2, 3]), &mut ctx ); - Ok(()) - } - - #[test] - fn test_chunk_offsets_trailing_empty_chunks() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let patch_value = 1u32 << 20; - let patch_indices = [100usize, 200, 1500]; - let mut values = vec![0u32; 5120usize]; - - patch_indices - .iter() - .for_each(|&idx| values[idx] = patch_value); - - let array = PrimitiveArray::from_iter(values); - let bitpacked = bitpack_encode(&array, 4, None, &mut ctx)?; - - let patches = bitpacked - .patches() - .ok_or_else(|| vortex_err!("expected patches"))?; - let chunk_offsets = patches - .chunk_offsets() - .as_ref() - .ok_or_else(|| vortex_err!("expected chunk offsets"))? - .clone() - .execute::(&mut ctx)?; - - // chunk 0 (0-1023): patches at 100, 200 -> starts at patch index 0 - // chunk 1 (1024-2047): patch at 1500 -> starts at patch index 2 - // chunk 2 (2048-3071): no patches -> points to patch index 3 - // chunk 3 (3072-4095): no patches -> points to patch index 3 (remaining chunks filled) - // chunk 4 (4096-5119): no patches -> points to patch index 3 (remaining chunks filled) + // Trailing chunks without patches all point past the last patch. assert_arrays_eq!( - chunk_offsets, + chunk_offsets_of(with_patches(5120, &[100, 200, 1500]))?, PrimitiveArray::from_iter([0u64, 2, 3, 3, 3]), &mut ctx ); - Ok(()) - } - - #[test] - fn test_chunk_offsets_single_chunk() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let patch_value = 1u32 << 20; - let patch_indices = [100usize, 200]; - let mut values = vec![0u32; 500usize]; - - patch_indices - .iter() - .for_each(|&idx| values[idx] = patch_value); - - let array = PrimitiveArray::from_iter(values); - let bitpacked = bitpack_encode(&array, 4, None, &mut ctx)?; - - let patches = bitpacked - .patches() - .ok_or_else(|| vortex_err!("expected patches"))?; - let chunk_offsets = patches - .chunk_offsets() - .as_ref() - .ok_or_else(|| vortex_err!("expected chunk offsets"))? - .clone() - .execute::(&mut ctx)?; - - // Single chunk starting at patch index 0. - assert_arrays_eq!(chunk_offsets, PrimitiveArray::from_iter([0u64]), &mut ctx); + assert_arrays_eq!( + chunk_offsets_of(with_patches(500, &[100, 200]))?, + PrimitiveArray::from_iter([0u64]), + &mut ctx + ); Ok(()) } } diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs index d192e50d04f..b73a8d9f0cd 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs @@ -23,8 +23,8 @@ 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( @@ -161,15 +161,16 @@ pub(crate) fn apply_patches_to_uninit_range, index: usize) -> Scalar { - let bit_width = array.bit_width() as usize; 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 scalar: Scalar = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |P| { - unsafe { - unpack_single_primitive::

(array.packed_slice::

(), bit_width, index_in_encoded) - .into() - } + let (packed_chunk, bit_width) = array.packed_chunk::

(chunk); + // SAFETY: `packed_chunk` is exactly one packed block at `bit_width`, and the index is + // within the chunk. + unsafe { BitPacking::unchecked_unpack_single(bit_width, packed_chunk, index_in_chunk) } + .into() }); // Cast to fix signedness and nullability scalar.cast(array.dtype()).vortex_expect("cast failure") @@ -228,7 +229,7 @@ mod tests { 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() diff --git a/encodings/fastlanes/src/bitpacking/array/mod.rs b/encodings/fastlanes/src/bitpacking/array/mod.rs index 05cbee8b3ef..356a92a29b3 100644 --- a/encodings/fastlanes/src/bitpacking/array/mod.rs +++ b/encodings/fastlanes/src/bitpacking/array/mod.rs @@ -3,11 +3,15 @@ use std::fmt::Display; use std::fmt::Formatter; +use std::hash::Hash; +use std::hash::Hasher; use std::mem::MaybeUninit; +use std::ops::Range; use fastlanes::BitPacking; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::arrays::Primitive; @@ -15,12 +19,15 @@ 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; use vortex_array::patches::PatchesData; use vortex_array::validity::Validity; use vortex_array::vtable::child_to_validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -31,9 +38,147 @@ 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] +pub const fn chunk_packed_bytes(bit_width: u8) -> usize { + (FL_CHUNK_SIZE / 8) * bit_width as usize +} + +/// The bit width of every FastLanes chunk of a bit-packed array. +/// +/// Each 1024-element chunk is packed independently at its own width, so an array always carries +/// one width per chunk. The byte offset of every chunk's packed block is cached so chunk lookup +/// stays O(1). +#[derive(Clone, Debug)] +pub struct ChunkWidths { + widths: Buffer, + /// Byte offset of every chunk's packed block, with a trailing entry holding the total. + byte_offsets: Buffer, + max_width: u8, +} + +impl ChunkWidths { + /// Build from one width per chunk. + pub fn new(widths: Buffer) -> Self { + let mut byte_offsets = BufferMut::::with_capacity(widths.len() + 1); + let mut total = 0u64; + let mut max_width = 0u8; + byte_offsets.push(0); + for &w in widths.iter() { + total += chunk_packed_bytes(w) as u64; + max_width = max_width.max(w); + byte_offsets.push(total); + } + Self { + widths, + byte_offsets: byte_offsets.freeze(), + max_width, + } + } + + /// `num_chunks` chunks all packed at `bit_width`. + pub fn uniform(bit_width: u8, num_chunks: usize) -> Self { + Self::new(Buffer::from_iter(std::iter::repeat_n( + bit_width, num_chunks, + ))) + } + + /// Number of chunks. + #[inline] + pub fn len(&self) -> usize { + self.widths.len() + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.widths.is_empty() + } + + /// The bit width of `chunk`. + #[inline] + pub fn width(&self, chunk: usize) -> u8 { + self.widths[chunk] + } + + /// One width per chunk. + #[inline] + pub fn as_slice(&self) -> &[u8] { + self.widths.as_slice() + } + + /// The widest chunk width. + #[inline] + pub fn max_width(&self) -> u8 { + self.max_width + } + + /// The single width shared by every chunk, if they all agree. + pub fn uniform_width(&self) -> Option { + let first = *self.widths.first()?; + self.widths.iter().all(|&w| w == first).then_some(first) + } + + /// Whether every chunk shares one width. An array with no chunks counts as uniform. + pub fn is_uniform(&self) -> bool { + self.widths + .first() + .is_none_or(|&first| self.widths.iter().all(|&w| w == first)) + } + + /// The widths as a buffer of one byte per chunk. + pub fn as_buffer(&self) -> &Buffer { + &self.widths + } + + /// Byte offset of `chunk`'s packed block. Passing the chunk count yields the total size. + #[inline] + pub fn byte_offset(&self, chunk: usize) -> usize { + self.byte_offsets[chunk] as usize + } + + /// Total packed bytes. + #[inline] + pub fn packed_bytes(&self) -> usize { + self.byte_offset(self.len()) + } + + /// Restrict to the given range of chunks. + pub fn slice(&self, chunks: Range) -> Self { + Self::new(self.widths.slice(chunks)) + } +} + +impl PartialEq for ChunkWidths { + fn eq(&self, other: &Self) -> bool { + self.widths.as_slice() == other.widths.as_slice() + } +} + +impl Eq for ChunkWidths {} + +impl Hash for ChunkWidths { + fn hash(&self, state: &mut H) { + self.widths.as_slice().hash(state); + } +} + +impl Display for ChunkWidths { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self.uniform_width() { + Some(w) => write!(f, "bit_width: {w}"), + None => write!( + f, + "bit_widths: {} chunks, max {}", + self.len(), + self.max_width + ), + } + } +} #[array_slots(crate::BitPacked)] pub struct BitPackedSlots { @@ -49,6 +194,10 @@ pub struct BitPackedSlots { /// The validity bitmap indicating which elements are non-null. #[slot(3)] pub validity_child: Option, + /// One width per 1024-element chunk, present exactly when the chunks' widths differ. A child + /// rather than metadata, so it can be compressed and metadata stays bounded. + #[slot(4)] + pub width_table: Option, } pub(crate) const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { @@ -57,9 +206,52 @@ 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); + +/// The width table child for `widths`. Absent when every chunk shares one width, since that case +/// serializes under the original format, which carries the width in its metadata. +pub(crate) fn width_table_child(widths: &ChunkWidths) -> Option { + (!widths.is_uniform()).then(|| { + PrimitiveArray::new(widths.as_buffer().clone(), Validity::NonNullable).into_array() + }) +} + +/// Check that `table` is the width table child `widths` requires: one `u8` per chunk, present +/// exactly when the chunks' widths differ. +pub(crate) fn validate_width_table( + widths: &ChunkWidths, + table: Option<&ArrayRef>, +) -> VortexResult<()> { + match table { + Some(table) => { + vortex_ensure!( + !widths.is_uniform(), + "BitPacked arrays with one shared width carry no width table" + ); + vortex_ensure!( + table.dtype() == &WIDTH_TABLE_DTYPE, + "BitPacked width table must be {WIDTH_TABLE_DTYPE}, got {}", + table.dtype() + ); + vortex_ensure!( + table.len() == widths.len(), + "BitPacked width table has {} entries for {} chunks", + table.len(), + widths.len() + ); + } + None => vortex_ensure!( + widths.is_uniform(), + "BitPacked arrays with differing chunk widths must carry a width table" + ), + } + Ok(()) +} + pub struct BitPackedDataParts { pub offset: u16, - pub bit_width: u8, + pub widths: ChunkWidths, pub len: usize, pub packed: BufferHandle, pub patches: Option, @@ -71,7 +263,7 @@ 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) widths: ChunkWidths, pub(super) packed: BufferHandle, /// Patch metadata for reconstructing Patches from slots. pub(super) patches_data: Option, @@ -79,67 +271,58 @@ 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.widths, 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 `widths` 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. + /// * `widths` must hold one width per chunk, each no wider than the `ptype` + /// * The `packed` buffer must be exactly the sum of the chunks' packed sizes. /// /// Any violation of these preconditions will result in an error. pub fn try_new( packed: BufferHandle, patches: Option, - bit_width: u8, + widths: ChunkWidths, 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}" + widths.max_width() <= 64, + "Unsupported bit width {}", + widths.max_width() + ); + vortex_ensure!( + (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, + widths, packed, patches_data: patches.as_ref().map(PatchesData::from_patches), }) @@ -150,12 +333,16 @@ impl BitPackedData { ptype: PType, validity: &Validity, patches: Option<&Patches>, - bit_width: u8, + widths: &ChunkWidths, length: usize, offset: u16, ) -> VortexResult<()> { vortex_ensure!(ptype.is_int(), MismatchedTypes: "integer", ptype); - vortex_ensure!(bit_width <= 64, "Unsupported bit width {bit_width}"); + vortex_ensure!( + widths.max_width() as usize <= ptype.bit_width(), + "Unsupported bit width {} for {ptype}", + widths.max_width() + ); if let Some(validity_len) = validity.maybe_len() { vortex_ensure!( @@ -169,9 +356,14 @@ 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); + // Validate chunk widths and the packed buffer + let num_chunks = (length + offset as usize).div_ceil(FL_CHUNK_SIZE); + vortex_ensure!( + widths.len() == num_chunks, + "Expected {num_chunks} chunk widths, got {}", + widths.len() + ); + let expected_packed_len = widths.packed_bytes(); vortex_ensure!( packed.len() == expected_packed_len, "Expected {} packed bytes, got {}", @@ -224,6 +416,18 @@ impl BitPackedData { unsafe { std::slice::from_raw_parts(packed_ptr, packed_len) } } + /// The packed FastLanes block of `chunk` as `T` words, along with that chunk's bit width. + #[inline] + pub fn packed_chunk(&self, chunk: usize) -> (&[T], usize) { + let bit_width = self.widths.width(chunk); + let start = self.widths.byte_offset(chunk) / size_of::(); + let len = chunk_packed_bytes(bit_width) / size_of::(); + ( + &self.packed_slice::()[start..][..len], + bit_width as usize, + ) + } + /// Accessor for bit unpacked chunks pub fn unpacked_chunks<'a, T: BitPackedIter>( &'a self, @@ -239,10 +443,16 @@ impl BitPackedData { BitUnpackedChunks::try_new(self, len, scratch) } - /// Bit-width of the packed values + /// The bit width of every chunk. + #[inline] + pub fn chunk_widths(&self) -> &ChunkWidths { + &self.widths + } + + /// The widest bit width used by any chunk. #[inline] pub fn bit_width(&self) -> u8 { - self.bit_width + self.widths.max_width() } #[inline] @@ -273,12 +483,18 @@ impl BitPackedData { bitpack_encode(&parray, bit_width, None, ctx) } - /// Calculate the maximum value that **can** be contained by this array, given its bit-width. + /// Calculate the maximum value that **can** be contained by this array, given its widest + /// chunk. /// /// 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 + let bit_width = self.bit_width() as u32; + if bit_width >= usize::BITS { + usize::MAX + } else { + (1usize << bit_width) - 1 + } } } @@ -288,6 +504,11 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { BitPackedData::packed(self) } + #[inline] + fn chunk_widths(&self) -> &ChunkWidths { + BitPackedData::chunk_widths(self) + } + #[inline] fn bit_width(&self) -> u8 { BitPackedData::bit_width(self) @@ -318,6 +539,11 @@ pub trait BitPackedArrayExt: BitPackedArraySlotsExt { BitPackedData::packed_slice::(self) } + #[inline] + fn packed_chunk(&self, chunk: usize) -> (&[T], usize) { + BitPackedData::packed_chunk::(self, chunk) + } + #[inline] fn unpacked_chunks<'a, T: BitPackedIter>( &'a self, @@ -343,8 +569,10 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_buffer::Buffer; + use vortex_buffer::buffer; use vortex_session::VortexSession; + use super::ChunkWidths; use crate::BitPackedData; use crate::bitpacking::array::BitPackedArrayExt; @@ -407,4 +635,22 @@ mod test { &mut ctx ); } + + #[test] + fn chunk_widths_offsets() { + assert_eq!(ChunkWidths::uniform(3, 3).uniform_width(), Some(3)); + assert_eq!(ChunkWidths::new(Buffer::::empty()).packed_bytes(), 0); + + let widths = ChunkWidths::new(buffer![3u8, 0, 16]); + assert_eq!(widths.uniform_width(), None); + assert_eq!(widths.len(), 3); + assert_eq!(widths.max_width(), 16); + assert_eq!(widths.width(1), 0); + assert_eq!(widths.byte_offset(0), 0); + assert_eq!(widths.byte_offset(1), 128 * 3); + assert_eq!(widths.byte_offset(2), 128 * 3); + assert_eq!(widths.packed_bytes(), 128 * 19); + assert_eq!(widths.slice(1..3), ChunkWidths::new(buffer![0u8, 16])); + assert_eq!(widths.slice(0..1).uniform_width(), Some(3)); + } } diff --git a/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking/array/unpack_iter.rs index 518495896e2..5b0b051f7fb 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::ChunkWidths; +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: &ChunkWidths, 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,21 @@ 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 mut unpacked_chunks: BitUnpackedChunks = array.unpacked_chunks(&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 +98,7 @@ impl> UnpackStrategy for BitPackingStr /// ``` pub struct UnpackedChunks<'a, T: PhysicalPType, S: UnpackStrategy> { strategy: S, - bit_width: usize, + widths: &'a ChunkWidths, offset: usize, len: usize, num_chunks: usize, @@ -109,7 +119,7 @@ impl<'a, T: BitPacked> BitUnpackedChunks<'a, T> { Self::try_new_with_strategy( BitPackingStrategy, array.packed_slice::(), - array.bit_width() as usize, + array.chunk_widths(), array.offset() as usize, len, scratch, @@ -117,14 +127,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 +143,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 ChunkWidths, 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 +164,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 +181,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]) } }) @@ -236,13 +244,17 @@ 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 ChunkWidths = self.widths; + let mut start = widths.byte_offset(range.start) / size_of::(); + for &bit_width in &widths.as_slice()[range] { + 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); } @@ -271,16 +283,18 @@ 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 &bit_width in &self.widths.as_slice()[range] { + 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; } @@ -295,15 +309,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]) } }) @@ -318,33 +331,48 @@ 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: &ChunkWidths, 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, &bit_width) in widths.as_slice().iter().enumerate() { + 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: &ChunkWidths, offset: usize, len: usize, ) -> VortexResult<(usize, usize)> { @@ -352,12 +380,16 @@ 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() + ); + 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)) } @@ -365,29 +397,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 ChunkWidths, 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 ChunkWidths, 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::(), } } } @@ -404,15 +437,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 new file mode 100644 index 00000000000..94133708b2e --- /dev/null +++ b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs @@ -0,0 +1,627 @@ +// 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 prost::Message; +use rstest::rstest; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::fns::is_constant::is_constant; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::buffer::BufferHandle; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; +use vortex_array::compute::conformance::cast::test_cast_conformance; +use vortex_array::compute::conformance::consistency::test_array_consistency; +use vortex_array::compute::conformance::filter::test_filter_conformance; +use vortex_array::compute::conformance::take::test_take_conformance; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::between::BetweenOptions; +use vortex_array::scalar_fn::fns::between::StrictComparison; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::serde::ArrayChildren; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::BitPacked; +use crate::BitPackedArray; +use crate::BitPackedArrayExt; +use crate::BitPackedArraySlotsExt; +use crate::BitPackedPlugin; +use crate::ChunkWidths; +use crate::FL_CHUNK_SIZE; +use crate::FoR; +use crate::bitpacked_v2_id; +use crate::bitpacking::bitpack_compress::bitpack_encode_with_widths; +use crate::bitpacking::bitpack_compress::bitpack_to_best_bit_width; +use crate::bitpacking::bitpack_compress::bitpack_to_best_chunk_widths; +use crate::bitpacking::bitpack_compress::bitpack_to_best_chunk_widths_multipass; +use crate::bitpacking::plugin::BitPackedV2Metadata; +use crate::bitpacking::vtable::BitPackedMetadata; + +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_chunk_widths(&PrimitiveArray::from_iter(values.iter().copied()), &mut ctx) +} + +fn primitive(values: &[u32]) -> ArrayRef { + PrimitiveArray::from_iter(values.iter().copied()).into_array() +} + +#[test] +fn picks_a_width_per_chunk() -> VortexResult<()> { + let packed = encode(&varied(100))?; + let widths = packed.chunk_widths(); + assert_eq!( + widths.uniform_width(), + None, + "chunks differ in magnitude: {widths}" + ); + assert_eq!(widths.len(), 5); + assert_eq!(widths.width(2), 0, "an all-zero chunk stores nothing"); + assert!(widths.width(0) < widths.width(1)); + assert!(widths.width(1) < widths.width(3)); + assert_eq!(widths.max_width(), widths.width(3)); + assert_eq!(packed.bit_width(), widths.max_width()); + assert!( + packed.patches().is_some(), + "chunk 1 outliers become patches" + ); + Ok(()) +} + +#[test] +fn uniform_data_gets_equal_widths() -> VortexResult<()> { + let values: Vec = (0..3000).map(|i| i % 128).collect(); + let packed = encode(&values)?; + assert_eq!(packed.chunk_widths().len(), 3); + assert_eq!(packed.chunk_widths().uniform_width(), Some(7)); + assert_eq!(packed.bit_width(), 7); + Ok(()) +} + +#[rstest] +#[case::exact_chunks(0)] +#[case::partial_tail(100)] +#[case::single_tail(1)] +fn roundtrip(#[case] tail: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(tail); + let packed = encode(&values)?; + assert_arrays_eq!(packed, primitive(&values), &mut ctx); + Ok(()) +} + +#[test] +fn scalar_at_every_chunk() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + for idx in [ + 0, + 5, + 1024, + 1024 + 7, + 1024 + 307, + 2048, + 2500, + 3072, + 4000, + 4095, + 4096, + 4195, + ] { + assert_eq!( + packed.execute_scalar(idx, &mut ctx)?, + Scalar::from(values[idx]), + "index {idx}" + ); + } + Ok(()) +} + +#[rstest] +#[case::within_first_chunk(10..900)] +#[case::across_first_boundary(900..1100)] +#[case::whole_middle_chunks(1024..3072)] +#[case::through_zero_chunk(1500..2600)] +#[case::into_tail(3000..4150)] +#[case::tail_only(4100..4196)] +fn slice_matches_primitive(#[case] range: std::ops::Range) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + let sliced = packed.slice(range.clone())?; + let expected = primitive(&values).slice(range.clone())?; + assert_arrays_eq!(sliced, expected, &mut ctx); + + // The slice is still bit-packed and keeps only the widths of the chunks it overlaps. + let sliced = sliced.execute::(&mut ctx)?; + if let Some(bp) = sliced.as_opt::() { + let expected_chunks = (range.end).div_ceil(FL_CHUNK_SIZE) - range.start / FL_CHUNK_SIZE; + assert_eq!(bp.chunk_widths().len(), expected_chunks); + } + assert_eq!( + sliced.execute_scalar(0, &mut ctx)?, + Scalar::from(values[range.start]) + ); + Ok(()) +} + +#[test] +fn take_sparse_indices() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + // Few enough indices that the kernel unpacks single values rather than the whole array. + let indices = [3usize, 1030, 1031, 1331, 2100, 3500, 4100]; + let taken = packed + .take(buffer![3u64, 1030, 1031, 1331, 2100, 3500, 4100].into_array())? + .execute::(&mut ctx)?; + assert_arrays_eq!( + taken, + PrimitiveArray::from_iter(indices.iter().map(|&i| values[i])), + &mut ctx + ); + Ok(()) +} + +#[test] +fn filter_sparse_mask() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + let indices = vec![3usize, 1030, 1031, 1331, 2100, 3500, 4100]; + let filtered = packed + .filter(Mask::from_indices(values.len(), indices.clone()))? + .execute::(&mut ctx)?; + assert_arrays_eq!( + filtered, + PrimitiveArray::from_iter(indices.iter().map(|&i| values[i])), + &mut ctx + ); + Ok(()) +} + +#[rstest] +#[case(Operator::Eq)] +#[case(Operator::Lt)] +#[case(Operator::Gte)] +fn compare_constant_matches_primitive(#[case] op: Operator) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + // Zero lands inside the all-zero chunk's range, exercising the zero-width fused path. + for rhs in [0u32, 5, 3000] { + let rhs = ConstantArray::new(rhs, values.len()).into_array(); + let got = packed + .clone() + .binary(rhs.clone(), op)? + .execute::(&mut ctx)?; + let want = primitive(&values) + .binary(rhs, op)? + .execute::(&mut ctx)?; + assert_arrays_eq!(got, want, &mut ctx); + } + Ok(()) +} + +#[test] +fn between_matches_primitive() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + let lower = ConstantArray::new(2u32, values.len()).into_array(); + let upper = ConstantArray::new(3000u32, values.len()).into_array(); + let options = BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::Strict, + }; + let got = packed + .between(lower.clone(), upper.clone(), options.clone())? + .execute::(&mut ctx)?; + let want = primitive(&values) + .between(lower, upper, options)? + .execute::(&mut ctx)?; + assert_arrays_eq!(got, want, &mut ctx); + Ok(()) +} + +#[test] +fn widening_cast_matches_primitive() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + let target = DType::Primitive(PType::U64, Nullability::NonNullable); + let got = packed + .cast(target.clone())? + .execute::(&mut ctx)?; + let want = primitive(&values) + .cast(target)? + .execute::(&mut ctx)?; + assert_arrays_eq!(got, want, &mut ctx); + Ok(()) +} + +#[test] +fn not_constant() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let packed = encode(&varied(100))?.into_array(); + assert!(!is_constant(&packed, &mut ctx)?); + Ok(()) +} + +#[test] +fn nullable_and_signed_roundtrip() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = varied(50).into_iter().map(|v| v as i32).collect(); + let validity = Validity::from_iter((0..values.len()).map(|i| i % 7 != 0)); + let array = PrimitiveArray::new(Buffer::from_iter(values.iter().copied()), validity.clone()); + let packed = bitpack_to_best_chunk_widths(&array, &mut ctx)?; + assert_eq!(packed.chunk_widths().uniform_width(), None); + assert_eq!( + packed.dtype(), + &DType::Primitive(PType::I32, Nullability::Nullable) + ); + assert_arrays_eq!( + packed, + PrimitiveArray::new(Buffer::from_iter(values.iter().copied()), validity), + &mut ctx + ); + Ok(()) +} + +#[test] +fn explicit_widths_including_full_width_chunk() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..2 * FL_CHUNK_SIZE + 10) + .map(|i| { + if i < FL_CHUNK_SIZE { + (i % 16) as u16 + } else { + u16::MAX - i as u16 + } + }) + .collect(); + let array = PrimitiveArray::from_iter(values.iter().copied()); + // Chunk 1 and the tail use the full 16 bits, which a single global width could never pick. + let widths = ChunkWidths::new(buffer![4u8, 16, 16]); + let packed = bitpack_encode_with_widths(&array, widths, &mut ctx)?; + assert!(packed.patches().is_none()); + assert_eq!(packed.bit_width(), 16); + assert_arrays_eq!(packed, array, &mut ctx); + Ok(()) +} + +#[test] +fn for_fused_decode() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + let for_array = FoR::try_new(packed, Scalar::from(1000u32))?; + assert_arrays_eq!( + for_array, + PrimitiveArray::from_iter(values.iter().map(|v| v + 1000)), + &mut ctx + ); + Ok(()) +} + +/// Serialize `array` through the session and read it back through the plugin registered for the +/// serialized ID, as a file reader would. +fn serde_roundtrip(array: &BitPackedArray) -> VortexResult<(ArrayId, Vec, ArrayRef)> { + let array_ref = array.as_array(); + let serialization = SESSION + .array_serialize(array_ref)? + .ok_or_else(|| vortex_err!("BitPacked must serialize"))?; + let children = array_ref.children(); + let buffers = array_ref + .buffers() + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + let plugin = SESSION + .arrays() + .registry() + .get(&serialization.serialized_id) + .ok_or_else(|| vortex_err!("no plugin for {}", serialization.serialized_id))?; + let parts = ArrayDeserialization::new( + serialization.serialized_id, + array_ref.dtype(), + array_ref.len(), + &serialization.metadata, + &buffers, + &children, + ); + let read = plugin.deserialize(parts, &SESSION)?; + Ok((serialization.serialized_id, serialization.metadata, read)) +} + +/// Differing chunk widths serialize under the v2 ID with the width table as a child, and read +/// back with the same widths. +#[test] +fn differing_widths_serialize_as_v2() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?; + let table = packed + .width_table() + .ok_or_else(|| vortex_err!("differing widths must carry a width table"))?; + assert_eq!( + table + .clone() + .execute::(&mut ctx)? + .as_slice::(), + packed.chunk_widths().as_slice() + ); + let (id, metadata, read) = serde_roundtrip(&packed)?; + assert_eq!(id, bitpacked_v2_id()); + assert_eq!(BitPackedV2Metadata::decode(metadata.as_slice())?.offset, 0); + assert_eq!( + read.as_::().chunk_widths(), + packed.chunk_widths() + ); + assert!(read.as_::().width_table().is_some()); + assert_arrays_eq!(read, primitive(&values), &mut ctx); + Ok(()) +} + +/// One shared width serializes under the original ID with the original metadata and no width +/// table, byte for byte. +#[test] +fn uniform_widths_serialize_as_original_format() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..3000).map(|i| i % 128).collect(); + let packed = encode(&values)?; + assert!(packed.width_table().is_none()); + assert!(packed.as_array().children().is_empty()); + let (id, metadata, read) = serde_roundtrip(&packed)?; + assert_eq!(id, ArrayVTable::id(&BitPacked)); + let original = BitPackedMetadata { + bit_width: 7, + offset: 0, + patches: None, + } + .encode_to_vec(); + assert_eq!(metadata, original); + assert_arrays_eq!(read, primitive(&values), &mut ctx); + Ok(()) +} + +/// An array with no chunks has nothing to tabulate and stays in the original format. +#[test] +fn empty_array_serializes_as_original_format() -> VortexResult<()> { + let packed = encode(&[])?; + assert!(packed.width_table().is_none()); + let (id, _, read) = serde_roundtrip(&packed)?; + assert_eq!(id, ArrayVTable::id(&BitPacked)); + assert!(read.is_empty()); + Ok(()) +} + +/// A compressor may re-encode the width table. The re-encoded child survives a round trip and +/// still yields the same widths. +#[test] +fn compressed_width_table_round_trips() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?; + let table = PrimitiveArray::new( + packed.chunk_widths().as_buffer().clone(), + Validity::NonNullable, + ); + let compressed_table = bitpack_to_best_bit_width(&table, &mut ctx)?.into_array(); + assert!(compressed_table.is::()); + let packed = BitPacked::with_width_table(packed, compressed_table)?; + let (id, _, read) = serde_roundtrip(&packed)?; + assert_eq!(id, bitpacked_v2_id()); + let view = read.as_::(); + assert_eq!(view.chunk_widths(), packed.chunk_widths()); + assert!( + view.width_table() + .is_some_and(|table| table.is::()) + ); + assert_arrays_eq!(read, primitive(&values), &mut ctx); + Ok(()) +} + +/// The width table must hold one `u8` per chunk, and only arrays whose chunk widths differ carry +/// one. +#[test] +fn width_table_is_validated() -> VortexResult<()> { + let packed = encode(&varied(100))?; + let num_chunks = packed.chunk_widths().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_none()); + let table = PrimitiveArray::from_iter(vec![7u8; uniform.chunk_widths().len()]).into_array(); + assert!(BitPacked::with_width_table(uniform, table).is_err()); + Ok(()) +} + +/// Children that report a dtype or length mismatch as an error, as a file reader does, instead +/// of panicking like the slice implementation. +struct StrictChildren(Vec); + +impl ArrayChildren for StrictChildren { + fn get(&self, index: usize, dtype: &DType, len: usize) -> VortexResult { + let child = + <[ArrayRef]>::get(&self.0, index).ok_or_else(|| vortex_err!("no child {index}"))?; + vortex_ensure!( + child.dtype() == dtype, + "child {index} has dtype {}, expected {dtype}", + child.dtype() + ); + vortex_ensure!( + child.len() == len, + "child {index} has length {}, expected {len}", + child.len() + ); + Ok(child.clone()) + } + + fn len(&self) -> usize { + self.0.len() + } +} + +/// The encoding's own serializer is labelled with the original ID, so it only ever writes arrays +/// that satisfy that ID's contract. +#[test] +fn bare_vtable_only_writes_the_original_format() -> VortexResult<()> { + let uniform = encode(&(0..3000u32).map(|i| i % 128).collect::>())?; + assert!(ArrayVTable::serialize(uniform.as_view(), &SESSION)?.is_some()); + let differing = encode(&varied(100))?; + assert!(ArrayVTable::serialize(differing.as_view(), &SESSION)?.is_none()); + Ok(()) +} + +/// Each ID keeps its contract: the original ID cannot read children that carry a width table, +/// and the v2 ID demands one. +#[test] +fn each_format_keeps_its_contract() -> VortexResult<()> { + let read_as = |array: &BitPackedArray, id: ArrayId| -> VortexResult<()> { + let array_ref = array.as_array(); + let serialization = SESSION + .array_serialize(array_ref)? + .ok_or_else(|| vortex_err!("BitPacked must serialize"))?; + let children = StrictChildren(array_ref.children()); + let buffers = array_ref + .buffers() + .into_iter() + .map(BufferHandle::new_host) + .collect::>(); + ArrayPlugin::deserialize( + &BitPackedPlugin, + ArrayDeserialization::new( + id, + array_ref.dtype(), + array_ref.len(), + &serialization.metadata, + &buffers, + &children, + ), + &SESSION, + ) + .map(|_| ()) + }; + + let differing = encode(&varied(100))?; + assert!( + read_as(&differing, ArrayVTable::id(&BitPacked)).is_err(), + "the original ID must reject a width table" + ); + let uniform = encode(&(0..3000u32).map(|i| i % 128).collect::>())?; + assert!( + read_as(&uniform, bitpacked_v2_id()).is_err(), + "the v2 ID must demand a width table" + ); + Ok(()) +} + +#[rstest] +#[case::varied(encode(&varied(100)).unwrap())] +#[case::varied_exact(encode(&varied(0)).unwrap())] +fn conformance(#[case] array: BitPackedArray) { + let mut ctx = SESSION.create_execution_ctx(); + let array = array.into_array(); + test_array_consistency(&array, &mut ctx); + test_take_conformance(&array, &mut ctx); + test_filter_conformance(&array, &mut ctx); + test_cast_conformance(&array, &mut ctx); + test_binary_numeric_array(&array, &mut ctx); +} + +/// The fused single-walk encoder must produce exactly what the multi-pass one does. +#[rstest] +#[case::varied(PrimitiveArray::from_iter(varied(100)))] +#[case::varied_exact(PrimitiveArray::from_iter(varied(0)))] +#[case::tiny(PrimitiveArray::from_iter([5u32, 1 << 20, 7]))] +#[case::nullable_signed(PrimitiveArray::new( + Buffer::from_iter(varied(50).into_iter().map(|v| v as i32)), + Validity::from_iter((0..4 * FL_CHUNK_SIZE + 50).map(|i| i % 7 != 0)), +))] +#[case::all_null(PrimitiveArray::new(Buffer::from_iter(varied(9)), Validity::AllInvalid))] +#[case::short_and_wide(short_and_wide())] +fn fused_matches_multipass(#[case] array: PrimitiveArray) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let fused = bitpack_to_best_chunk_widths(&array, &mut ctx)?; + let multipass = bitpack_to_best_chunk_widths_multipass(&array, &mut ctx)?; + assert_eq!(fused.chunk_widths(), multipass.chunk_widths()); + assert_eq!(fused.packed().as_host(), multipass.packed().as_host()); + assert_eq!( + fused.patches().map(|p| p.num_patches()), + multipass.patches().map(|p| p.num_patches()) + ); + assert_eq!(fused.nbytes(), multipass.nbytes()); + assert_arrays_eq!(fused, array, &mut ctx); + Ok(()) +} + +/// 200 u8 values needing 7 bits: the single padded block (896 bytes) is larger than the raw +/// array (200 bytes), which an encoder sizing its output by the raw length overflows. +fn short_and_wide() -> PrimitiveArray { + PrimitiveArray::from_iter((0..200u8).map(|i| i.wrapping_mul(97) % 128)) +} + +#[test] +fn short_chunk_packs_wider_than_raw() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = short_and_wide(); + let packed = bitpack_to_best_chunk_widths(&array, &mut ctx)?; + assert_eq!(packed.chunk_widths().as_slice(), &[7]); + assert!(packed.packed().len() > array.nbytes() as usize); + assert_arrays_eq!(packed, array, &mut ctx); + 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/cast.rs b/encodings/fastlanes/src/bitpacking/compute/cast.rs index f7cfc56acf0..ab04c0419c6 100644 --- a/encodings/fastlanes/src/bitpacking/compute/cast.rs +++ b/encodings/fastlanes/src/bitpacking/compute/cast.rs @@ -44,7 +44,7 @@ fn build_with_validity( .patches() .map(|patches| patches.map_values(|values| values.cast(dtype.clone()))) .transpose()?, - array.bit_width(), + array.chunk_widths().clone(), array.len(), array.offset(), )? 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..d62df6fe95a 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_widths(); 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..9f18c49efdf 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -123,21 +123,17 @@ fn filter_with_indices( 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::(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..0ee34f44985 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)] 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 996565a2672..e808168118e 100644 --- a/encodings/fastlanes/src/bitpacking/compute/slice.rs +++ b/encodings/fastlanes/src/bitpacking/compute/slice.rs @@ -54,15 +54,18 @@ fn slice_bitpacked( 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 chunk_start = block_start / 1024; + let chunk_stop = block_stop / 1024; + let widths = array.chunk_widths(); + let encoded_start = widths.byte_offset(chunk_start); + let encoded_stop = widths.byte_offset(chunk_stop); Ok(BitPacked::try_new( array.packed().slice(encoded_start..encoded_stop), array.dtype().as_ptype(), array.validity()?.slice(range.clone())?, patches, - array.bit_width(), + widths.slice(chunk_start..chunk_stop), range.len(), offset as u16, )? @@ -79,7 +82,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..a73c36c11dc 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`]. diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 86e97623cf6..36785eb280f 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 @@ -81,9 +81,6 @@ fn take_primitive( } 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 +90,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.packed_chunk::(chunk_idx); let mut have_unpacked = false; let (offset_chunks, remainder) = indices_within_chunk.as_chunks::(); diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index efa0677a91e..61912717c2c 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -7,16 +7,22 @@ pub use array::BitPackedArraySlotsExt; pub use array::BitPackedData; pub use array::BitPackedDataParts; pub use array::BitPackedSlots; +pub use array::ChunkWidths; pub use array::bitpack_compress; pub use array::bitpack_decompress; +pub use array::chunk_packed_bytes; pub use array::unpack_iter; +#[cfg(test)] +mod chunk_widths_tests; pub(crate) mod compute; mod plugin; mod vtable; pub(crate) use plugin::BitPackedPatchedPlugin; +pub use plugin::BitPackedPlugin; +pub use plugin::bitpacked_v2_id; pub use vtable::BitPacked; pub use vtable::BitPackedArray; diff --git a/encodings/fastlanes/src/bitpacking/plugin.rs b/encodings/fastlanes/src/bitpacking/plugin.rs index 3ff07db7e5c..684e814b54c 100644 --- a/encodings/fastlanes/src/bitpacking/plugin.rs +++ b/encodings/fastlanes/src/bitpacking/plugin.rs @@ -1,11 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A custom [`ArrayPlugin`] that lets you load in and deserialize a `BitPacked` array with interior -//! patches as a `PatchedArray` that wraps a patchless `BitPacked` array. +//! [`ArrayPlugin`]s for bit-packed arrays. //! -//! This enables zero-cost backward compatibility with previously written datasets. +//! [`BitPackedPlugin`] owns the wire history of `BitPacked`: the frozen `fastlanes.bitpacked` +//! format for arrays whose chunks share one width, and `fastlanes.bitpacked_v2`, whose width +//! table child carries one width per chunk. [`BitPackedPatchedPlugin`] reads both and lifts +//! interior patches into a `Patched` array. +use prost::Message; use vortex_array::Array; use vortex_array::ArrayDeserialization; use vortex_array::ArrayId; @@ -16,13 +19,156 @@ use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Patched; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::patches::PatchesMetadata; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::BitPacked; use crate::BitPackedArrayExt; +use crate::ChunkWidths; +use crate::FL_CHUNK_SIZE; +use crate::bitpacking::array::WIDTH_TABLE_DTYPE; +use crate::bitpacking::vtable::deserialize_children; +use crate::bitpacking::vtable::offset_from_metadata; +use crate::bitpacking::vtable::single_buffer; + +/// The serialized format for arrays whose chunks do not all share one bit width. +/// +/// The original `fastlanes.bitpacked` format carries a single `bit_width`, and readers of that +/// format assume every chunk uses it. Arrays with differing chunk widths therefore serialize under +/// this successor ID, which older readers reject as unknown instead of misreading. +pub fn bitpacked_v2_id() -> ArrayId { + static ID: CachedId = CachedId::new("fastlanes.bitpacked_v2"); + *ID +} + +/// Metadata of the `fastlanes.bitpacked_v2` format. The chunk widths travel in the width table +/// child, so there is no width here. +/// +/// Tag 1 is left unused: it is `bit_width` in the original format, so metadata misdirected across +/// the two IDs decodes to the right fields and fails on the child layout instead. +#[derive(Clone, prost::Message)] +pub(crate) struct BitPackedV2Metadata { + #[prost(uint32, tag = "2")] + pub(crate) offset: u32, + #[prost(message, optional, tag = "3")] + pub(crate) patches: Option, +} + +/// The [`ArrayPlugin`] for `BitPacked`, owning both of its wire formats. +/// +/// Arrays whose chunks share one width serialize through the encoding's own serializer as the +/// frozen `fastlanes.bitpacked` format, byte for byte. Differing widths serialize as +/// `fastlanes.bitpacked_v2`, with the width table as the last child. +#[derive(Debug, Clone)] +pub struct BitPackedPlugin; + +impl ArrayPlugin for BitPackedPlugin { + fn id(&self) -> ArrayId { + ArrayVTable::id(&BitPacked) + } + + fn serialized_ids(&self) -> Vec { + vec![self.id(), bitpacked_v2_id()] + } + + fn serialize( + &self, + array: &ArrayRef, + session: &VortexSession, + ) -> VortexResult> { + let view = array.as_::(); + if view.chunk_widths().is_uniform() { + return ArrayPlugin::serialize(&BitPacked, array, session); + } + let metadata = BitPackedV2Metadata { + offset: view.offset() as u32, + patches: view + .patches() + .map(|p| p.to_metadata(view.len(), view.dtype())) + .transpose()?, + } + .encode_to_vec(); + // The array's children already run patches, validity, then the width table. + Ok(Some(ArraySerialization::from_array( + bitpacked_v2_id(), + array, + metadata, + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + session: &VortexSession, + ) -> VortexResult { + if parts.serialized_id == self.id() { + return Ok(Array::::try_from_parts(ArrayVTable::deserialize( + &BitPacked, + parts.dtype, + parts.len, + parts.metadata, + parts.buffers, + parts.children, + session, + )?)? + .into_array()); + } + vortex_ensure!( + parts.serialized_id == bitpacked_v2_id(), + "BitPacked plugin does not recognize serialized ID {}", + parts.serialized_id, + ); + deserialize_v2(parts, session) + } +} + +/// Read the `fastlanes.bitpacked_v2` format: [`BitPackedV2Metadata`], one packed buffer, and +/// children running patches, validity, then the width table. +fn deserialize_v2( + parts: ArrayDeserialization<'_>, + session: &VortexSession, +) -> VortexResult { + let ArrayDeserialization { + dtype, + len, + metadata, + buffers, + children, + .. + } = parts; + let metadata = BitPackedV2Metadata::decode(metadata)?; + let packed = single_buffer(buffers)?; + let offset = offset_from_metadata(metadata.offset)?; + let num_chunks = (len + offset as usize).div_ceil(FL_CHUNK_SIZE); + let (patches, validity, table_idx) = + deserialize_children(children, metadata.patches, dtype, len, 1)?; + let table = children.get(table_idx, &WIDTH_TABLE_DTYPE, num_chunks)?; + let widths = ChunkWidths::new( + table + .clone() + .execute::(&mut session.create_execution_ctx())? + .into_buffer::(), + ); + let array = BitPacked::try_new( + packed, + dtype.as_ptype(), + validity, + patches, + widths, + len, + offset, + )?; + // A table whose entries all agree describes a uniform array, which carries no table in + // memory. Otherwise keep the table as read, since a compressor may have re-encoded it. + if array.chunk_widths().is_uniform() { + return Ok(array.into_array()); + } + Ok(BitPacked::with_width_table(array, table)?.into_array()) +} /// Custom deserialization plugin that converts a BitPacked array with interior /// Patches into a PatchedArray holding a BitPacked array. @@ -33,8 +179,11 @@ impl ArrayPlugin for BitPackedPatchedPlugin { fn id(&self) -> ArrayId { // We reuse the existing `BitPacked` ID so that we can take over its // deserialization pathway. - // TODO(joe): dedup method name - ArrayVTable::id(&BitPacked) + BitPackedPlugin.id() + } + + fn serialized_ids(&self) -> Vec { + BitPackedPlugin.serialized_ids() } fn serialize( @@ -42,8 +191,7 @@ impl ArrayPlugin for BitPackedPatchedPlugin { array: &ArrayRef, session: &VortexSession, ) -> VortexResult> { - // delegate to BitPacked VTable for serialization - ArrayPlugin::serialize(&BitPacked, array, session) + BitPackedPlugin.serialize(array, session) } fn deserialize( @@ -51,21 +199,8 @@ impl ArrayPlugin for BitPackedPatchedPlugin { parts: ArrayDeserialization<'_>, session: &VortexSession, ) -> VortexResult { - vortex_ensure!( - parts.serialized_id == self.id(), - "BitPacked plugin does not recognize serialized ID {}", - parts.serialized_id, - ); - let bitpacked = Array::::try_from_parts(ArrayVTable::deserialize( - &BitPacked, - parts.dtype, - parts.len, - parts.metadata, - parts.buffers, - parts.children, - session, - )?) - .map_err(|_| vortex_err!("BitPacked plugin should only deserialize fastlanes.bitpacked"))?; + let bitpacked = BitPackedPlugin.deserialize(parts, session)?; + let bitpacked = bitpacked.as_::().into_owned(); // Create a new BitPackedArray without the interior patches installed. let Some(patches) = bitpacked.patches() else { @@ -75,12 +210,12 @@ 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 widths = bitpacked.chunk_widths().clone(); let len = bitpacked.len(); let offset = bitpacked.offset(); let bitpacked_without_patches = - BitPacked::try_new(packed, ptype, validity, None, bw, len, offset)?.into_array(); + BitPacked::try_new(packed, ptype, validity, None, widths, len, offset)?.into_array(); let patched = Patched::from_array_and_patches( bitpacked_without_patches, diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..51b61139ffb 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -44,11 +44,15 @@ use vortex_session::registry::CachedId; use crate::BitPackedArrayExt; use crate::BitPackedData; use crate::BitPackedDataParts; -use crate::bitpack_decompress::unpack_array; -use crate::bitpack_decompress::unpack_into_primitive_builder; +use crate::ChunkWidths; +use crate::FL_CHUNK_SIZE; use crate::bitpacking::array::BitPackedSlots; use crate::bitpacking::array::BitPackedSlotsView; use crate::bitpacking::array::PATCH_SLOTS; +use crate::bitpacking::array::validate_width_table; +use crate::bitpacking::array::width_table_child; +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,6 +66,10 @@ pub(crate) fn initialize(session: &VortexSession) { kernels::initialize(session); } +/// Metadata of the frozen `fastlanes.bitpacked` format: every chunk is packed at `bit_width`. +/// +/// Arrays whose chunks differ in width serialize as `fastlanes.bitpacked_v2` through +/// `BitPackedPlugin`, which carries the widths in a child instead. #[derive(Clone, prost::Message)] pub struct BitPackedMetadata { #[prost(uint32, tag = "1")] @@ -75,7 +83,7 @@ pub struct BitPackedMetadata { impl ArrayHash for BitPackedData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.offset.hash(state); - self.bit_width.hash(state); + self.widths.hash(state); self.packed.array_hash(state, accuracy); self.patches_data.hash(state); } @@ -84,7 +92,7 @@ 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.widths == other.widths && self.packed.array_eq(&other.packed, accuracy) && self.patches_data == other.patches_data } @@ -118,10 +126,11 @@ impl VTable for BitPacked { dtype.as_ptype(), &validity, patches.as_ref(), - data.bit_width, + &data.widths, len, data.offset, - ) + )?; + validate_width_table(&data.widths, bp_slots.width_table) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -164,6 +173,12 @@ impl VTable for BitPacked { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { + // This output is labelled with the encoding's own ID, whose frozen contract is one width + // for every chunk. Differing widths need the v2 format, which only `BitPackedPlugin` + // emits. + if !array.chunk_widths().is_uniform() { + return Ok(None); + } Ok(Some( BitPackedMetadata { bit_width: array.bit_width() as u32, @@ -187,71 +202,27 @@ impl VTable for BitPacked { _session: &VortexSession, ) -> VortexResult> { let metadata = BitPackedMetadata::decode(metadata)?; - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - let packed = buffers[0].clone(); - - let load_validity = |child_idx: usize| { - if children.len() == child_idx { - Ok(Validity::from(dtype.nullability())) - } else if children.len() == child_idx + 1 { - let validity = children.get(child_idx, &Validity::DTYPE, len)?; - Ok(Validity::Array(validity)) - } else { - vortex_bail!( - "Expected {} or {} children, got {}", - child_idx, - child_idx + 1, - children.len() - ); - } - }; - - let validity_idx = match &metadata.patches { - None => 0, - Some(patches_meta) if patches_meta.chunk_offsets_dtype()?.is_some() => 3, - Some(_) => 2, - }; - - let validity = load_validity(validity_idx)?; - - let patches = metadata - .patches - .map(|p| { - let indices = children.get(0, &p.indices_dtype()?, p.len()?)?; - let values = children.get(1, dtype, p.len()?)?; - let chunk_offsets = p - .chunk_offsets_dtype()? - .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize)) - .transpose()?; - - Patches::new(len, p.offset()?, indices, values, chunk_offsets) - }) - .transpose()?; + let packed = single_buffer(buffers)?; + let offset = offset_from_metadata(metadata.offset)?; + let bit_width = u8::try_from(metadata.bit_width).map_err(|_| { + vortex_err!( + "BitPackedMetadata bit_width {} does not fit in u8", + 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(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(None); s }; - let data = BitPackedData::try_new( - packed, - patches, - u8::try_from(metadata.bit_width).map_err(|_| { - vortex_err!( - "BitPackedMetadata bit_width {} does not fit in u8", - metadata.bit_width - ) - })?, - u16::try_from(metadata.offset).map_err(|_| { - vortex_err!( - "BitPackedMetadata offset {} does not fit in u16", - metadata.offset - ) - })?, - )?; + let widths = ChunkWidths::uniform(bit_width, num_chunks); + let data = BitPackedData::try_new(packed, patches, widths, offset)?; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) } @@ -299,30 +270,115 @@ impl VTable for BitPacked { } } +/// The single packed buffer of a serialized bit-packed array. +pub(crate) fn single_buffer(buffers: &[BufferHandle]) -> VortexResult { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + Ok(buffers[0].clone()) +} + +/// The offset into the first chunk, which the metadata stores as a `u32`. +pub(crate) fn offset_from_metadata(offset: u32) -> VortexResult { + u16::try_from(offset) + .map_err(|_| vortex_err!("BitPackedMetadata offset {offset} does not fit in u16")) +} + +/// Read the patches and validity children that both wire formats share. +/// +/// Children run: the patches, then a validity bitmap if there is one, then `trailing` children +/// the caller reads itself. Returns the index of the first trailing child. +pub(crate) fn deserialize_children( + children: &dyn ArrayChildren, + patches: Option, + dtype: &DType, + len: usize, + trailing: usize, +) -> VortexResult<(Option, Validity, usize)> { + let num_patch_children = match &patches { + None => 0, + Some(patches_meta) if patches_meta.chunk_offsets_dtype()?.is_some() => 3, + Some(_) => 2, + }; + let num_fixed = num_patch_children + trailing; + let has_validity = match children.len().checked_sub(num_fixed) { + Some(0) => false, + Some(1) => true, + _ => vortex_bail!( + "Expected {num_fixed} or {} children, got {}", + num_fixed + 1, + children.len() + ), + }; + let validity = if has_validity { + Validity::Array(children.get(num_patch_children, &Validity::DTYPE, len)?) + } else { + Validity::from(dtype.nullability()) + }; + let patches = patches + .map(|p| { + let indices = children.get(0, &p.indices_dtype()?, p.len()?)?; + let values = children.get(1, dtype, p.len()?)?; + let chunk_offsets = p + .chunk_offsets_dtype()? + .map(|dtype| children.get(2, &dtype, p.chunk_offsets_len() as usize)) + .transpose()?; + Patches::new(len, p.offset()?, indices, values, chunk_offsets) + }) + .transpose()?; + Ok(( + patches, + validity, + num_patch_children + usize::from(has_validity), + )) +} + #[derive(Clone, Debug)] pub struct BitPacked; impl BitPacked { + /// Build a bit-packed array from its parts, with one width per chunk. pub fn try_new( packed: BufferHandle, ptype: PType, validity: Validity, patches: Option, - bit_width: u8, + widths: ChunkWidths, 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(width_table_child(&widths)); s }; - let data = BitPackedData::try_new(packed, patches, bit_width, offset)?; + let data = BitPackedData::try_new(packed, patches, widths, offset)?; Array::try_from_parts(ArrayParts::new(BitPacked, dtype, len, data).with_slots(slots)) } + /// Replace the width table child of an array whose chunk widths differ with `table`, an + /// array equal to the current table, so a compressor can re-encode it before the array is + /// written. The table must hold one `u8` per chunk. + 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(); + let stats = array.statistics().to_owned(); + Ok(Array::try_from_parts( + ArrayParts::new(BitPacked, dtype, len, array.into_data()).with_slots(slots), + )? + .with_stats_set(stats)) + } + pub fn into_parts(array: BitPackedArray) -> BitPackedDataParts { let len = array.len(); let patches = array.patches(); @@ -330,7 +386,7 @@ impl BitPacked { let data = array.into_data(); BitPackedDataParts { offset: data.offset, - bit_width: data.bit_width, + widths: data.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 e14b27323c1..3fbbec12b76 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/operations.rs @@ -8,8 +8,8 @@ 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 { fn scalar_at( array: ArrayView<'_, BitPacked>, @@ -54,6 +54,7 @@ mod test { use crate::BitPacked; use crate::BitPackedArray; use crate::BitPackedData; + use crate::ChunkWidths; use crate::bitpacking::array::BitPackedArrayExt; use crate::test::SESSION; @@ -253,7 +254,7 @@ mod test { ) .unwrap(), ), - 1, + ChunkWidths::uniform(1, 1), 8, 0, ) diff --git a/encodings/fastlanes/src/for/array/for_decompress.rs b/encodings/fastlanes/src/for/array/for_decompress.rs index 6e429abc419..723d576642a 100644 --- a/encodings/fastlanes/src/for/array/for_decompress.rs +++ b/encodings/fastlanes/src/for/array/for_decompress.rs @@ -103,7 +103,7 @@ pub(crate) fn fused_decompress< let mut unpacked = UnpackedChunks::try_new_with_strategy( strategy, bp.packed_slice::(), - bp.bit_width() as usize, + bp.chunk_widths(), bp.offset() as usize, bp.len(), &mut scratch, diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 43d83c6fc7f..ee7a6838883 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -84,7 +84,7 @@ pub fn initialize(session: &VortexSession) { if use_experimental_patches() { session.arrays().register(BitPackedPatchedPlugin); } else { - session.arrays().register(BitPacked); + session.arrays().register(BitPackedPlugin); } session.arrays().register(Delta); session.arrays().register(FoR); 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-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index 43f6a5a3712..8a8a2c81cb4 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -559,7 +559,11 @@ 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 (packed, widths, bitpacked_offset, patch_range) = + bitpacked_slice_view(bp, 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 +571,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 +619,9 @@ impl FusedPlan { fn walk_bitpacked(&mut self, array: ArrayRef) -> VortexResult { let bp = array.as_::(); + let Some(bit_width) = bp.chunk_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 +629,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..6126887d0a4 100644 --- a/vortex-cuda/src/kernel/encodings/bitpacked.rs +++ b/vortex-cuda/src/kernel/encodings/bitpacked.rs @@ -26,8 +26,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; @@ -53,7 +55,7 @@ pub(crate) fn bitpacked_slice_view( bp: ArrayView<'_, BitPacked>, 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 +63,15 @@ 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 widths = bp.chunk_widths(); + 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, )) @@ -90,13 +96,14 @@ 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 (packed, widths, bitpacked_offset, patch_range) = + bitpacked_slice_view(bp, offset, len)?; let sliced = BitPacked::try_new( packed, bp.ptype(bp.dtype()), child.validity()?.slice(patch_range.clone())?, bp.patches(), - bp.bit_width(), + widths, len, bitpacked_offset, )?; @@ -162,7 +169,7 @@ where { let BitPackedDataParts { offset, - bit_width, + widths, len, packed, patches, @@ -170,6 +177,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?; From dcd7e852fc27222ad99a8c4f5f76e8d120b7dcfd Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 3 Sep 2026 15:15:44 -0400 Subject: [PATCH 2/2] bench(fastlanes): sweep global against per-chunk bit widths on synthetic patterns Signed-off-by: Matt Katz --- encodings/fastlanes/Cargo.toml | 4 + .../fastlanes/benches/bitpack_chunk_widths.rs | 250 ++++++++++++++++++ 2 files changed, 254 insertions(+) create mode 100644 encodings/fastlanes/benches/bitpack_chunk_widths.rs diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..8170bb5ed12 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -70,3 +70,7 @@ harness = false name = "cast_bitpacked" harness = false required-features = ["_test-harness"] + +[[bench]] +name = "bitpack_chunk_widths" +harness = false diff --git a/encodings/fastlanes/benches/bitpack_chunk_widths.rs b/encodings/fastlanes/benches/bitpack_chunk_widths.rs new file mode 100644 index 00000000000..8f82543684f --- /dev/null +++ b/encodings/fastlanes/benches/bitpack_chunk_widths.rs @@ -0,0 +1,250 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Synthetic sweep comparing one global bit width (`bitpack_to_best_bit_width`, the +//! `fastlanes.bitpacked` wire format) against a width per 1024-element chunk +//! (`bitpack_to_best_chunk_widths`, `fastlanes.bitpacked_v2`). +//! +//! Each case is 32 FastLanes chunks of one integer type, generated from a per-chunk width +//! pattern with optional exceptions and nulls. `compress_*` benches the width selection plus +//! packing, `decompress_*` the unpack back to a primitive array. Compressed sizes for every +//! case are printed to stderr before the timings. +//! +//! Run with: cargo bench -p vortex-fastlanes --bench bitpack_chunk_widths + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::fmt; +use std::sync::LazyLock; + +use divan::Bencher; +use num_traits::NumCast; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::validity::Validity; +use vortex_buffer::BufferMut; +use vortex_fastlanes::BitPackedArray; +use vortex_fastlanes::BitPackedArrayExt; +use vortex_fastlanes::FL_CHUNK_SIZE; +use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width; +use vortex_fastlanes::bitpack_compress::bitpack_to_best_chunk_widths; +use vortex_session::VortexSession; + +const NUM_CHUNKS: usize = 32; +const LEN: usize = NUM_CHUNKS * FL_CHUNK_SIZE; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_fastlanes::initialize(&session); + session +}); + +/// How the bit width needed by a chunk's values varies across the array. +#[derive(Clone, Copy)] +enum Pattern { + /// Every chunk needs half the type width. + Uniform, + /// Widths grow linearly from 1 bit to nearly the full type width. + Drift, + /// Every chunk draws its own width at random. + Random, + /// Every other chunk is all zeros. + ZeroHeavy, + /// Narrow everywhere except one nearly full-width chunk in eight. + Spiky, +} + +impl Pattern { + fn width(self, chunk: usize, bits: usize, rng: &mut StdRng) -> usize { + match self { + Pattern::Uniform => bits / 2, + Pattern::Drift => 1 + chunk * (bits - 2) / NUM_CHUNKS, + Pattern::Random => rng.random_range(1..bits), + Pattern::ZeroHeavy => { + if chunk.is_multiple_of(2) { + 0 + } else { + bits / 2 + } + } + Pattern::Spiky => { + if chunk % 8 == 7 { + bits - 1 + } else { + bits / 4 + } + } + } + } +} + +#[derive(Clone, Copy)] +struct Case { + pattern: Pattern, + /// Fraction of values pushed above their chunk's width. + exceptions: f64, + /// Fraction of null values. + nulls: f64, +} + +impl fmt::Debug for Case { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let name = match self.pattern { + Pattern::Uniform => "uniform", + Pattern::Drift => "drift", + Pattern::Random => "random", + Pattern::ZeroHeavy => "zero_heavy", + Pattern::Spiky => "spiky", + }; + write!(f, "{name}")?; + if self.exceptions > 0.0 { + write!(f, "+exc{}%", (self.exceptions * 100.0) as usize)?; + } + if self.nulls > 0.0 { + write!(f, "+null{}%", (self.nulls * 100.0) as usize)?; + } + Ok(()) + } +} + +const fn case(pattern: Pattern, exceptions: f64, nulls: f64) -> Case { + Case { + pattern, + exceptions, + nulls, + } +} + +const CASES: &[Case] = &[ + case(Pattern::Uniform, 0.0, 0.0), + case(Pattern::Uniform, 0.01, 0.0), + case(Pattern::Uniform, 0.0, 0.1), + case(Pattern::Drift, 0.0, 0.0), + case(Pattern::Drift, 0.01, 0.0), + case(Pattern::Drift, 0.0, 0.1), + case(Pattern::Random, 0.0, 0.0), + case(Pattern::ZeroHeavy, 0.0, 0.0), + case(Pattern::Spiky, 0.0, 0.0), + case(Pattern::Spiky, 0.01, 0.0), +]; + +fn fixture(case: Case) -> PrimitiveArray { + let bits = T::PTYPE.bit_width(); + let mut rng = StdRng::seed_from_u64(42); + let mut values = BufferMut::::with_capacity(LEN); + for chunk in 0..NUM_CHUNKS { + let width = case.pattern.width(chunk, bits, &mut rng); + for _ in 0..FL_CHUNK_SIZE { + // An exception needs more bits than the chunk width but must still fit the type, so + // a chunk already at the widest supported width cannot have any. + let can_except = case.exceptions > 0.0 && width + 1 < bits; + let v: u64 = if can_except && rng.random_bool(case.exceptions) { + let exc_bits = (width + 8).min(bits - 1); + rng.random_range((1u64 << width)..(1u64 << exc_bits)) + } else if width == 0 { + 0 + } else { + rng.random_range(0..(1u64 << width)) + }; + values.push(T::from(v).unwrap()); + } + } + let validity = if case.nulls > 0.0 { + Validity::from_iter((0..LEN).map(|_| !rng.random_bool(case.nulls))) + } else { + Validity::NonNullable + }; + PrimitiveArray::new(values.freeze(), validity) +} + +fn pack_v1(array: &PrimitiveArray) -> BitPackedArray { + bitpack_to_best_bit_width(array, &mut SESSION.create_execution_ctx()).unwrap() +} + +fn pack_v2(array: &PrimitiveArray) -> BitPackedArray { + bitpack_to_best_chunk_widths(array, &mut SESSION.create_execution_ctx()).unwrap() +} + +/// Serialized bytes: packed data, patches, and one width byte per chunk when widths differ. +fn compressed_bytes(array: &BitPackedArray) -> u64 { + let widths = array.chunk_widths(); + let width_bytes = if widths.uniform_width().is_some() { + 0 + } else { + widths.len() as u64 + }; + array.nbytes() + width_bytes +} + +fn exceptions(array: &BitPackedArray) -> usize { + array.patches().map_or(0, |p| p.num_patches()) +} + +fn report_sizes() { + for &case in CASES { + let array = fixture::(case); + let v1 = pack_v1(&array); + let v2 = pack_v2(&array); + let (b1, b2) = (compressed_bytes(&v1), compressed_bytes(&v2)); + eprintln!( + "{:<4} {:<22} raw {:>7} v1 {:>7} B (bw {:>2}, {:>5} exc) v2 {:>7} B (max {:>2}, {:>5} exc) saving {:>6.2}%", + T::PTYPE, + format!("{case:?}"), + array.nbytes(), + b1, + v1.bit_width(), + exceptions(&v1), + b2, + v2.bit_width(), + exceptions(&v2), + 100.0 * (b1 as f64 - b2 as f64) / b1 as f64, + ); + } +} + +fn main() { + eprintln!("compressed sizes ({NUM_CHUNKS} chunks of {FL_CHUNK_SIZE} values per case):"); + report_sizes::(); + report_sizes::(); + report_sizes::(); + report_sizes::(); + divan::main(); +} + +#[divan::bench(types = [u8, u16, u32, u64], args = CASES)] +fn compress_v1(bencher: Bencher, case: Case) { + let array = fixture::(case); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| bitpack_to_best_bit_width(array, &mut ctx).unwrap()) +} + +#[divan::bench(types = [u8, u16, u32, u64], args = CASES)] +fn compress_v2(bencher: Bencher, case: Case) { + let array = fixture::(case); + bencher + .with_inputs(|| (&array, SESSION.create_execution_ctx())) + .bench_values(|(array, mut ctx)| bitpack_to_best_chunk_widths(array, &mut ctx).unwrap()) +} + +#[divan::bench(types = [u8, u16, u32, u64], args = CASES)] +fn decompress_v1(bencher: Bencher, case: Case) { + let packed = pack_v1(&fixture::(case)).into_array(); + bencher + .with_inputs(|| (packed.clone(), SESSION.create_execution_ctx())) + .bench_values(|(packed, mut ctx)| packed.execute::(&mut ctx).unwrap()) +} + +#[divan::bench(types = [u8, u16, u32, u64], args = CASES)] +fn decompress_v2(bencher: Bencher, case: Case) { + let packed = pack_v2(&fixture::(case)).into_array(); + bencher + .with_inputs(|| (packed.clone(), SESSION.create_execution_ctx())) + .bench_values(|(packed, mut ctx)| packed.execute::(&mut ctx).unwrap()) +}