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()) +} 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?;