diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index d10a70d32da..426d2a9cfd1 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,10 +15,12 @@ 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; @@ -35,15 +38,222 @@ use crate::bitpack_decompress::count_exceptions; use crate::bitpacking::array::ChunkWidths; use crate::bitpacking::array::chunk_packed_bytes; -/// Choose a cost-model width for each chunk, then pack values and gather exceptions. +/// 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 { + 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 offsets = widths.offsets_array(); + let bitpacked = BitPacked::try_new( + BufferHandle::new_host(packed), + array.ptype(), + validity, + patches, + widths.into_array(), + offsets, + len, + 0, + )?; + bitpacked.statistics().inherit_from(array.statistics()); + Ok(bitpacked) +} + +/// 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, + ctx: &mut ExecutionCtx, ) -> VortexResult { 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)); + } + }; + + // 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; + }); + + 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); + } + } + } +} + +/// 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()) } +} + /// The cost-model-optimal bit width of every 1024-element chunk of `array`. pub fn best_chunk_widths( array: ArrayView<'_, Primitive>, diff --git a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs index 8c7d330e41c..f7302b82423 100644 --- a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs +++ b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs @@ -65,6 +65,7 @@ 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::BitPackedMetadata; use crate::bitpacking::plugin::BitPackedV2Metadata; @@ -833,3 +834,55 @@ fn conformance(#[case] array: BitPackedArray) { 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(&mut SESSION.create_execution_ctx())?, + multipass.chunk_widths(&mut SESSION.create_execution_ctx())? + ); + 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(&mut SESSION.create_execution_ctx())? + .as_buffer() + .as_slice(), + &[7] + ); + assert!(packed.packed().len() > array.nbytes() as usize); + assert_arrays_eq!(packed, array, &mut ctx); + Ok(()) +}