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..90fb6dbf2e8 --- /dev/null +++ b/encodings/fastlanes/benches/bitpack_chunk_widths.rs @@ -0,0 +1,254 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Synthetic sweep comparing `BitPacked` (one bit width) against `BitPackedV2` (a width per +//! 1024-element chunk). +//! +//! 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::BitPackedV2Array; +use vortex_fastlanes::BitPackedV2ArrayExt; +use vortex_fastlanes::FL_CHUNK_SIZE; +use vortex_fastlanes::bitpack_compress::bitpack_to_best_bit_width; +use vortex_fastlanes::bitpacking_v2::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) -> BitPackedV2Array { + bitpack_to_best_chunk_widths(array, &mut SESSION.create_execution_ctx()).unwrap() +} + +/// Compressed bytes of a v1 array: packed data and patches. +fn v1_bytes(array: &BitPackedArray) -> u64 { + array.nbytes() +} + +/// Compressed bytes of a v2 array: packed data, patches, and one width byte per chunk. +fn v2_bytes(array: &BitPackedV2Array) -> u64 { + array.nbytes() + array.chunk_widths().len() as u64 +} + +fn v1_exceptions(array: &BitPackedArray) -> usize { + array.patches().map_or(0, |p| p.num_patches()) +} + +fn v2_exceptions(array: &BitPackedV2Array) -> 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) = (v1_bytes(&v1), v2_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(), + v1_exceptions(&v1), + b2, + v2.bit_width(), + v2_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/goldenfiles/bitpacked_v2.metadata b/encodings/fastlanes/goldenfiles/bitpacked_v2.metadata new file mode 100644 index 00000000000..868731b909c Binary files /dev/null and b/encodings/fastlanes/goldenfiles/bitpacked_v2.metadata differ diff --git a/encodings/fastlanes/src/bitpacking_v2/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking_v2/array/bitpack_compress.rs new file mode 100644 index 00000000000..a8725e756b9 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/array/bitpack_compress.rs @@ -0,0 +1,761 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use fastlanes::BitPacking; +use num_traits::PrimInt; +use num_traits::Zero; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::primitive::PrimitiveArrayExt; +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::BitPackedV2; +use crate::BitPackedV2Array; +use crate::FL_CHUNK_SIZE; +use crate::bitpack_decompress::count_exceptions; +use crate::bitpacking_v2::array::ChunkWidths; +use crate::bitpacking_v2::array::chunk_packed_bytes; + +/// 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 bitpacked = BitPackedV2::try_new( + BufferHandle::new_host(packed), + array.ptype(), + validity, + patches, + widths, + 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>, + 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 every chunk of `array` at the same `bit_width`, which must be narrower than the type. +pub fn bitpack_encode( + array: &PrimitiveArray, + bit_width: u8, + 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() + ) + } + let widths = ChunkWidths::uniform(bit_width, array.len().div_ceil(FL_CHUNK_SIZE)); + bitpack_encode_with_widths(array, widths, 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)); + } + } + } + + 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 = BitPackedV2::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 as usize, + input, + &mut output[output_len..][..packed_len], + ); + } + } + + output.freeze() +} + +/// Gather the values that do not fit their chunk's bit width into patches. +pub fn gather_patches_with_widths( + parray: &PrimitiveArray, + widths: &ChunkWidths, + num_exceptions_hint: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let patch_validity = match parray.validity()? { + Validity::NonNullable => Validity::NonNullable, + _ => Validity::AllValid, + }; + + let array_len = parray.len(); + let validity_mask = parray + .as_ref() + .validity()? + .execute_mask(parray.len(), ctx)?; + + let patches = if array_len < u8::MAX as usize { + match_each_integer_ptype!(parray.ptype(), |T| { + gather_patches_impl::( + parray.as_slice::(), + widths, + num_exceptions_hint, + patch_validity, + validity_mask, + )? + }) + } else if array_len < u16::MAX as usize { + match_each_integer_ptype!(parray.ptype(), |T| { + gather_patches_impl::( + parray.as_slice::(), + widths, + num_exceptions_hint, + patch_validity, + validity_mask, + )? + }) + } else if array_len < u32::MAX as usize { + match_each_integer_ptype!(parray.ptype(), |T| { + gather_patches_impl::( + parray.as_slice::(), + widths, + num_exceptions_hint, + patch_validity, + validity_mask, + )? + }) + } else { + match_each_integer_ptype!(parray.ptype(), |T| { + gather_patches_impl::( + parray.as_slice::(), + widths, + num_exceptions_hint, + patch_validity, + validity_mask, + )? + }) + }; + + Ok(patches) +} + +fn gather_patches_impl( + data: &[T], + widths: &ChunkWidths, + num_exceptions_hint: usize, + patch_validity: Validity, + validity_mask: Mask, +) -> VortexResult> +where + T: PrimInt + NativePType, + P: IntegerPType, +{ + 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(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.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) < overflow_leading_zeros && valid { + indices.push(P::from(idx).vortex_expect("cast index from usize")); + values.push(*value); + } + } + + if indices.is_empty() { + Ok(None) + } else { + Ok(Some(Patches::new( + data.len(), + 0, + indices.into_array(), + PrimitiveArray::new(values, patch_validity).into_array(), + Some(chunk_offsets.into_array()), + )?)) + } +} + +/// 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 +} + +/// 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 +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ChunkedArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::PrimitiveBuilder; + use vortex_buffer::Buffer; + use vortex_error::VortexError; + use vortex_error::vortex_err; + use vortex_session::VortexSession; + + use super::*; + use crate::BitPackedV2ArrayExt; + use crate::BitPackedV2Data; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[test] + 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] + fn null_patches() { + let mut ctx = SESSION.create_execution_ctx(); + let valid_values = (0..24).map(|v| v < 1 << 4).collect::>(); + let values = PrimitiveArray::new( + (0u32..24).collect::>(), + Validity::from_iter(valid_values), + ); + assert!(values.ptype().is_unsigned_int()); + let compressed = BitPackedV2Data::encode(&values.into_array(), 4, &mut ctx).unwrap(); + assert!(compressed.patches().is_none()); + assert_eq!( + (0..(1 << 4)).collect::>(), + compressed + .as_ref() + .validity() + .unwrap() + .execute_mask(compressed.as_ref().len(), &mut ctx) + .unwrap() + .to_bit_buffer() + .set_indices() + .collect::>() + ) + } + + #[test] + fn compress_signed_fails() { + let mut ctx = SESSION.create_execution_ctx(); + let values: Buffer = (-500..500).collect(); + let array = PrimitiveArray::new(values, Validity::AllValid); + assert!(array.ptype().is_signed_int()); + + let err = BitPackedV2Data::encode(&array.into_array(), 1024u32.ilog2() as u8, &mut ctx) + .unwrap_err(); + 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 chunks = (0..10) + .map(|seed| { + bitpack_encode(&patchy_nullable(100, seed), 12, &mut ctx).map(|a| a.into_array()) + }) + .collect::>>()?; + let chunked = ChunkedArray::from_iter(chunks).into_array(); + + let into_ca = chunked.clone().execute::(&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(()) + } + + fn chunk_offsets_of(values: Vec) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + let array = PrimitiveArray::from_iter(values); + let bitpacked = bitpack_encode(&array, 4, &mut ctx)?; + let patches = bitpacked + .patches() + .ok_or_else(|| vortex_err!("expected patches"))?; + patches + .chunk_offsets() + .as_ref() + .ok_or_else(|| vortex_err!("expected chunk offsets"))? + .clone() + .execute::(&mut ctx) + } + + 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() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + // chunk 0: patches at 100, 200; chunk 1: none; chunk 2: 3000; chunk 3: 3100 + assert_arrays_eq!( + chunk_offsets_of(with_patches(4096, &[100, 200, 3000, 3100]))?, + PrimitiveArray::from_iter([0u64, 2, 2, 3]), + &mut ctx + ); + // Trailing chunks without patches all point past the last patch. + assert_arrays_eq!( + chunk_offsets_of(with_patches(5120, &[100, 200, 1500]))?, + PrimitiveArray::from_iter([0u64, 2, 3, 3, 3]), + &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_v2/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs new file mode 100644 index 00000000000..b5006344f1a --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs @@ -0,0 +1,771 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::MaybeUninit; + +use fastlanes::BitPacking; +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::PrimitiveBuilder; +use vortex_array::builders::UninitRange; +use vortex_array::dtype::NativePType; +use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::patches::Patches; +use vortex_array::scalar::Scalar; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::BitPackedV2ArrayExt; +use crate::FL_CHUNK_SIZE; +use crate::bitpacking_v2::unpack_iter::BitPackedV2 as BitPackedV2Unpack; +use crate::bitpacking_v2::unpack_iter::BitUnpackedChunks; + +/// Unpacks a bit-packed array into a primitive array. +pub fn unpack_array( + array: ArrayView<'_, BitPackedV2>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match_each_integer_ptype!(array.dtype().as_ptype(), |P| { + unpack_primitive_array::

(array, ctx) + }) +} + +pub fn unpack_primitive_array( + array: ArrayView<'_, BitPackedV2>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len()); + unpack_into_primitive_builder::(array, &mut builder, ctx)?; + assert_eq!(builder.len(), array.len()); + Ok(builder.finish_into_primitive()) +} + +/// Unpack a bit-packed array directly into a same-typed `PrimitiveBuilder`. +/// +/// This is the fast path for ordinary decompression: full FastLanes chunks are unpacked straight +/// into the final output buffer, avoiding the scratch chunk and copy needed by mapped decode. +pub(crate) fn unpack_into_primitive_builder( + array: ArrayView<'_, BitPackedV2>, + builder: &mut PrimitiveBuilder, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + unpack_into_builder_with( + array, + builder, + ctx, + |v: T| v, + |chunks, output, _| { + chunks.decode_into(output); + }, + ) +} + +/// Unpack a bit-packed array of physical type `F` into a `PrimitiveBuilder`, applying `map` +/// to each value during decompression. +/// +/// Use [`unpack_into_primitive_builder`] for same-type plain decompression. This mapped path is +/// for widening casts or other element-wise transforms: each 1024-element FastLanes chunk is +/// unpacked into a cache-resident scratch buffer and written through `map` directly into the `T` +/// output, so when `F != T` no full-length `F`-typed intermediate is materialized. +/// +/// The caller must ensure that every valid source value is representable in `T` under `map`; no +/// per-value bounds check is performed. +pub(crate) fn unpack_map_into_builder( + array: ArrayView<'_, BitPackedV2>, + builder: &mut PrimitiveBuilder, + ctx: &mut ExecutionCtx, + map: M, +) -> VortexResult<()> +where + F: BitPackedV2Unpack, + T: NativePType, + M: Fn(F) -> T, +{ + unpack_into_builder_with(array, builder, ctx, map, |chunks, output, map| { + chunks.decode_map_into(output, map); + }) +} + +fn unpack_into_builder_with( + array: ArrayView<'_, BitPackedV2>, + builder: &mut PrimitiveBuilder, + ctx: &mut ExecutionCtx, + map: M, + decode: D, +) -> VortexResult<()> +where + F: BitPackedV2Unpack, + T: NativePType, + M: Fn(F) -> T, + D: FnOnce(&mut BitUnpackedChunks<'_, F>, &mut [MaybeUninit], &M), +{ + if array.is_empty() { + return Ok(()); + } + + let len = array.len(); + let mut uninit_range = builder.uninit_range(len); + + // SAFETY: We initialize all `len` values below via `decode` and the patch loop. + unsafe { + uninit_range.append_mask(&array.validity()?.execute_mask(len, ctx)?); + } + + // SAFETY: `decode` writes a value to every slot in this range. + let uninit_slice = unsafe { uninit_range.slice_uninit_mut(0, len) }; + + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + let mut chunks = array.unpacked_chunks::(&mut scratch)?; + decode(&mut chunks, uninit_slice, &map); + + if let Some(patches) = array.patches() { + apply_patches_to_uninit_range(&mut uninit_range, &patches, ctx, &map)?; + } + + // SAFETY: A correct validity mask of `len` values was set via `append_mask`, and the same + // number of values was initialized via `decode` (and overwritten by patches). + unsafe { + uninit_range.finish(); + } + Ok(()) +} + +pub(crate) fn apply_patches_to_uninit_range T>( + dst: &mut UninitRange, + patches: &Patches, + ctx: &mut ExecutionCtx, + f: F, +) -> VortexResult<()> { + assert_eq!(patches.array_len(), dst.len()); + + let indices = patches.indices().clone().execute::(ctx)?; + let values = patches.values().clone().execute::(ctx)?; + assert!(values.all_valid(ctx)?, "Patch values must be all valid"); + let values = values.as_slice::(); + + match_each_unsigned_integer_ptype!(indices.ptype(), |P| { + for (index, &value) in indices.as_slice::

().iter().zip_eq(values) { + dst.set_value( +

>::as_(*index) - patches.offset(), + f(value), + ); + } + }); + Ok(()) +} + +pub fn unpack_single(array: ArrayView<'_, BitPackedV2>, index: usize) -> Scalar { + let ptype = array.dtype().as_ptype(); + 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| { + 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") +} + +/// # Safety +/// +/// The caller must ensure the following invariants hold: +/// * `packed.len() == (length + 1023) / 1024 * 128 * bit_width` +/// * `index_to_decode < length` +/// +/// Where `length` is the length of the array/slice backed by `packed` +/// (but is not provided to this function). +pub unsafe fn unpack_single_primitive( + packed: &[T], + bit_width: usize, + index_to_decode: usize, +) -> T { + let chunk_index = index_to_decode / 1024; + let index_in_chunk = index_to_decode % 1024; + let elems_per_chunk: usize = 128 * bit_width / size_of::(); + + let packed_chunk = &packed[chunk_index * elems_per_chunk..][0..elems_per_chunk]; + unsafe { BitPacking::unchecked_unpack_single(bit_width, packed_chunk, index_in_chunk) } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::LazyLock; + + use vortex_array::Canonical; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builders::ListBuilder; + use vortex_array::builders::ListViewBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::BufferMut; + use vortex_buffer::buffer; + use vortex_session::VortexSession; + + use super::*; + use crate::BitPackedV2Array; + use crate::BitPackedV2Data; + use crate::bitpacking_v2::bitpack_compress::bitpack_encode; + + fn encode(array: &PrimitiveArray, bit_width: u8) -> BitPackedV2Array { + bitpack_encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() + } + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + fn unpack(bitpacked: &BitPackedV2Array) -> VortexResult { + unpack_array(bitpacked.as_view(), &mut SESSION.create_execution_ctx()) + } + + fn compression_roundtrip(n: usize) { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter((0..n).map(|i| (i % 2047) as u16)); + let compressed = + BitPackedV2Data::encode(&values.clone().into_array(), 11, &mut ctx).unwrap(); + assert_arrays_eq!(compressed, values, &mut ctx); + + values + .as_slice::() + .iter() + .enumerate() + .for_each(|(i, v)| { + let scalar: u16 = (&unpack_single(compressed.as_view(), i)) + .try_into() + .unwrap(); + assert_eq!(scalar, *v); + }); + } + + #[test] + fn test_compression_roundtrip_fast() { + compression_roundtrip(125); + } + + #[test] + #[cfg_attr(miri, ignore)] // This test is too slow on miri + fn test_compression_roundtrip() { + compression_roundtrip(1024); + compression_roundtrip(10_000); + compression_roundtrip(10_240); + } + + /// List builders bulk-append the source's elements into their internal elements builder. + /// Fixed-width builder appends assume the caller reserved capacity, so the list builders + /// must grow the elements builder themselves; encoded elements exercise that because + /// BitPackedV2's `append_to_builder` writes through `uninit_range`. + #[test] + fn test_list_append_grows_elements_builder() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + let elements = encode(&PrimitiveArray::from_iter((0..3072u32).map(|i| i % 256)), 8); + let element_dtype: Arc = Arc::new(PType::U32.into()); + + // 48 lists of 64 elements each. + let offsets = Buffer::from_iter((0..=48u64).map(|i| i * 64)).into_array(); + let list = ListArray::try_new( + elements.clone().into_array(), + offsets, + Validity::NonNullable, + )?; + + let mut listview_builder = ListViewBuilder::::with_capacity( + Arc::clone(&element_dtype), + Nullability::NonNullable, + 0, + 0, + ); + list.clone() + .into_array() + .append_to_builder(&mut listview_builder, &mut ctx)?; + assert_arrays_eq!(listview_builder.finish(), list, &mut ctx); + + let mut list_builder = ListBuilder::::with_capacity( + Arc::clone(&element_dtype), + Nullability::NonNullable, + 0, + 0, + ); + list.clone() + .into_array() + .append_to_builder(&mut list_builder, &mut ctx)?; + assert_arrays_eq!(list_builder.finish(), list, &mut ctx); + + // A `ListViewArray` source appended into a `ListBuilder` appends elements list by list. + let listview = ListViewArray::try_new( + elements.into_array(), + Buffer::from_iter((0..48u64).map(|i| i * 64)).into_array(), + Buffer::from_iter(std::iter::repeat_n(64u32, 48)).into_array(), + Validity::NonNullable, + )?; + let mut list_builder = + ListBuilder::::with_capacity(element_dtype, Nullability::NonNullable, 0, 0); + listview + .into_array() + .append_to_builder(&mut list_builder, &mut ctx)?; + assert_arrays_eq!(list_builder.finish(), list, &mut ctx); + + Ok(()) + } + + #[test] + fn test_all_zeros() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = buffer![0u16, 0, 0, 0] + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 0); + let actual = unpack(&bitpacked)?; + assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 0, 0, 0]), &mut ctx); + Ok(()) + } + + #[test] + fn test_simple_patches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = buffer![0u16, 1, 0, 1] + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 0); + let actual = unpack(&bitpacked)?; + assert_arrays_eq!(actual, PrimitiveArray::from_iter([0u16, 1, 0, 1]), &mut ctx); + Ok(()) + } + + #[test] + fn test_one_full_chunk() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = BufferMut::from_iter(0u16..1024) + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 10); + let actual = unpack(&bitpacked)?; + assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1024), &mut ctx); + Ok(()) + } + + #[test] + fn test_three_full_chunks_with_patches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = BufferMut::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029)) + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 10); + assert!(bitpacked.patches().is_some()); + let actual = unpack(&bitpacked)?; + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter((5u16..1029).chain(5u16..1029).chain(5u16..1029)), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_one_full_chunk_and_one_short_chunk_no_patch() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = BufferMut::from_iter(0u16..1025) + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 11); + assert!(bitpacked.patches().is_none()); + let actual = unpack(&bitpacked)?; + assert_arrays_eq!(actual, PrimitiveArray::from_iter(0u16..1025), &mut ctx); + Ok(()) + } + + #[test] + fn test_one_full_chunk_and_one_short_chunk_with_patches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = BufferMut::from_iter(512u16..1537) + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 10); + assert_eq!(bitpacked.len(), 1025); + assert!(bitpacked.patches().is_some()); + let actual = unpack(&bitpacked)?; + assert_arrays_eq!(actual, PrimitiveArray::from_iter(512u16..1537), &mut ctx); + Ok(()) + } + + #[test] + fn test_offset_and_short_chunk_and_patches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = BufferMut::from_iter(512u16..1537) + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 10); + assert_eq!(bitpacked.len(), 1025); + assert!(bitpacked.patches().is_some()); + let slice_ref = bitpacked.into_array().slice(1023..1025)?; + let actual = slice_ref.execute::(&mut ctx)?.into_primitive(); + assert_arrays_eq!(actual, PrimitiveArray::from_iter([1535u16, 1536]), &mut ctx); + Ok(()) + } + + #[test] + fn test_offset_and_short_chunk_with_chunks_between_and_patches() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let zeros = BufferMut::from_iter(512u16..2741) + .into_array() + .execute::(&mut ctx)?; + let bitpacked = encode(&zeros, 10); + assert_eq!(bitpacked.len(), 2229); + assert!(bitpacked.patches().is_some()); + let slice_ref = bitpacked.into_array().slice(1023..2049)?; + let actual = slice_ref.execute::(&mut ctx)?.into_primitive(); + assert_arrays_eq!( + actual, + PrimitiveArray::from_iter((1023u16..2049).map(|x| x + 512)), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_unpack_into_empty_array() -> VortexResult<()> { + let empty: PrimitiveArray = PrimitiveArray::from_iter(Vec::::new()); + let bitpacked = encode(&empty, 0); + + let mut builder = PrimitiveBuilder::::new(Nullability::NonNullable); + unpack_map_into_builder( + bitpacked.as_view(), + &mut builder, + &mut SESSION.create_execution_ctx(), + |v: u32| v, + )?; + + let result = builder.finish_into_primitive(); + assert_eq!( + result.len(), + 0, + "Empty array should result in empty builder" + ); + Ok(()) + } + + /// This test ensures that the mask is properly appended to the range, not the builder. + #[test] + fn test_unpack_into_with_validity_mask() -> VortexResult<()> { + // Create an array with some null values. + let values = Buffer::from_iter([1u32, 0, 3, 4, 0]); + let validity = Validity::from_iter([true, false, true, true, false]); + let array = PrimitiveArray::new(values, validity); + + // Bitpack the array. + let bitpacked = encode(&array, 3); + + // Unpack into a new builder. + let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); + unpack_map_into_builder( + bitpacked.as_view(), + &mut builder, + &mut SESSION.create_execution_ctx(), + |v: u32| v, + )?; + + let result = builder.finish_into_primitive(); + + // Verify the validity mask was correctly applied. + assert_eq!(result.len(), 5); + let mut ctx = SESSION.create_execution_ctx(); + assert!(!result.execute_scalar(0, &mut ctx)?.is_null()); + assert!(result.execute_scalar(1, &mut ctx)?.is_null()); + assert!(!result.execute_scalar(2, &mut ctx)?.is_null()); + assert!(!result.execute_scalar(3, &mut ctx)?.is_null()); + assert!(result.execute_scalar(4, &mut ctx)?.is_null()); + Ok(()) + } + + /// Test that `unpack_into` correctly handles arrays with patches. + #[test] + fn test_unpack_into_with_patches() -> VortexResult<()> { + // Create an array where most values fit in 4 bits but some need patches. + let values: Vec = (0..100) + .map(|i| if i % 20 == 0 { 1000 + i } else { i % 16 }) + .collect(); + let array = PrimitiveArray::from_iter(values.clone()); + + // Bitpack with a bit width that will require patches. + let bitpacked = encode(&array, 4); + assert!( + bitpacked.patches().is_some(), + "Should have patches for values > 15" + ); + + // Unpack into a new builder. + let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 100); + unpack_map_into_builder( + bitpacked.as_view(), + &mut builder, + &mut SESSION.create_execution_ctx(), + |v: u32| v, + )?; + + let result = builder.finish_into_primitive(); + + // Verify all values were correctly unpacked including patches. + assert_arrays_eq!( + result, + PrimitiveArray::from_iter(values), + &mut SESSION.create_execution_ctx() + ); + Ok(()) + } + + /// Test unpacking with patches at various positions. + #[test] + fn test_unpack_to_primitive_with_patches() -> VortexResult<()> { + // Create an array where patches are needed at start, middle, and end. + let values = buffer![ + 2000u32, // Patch at start + 5, 10, 15, 20, 25, 30, 3000, // Patch in middle + 35, 40, 45, 50, 55, 4000, // Patch at end + ]; + let array = PrimitiveArray::new(values, Validity::NonNullable); + + // Bitpack with a small bit width to force patches. + let bitpacked = encode(&array, 6); + assert!(bitpacked.patches().is_some(), "Should have patches"); + + // Test with a larger array with multiple patches across chunks. + let large_values: Vec = (0..3072) + .map(|i| { + if i % 500 == 0 { + 2000 + i as u16 // Values that need patches + } else { + (i % 256) as u16 // Values that fit in 8 bits + } + }) + .collect(); + let large_array = PrimitiveArray::from_iter(large_values); + let large_bitpacked = encode(&large_array, 8); + assert!(large_bitpacked.patches().is_some()); + + let large_result = unpack(&large_bitpacked)?; + assert_eq!(large_result.len(), 3072); + Ok(()) + } + + /// Test unpacking with nullability and validity masks. + #[test] + fn test_unpack_to_primitive_nullability() { + // Test with null values at various positions. + let values = Buffer::from_iter([100u32, 0, 200, 0, 300, 0, 400]); + let validity = Validity::from_iter([true, false, true, false, true, false, true]); + let array = PrimitiveArray::new(values, validity); + + let bitpacked = encode(&array, 9); + let result = unpack(&bitpacked).vortex_expect("unpack"); + + // Verify length. + assert_eq!(result.len(), 7); + // Validity should be preserved when unpacking. + let mut ctx = SESSION.create_execution_ctx(); + assert!(!result.execute_scalar(0, &mut ctx).unwrap().is_null()); + assert!(result.execute_scalar(1, &mut ctx).unwrap().is_null()); + assert!(!result.execute_scalar(2, &mut ctx).unwrap().is_null()); + + // Test combining patches with nullability. + let patch_values = Buffer::from_iter([10u16, 0, 2000, 0, 30, 3000, 0]); + let patch_validity = Validity::from_iter([true, false, true, false, true, true, false]); + let patch_array = PrimitiveArray::new(patch_values, patch_validity); + + let patch_bitpacked = encode(&patch_array, 5); + assert!(patch_bitpacked.patches().is_some()); + + let patch_result = unpack(&patch_bitpacked).vortex_expect("unpack"); + assert_eq!(patch_result.len(), 7); + + // Test all nulls edge case. + let all_nulls = PrimitiveArray::new( + Buffer::from_iter([0u32, 0, 0, 0]), + Validity::from_iter([false, false, false, false]), + ); + let all_nulls_bp = encode(&all_nulls, 0); + let all_nulls_result = unpack(&all_nulls_bp).vortex_expect("unpack"); + assert_eq!(all_nulls_result.len(), 4); + } + + /// Test that the execute method produces consistent results with other unpacking methods. + #[test] + fn test_execute_method_consistency() -> VortexResult<()> { + // Test that execute(), unpack_to_primitive(), and unpack_array() all produce consistent results. + let test_consistency = |array: &PrimitiveArray, bit_width: u8| -> VortexResult<()> { + let bitpacked = encode(array, bit_width); + + let unpacked_array = unpack(&bitpacked)?; + + let executed = { + let mut ctx = SESSION.create_execution_ctx(); + bitpacked.into_array().execute::(&mut ctx)? + }; + + assert_eq!( + unpacked_array.len(), + array.len(), + "unpacked array length mismatch" + ); + + // The executed canonical should also have the correct length. + let executed_primitive = executed.into_primitive(); + assert_eq!( + executed_primitive.len(), + array.len(), + "executed primitive length mismatch" + ); + + // Verify that the execute() method works correctly by comparing with unpack_array. + // We convert unpack_array result to canonical to compare. + let unpacked_executed = { + let mut ctx = SESSION.create_execution_ctx(); + unpacked_array + .into_array() + .execute::(&mut ctx)? + .into_primitive() + }; + assert_eq!( + executed_primitive.len(), + unpacked_executed.len(), + "execute() and unpack_array().execute() produced different lengths" + ); + // Both should produce identical arrays since they represent the same data. + Ok(()) + }; + + // Test various scenarios without patches. + test_consistency(&PrimitiveArray::from_iter(0u16..100), 7)?; + test_consistency(&PrimitiveArray::from_iter(0u32..1024), 10)?; + + // Test with values that will create patches. + test_consistency(&PrimitiveArray::from_iter((0i16..2048).map(|x| x % 128)), 7)?; + + // Test with an array that definitely has patches. + let patch_values: Vec = (0..100) + .map(|i| if i % 20 == 0 { 1000 + i } else { i % 16 }) + .collect(); + let patch_array = PrimitiveArray::from_iter(patch_values); + test_consistency(&patch_array, 4)?; + + // Test with sliced array (offset > 0). + let values = PrimitiveArray::from_iter(0u32..2048); + let bitpacked = encode(&values, 11); + let slice_ref = bitpacked.into_array().slice(500..1500)?; + let sliced = { + let mut ctx = SESSION.create_execution_ctx(); + slice_ref + .clone() + .execute::(&mut ctx)? + .into_primitive() + }; + + // Test all three methods on the sliced array. + let primitive_result = sliced.clone(); + let unpacked_array = sliced; + let executed = { + let mut ctx = SESSION.create_execution_ctx(); + slice_ref.execute::(&mut ctx)? + }; + + assert_eq!( + primitive_result.len(), + 1000, + "sliced primitive length should be 1000" + ); + assert_eq!( + unpacked_array.len(), + 1000, + "sliced unpacked array length should be 1000" + ); + + let executed_primitive = executed.into_primitive(); + assert_eq!( + executed_primitive.len(), + 1000, + "sliced executed primitive length should be 1000" + ); + Ok(()) + } + + /// Test edge cases for unpacking. + #[test] + fn test_unpack_edge_cases() -> VortexResult<()> { + // Empty array. + let empty: PrimitiveArray = PrimitiveArray::from_iter(Vec::::new()); + let empty_bp = encode(&empty, 0); + let empty_result = unpack(&empty_bp)?; + assert_eq!(empty_result.len(), 0); + + // All zeros (bit_width = 0). + let zeros = PrimitiveArray::from_iter([0u32; 100]); + let zeros_bp = encode(&zeros, 0); + let zeros_result = unpack(&zeros_bp)?; + assert_eq!(zeros_result.len(), 100); + // Verify consistency with unpack_array. + let zeros_array = unpack(&zeros_bp)?; + assert_eq!(zeros_result.len(), zeros_array.len()); + assert_arrays_eq!( + zeros_result, + zeros_array, + &mut SESSION.create_execution_ctx() + ); + + // Maximum bit width for u16 (15 bits, since bitpacking requires bit_width < type bit width). + let max_values = PrimitiveArray::from_iter([32767u16; 50]); // 2^15 - 1 + let max_bp = encode(&max_values, 15); + let max_result = unpack(&max_bp)?; + assert_eq!(max_result.len(), 50); + + // Exactly 3072 elements with patches across chunks. + let boundary_values: Vec = (0..3072) + .map(|i| { + if i == 1023 || i == 1024 || i == 2047 || i == 2048 { + 50000 // Force patches at chunk boundaries + } else { + (i % 128) as u32 + } + }) + .collect(); + let boundary_array = PrimitiveArray::from_iter(boundary_values); + let boundary_bp = encode(&boundary_array, 7); + assert!(boundary_bp.patches().is_some()); + + let boundary_result = unpack(&boundary_bp)?; + assert_eq!(boundary_result.len(), 3072); + // Verify consistency. + let boundary_unpacked = unpack(&boundary_bp)?; + assert_eq!(boundary_result.len(), boundary_unpacked.len()); + assert_arrays_eq!( + boundary_result, + boundary_unpacked, + &mut SESSION.create_execution_ctx() + ); + + // Single element. + let single = PrimitiveArray::from_iter([42u8]); + let single_bp = encode(&single, 6); + let single_result = unpack(&single_bp)?; + assert_eq!(single_result.len(), 1); + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/array/mod.rs b/encodings/fastlanes/src/bitpacking_v2/array/mod.rs new file mode 100644 index 00000000000..d2b06ed93ca --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/array/mod.rs @@ -0,0 +1,595 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +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::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::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; + +pub mod bitpack_compress; +pub mod bitpack_decompress; +pub mod unpack_iter; + +use crate::BitPackedV2Array; +use crate::FL_CHUNK_SIZE; +use crate::bitpacking_v2::bitpack_compress::bitpack_encode; +use crate::bitpacking_v2::unpack_iter::BitPackedV2 as BitPackedV2Iter; +use crate::bitpacking_v2::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) + } + + /// 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_widths: {} chunks x {w}", self.len()), + None => write!( + f, + "bit_widths: {} chunks, max {}", + self.len(), + self.max_width + ), + } + } +} + +#[array_slots(crate::BitPackedV2)] +pub struct BitPackedV2Slots { + /// The indices of exception values that don't fit in the bit-packed representation. + #[slot(0)] + pub patch_indices: Option, + /// The exception values that don't fit in the bit-packed representation. + #[slot(1)] + pub patch_values: Option, + /// Chunk offsets for the patch indices/values. + #[slot(2)] + pub patch_chunk_offsets: Option, + /// The validity bitmap indicating which elements are non-null. + #[slot(3)] + pub validity_child: Option, +} + +pub(crate) const PATCH_SLOTS: PatchSlotIndices = PatchSlotIndices { + indices: BitPackedV2Slots::PATCH_INDICES, + values: BitPackedV2Slots::PATCH_VALUES, + chunk_offsets: BitPackedV2Slots::PATCH_CHUNK_OFFSETS, +}; + +pub struct BitPackedV2DataParts { + pub offset: u16, + pub widths: ChunkWidths, + pub len: usize, + pub packed: BufferHandle, + pub patches: Option, + pub validity: Validity, +} + +#[derive(Clone, Debug)] +pub struct BitPackedV2Data { + /// The offset within the first block (created with a slice). + /// 0 <= offset < 1024 + pub(super) offset: u16, + pub(super) widths: ChunkWidths, + pub(super) packed: BufferHandle, + /// Patch metadata for reconstructing Patches from slots. + pub(super) patches_data: Option, +} + +impl Display for BitPackedV2Data { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}, offset: {}", self.widths, self.offset) + } +} + +impl BitPackedV2Data { + /// Create a new bitpacked array using a buffer of packed data. + /// + /// 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 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 `BitPackedV2Array` from parts. + /// + /// See also the [`encode`][Self::encode] method on this type for a safe path to create a new + /// bit-packed array. + /// + /// # 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` + /// * `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, + widths: ChunkWidths, + offset: u16, + ) -> VortexResult { + vortex_ensure!( + 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, + widths, + packed, + patches_data: patches.as_ref().map(PatchesData::from_patches), + }) + } + + pub(crate) fn validate( + packed: &BufferHandle, + ptype: PType, + validity: &Validity, + patches: Option<&Patches>, + widths: &ChunkWidths, + length: usize, + offset: u16, + ) -> VortexResult<()> { + vortex_ensure!(ptype.is_int(), MismatchedTypes: "integer", ptype); + 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!( + validity_len == length, + "BitPackedV2Array validity length {validity_len} != array length {length}", + ); + } + + // Validate patches + if let Some(patches) = patches { + Self::validate_patches(patches, ptype, length)?; + } + + // 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 {}", + expected_packed_len, + packed.len() + ); + + Ok(()) + } + + fn validate_patches(patches: &Patches, ptype: PType, len: usize) -> VortexResult<()> { + // Ensure that array and patches have same ptype + vortex_ensure!( + patches.dtype().eq_ignore_nullability(ptype.into()), + "Patches DType {} does not match BitPackedV2Array dtype {}", + patches.dtype().as_nonnullable(), + ptype + ); + + vortex_ensure!( + patches.array_len() == len, + "BitPackedV2Array patches length {} != expected {len}", + patches.array_len(), + ); + + Ok(()) + } + + pub fn ptype(&self, dtype: &DType) -> PType { + dtype.as_ptype() + } + + /// Underlying bit packed values as byte array + #[inline] + pub fn packed(&self) -> &BufferHandle { + &self.packed + } + + /// Access the slice of packed values as an array of `T` + #[inline] + pub fn packed_slice(&self) -> &[T] { + let packed_bytes = self.packed().as_host(); + let packed_ptr: *const T = packed_bytes.as_ptr().cast(); + // Return number of elements of type `T` packed in the buffer + let packed_len = packed_bytes.len() / size_of::(); + + // SAFETY: as_slice points to buffer memory that outlives the lifetime of `self`. + // Unfortunately Rust cannot understand this, so we reconstruct the slice from raw parts + // to get it to reinterpret the lifetime. + 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: BitPackedV2Iter>( + &'a self, + dtype: &DType, + len: usize, + scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], + ) -> VortexResult> { + assert_eq!( + T::PTYPE, + self.ptype(dtype), + "Requested type doesn't match the array ptype" + ); + BitUnpackedChunks::try_new(self, len, scratch) + } + + /// 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.widths.max_width() + } + + #[inline] + pub fn offset(&self) -> u16 { + self.offset + } + + /// Bit-pack an array of primitive integers down to the target bit-width using the FastLanes + /// SIMD-accelerated packing kernels. + /// + /// # Errors + /// + /// If the provided array is not an integer type, an error will be returned. + /// + /// If the provided array contains negative values, an error will be returned. + /// + /// If the requested bit-width for packing is larger than the array's native width, an + /// error will be returned. + pub fn encode( + array: &ArrayRef, + bit_width: u8, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let parray: PrimitiveArray = array + .clone() + .try_downcast::() + .map_err(|a| vortex_err!(InvalidArgument: "Bitpacking can only encode primitive arrays, got {}", a.encoding_id()))?; + bitpack_encode(&parray, bit_width, ctx) + } + + /// 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 { + let bit_width = self.bit_width() as u32; + if bit_width >= usize::BITS { + usize::MAX + } else { + (1usize << bit_width) - 1 + } + } +} + +pub trait BitPackedV2ArrayExt: BitPackedV2ArraySlotsExt { + #[inline] + fn packed(&self) -> &BufferHandle { + BitPackedV2Data::packed(self) + } + + #[inline] + fn chunk_widths(&self) -> &ChunkWidths { + BitPackedV2Data::chunk_widths(self) + } + + #[inline] + fn bit_width(&self) -> u8 { + BitPackedV2Data::bit_width(self) + } + + #[inline] + fn offset(&self) -> u16 { + BitPackedV2Data::offset(self) + } + + #[inline] + fn patches(&self) -> Option { + PatchesData::patches_from_slots( + self.patches_data.as_ref(), + self.as_ref().len(), + self.as_ref().slots(), + PATCH_SLOTS, + ) + } + + #[inline] + fn validity(&self) -> Validity { + child_to_validity(self.validity_child(), self.as_ref().dtype().nullability()) + } + + #[inline] + fn packed_slice(&self) -> &[T] { + BitPackedV2Data::packed_slice::(self) + } + + #[inline] + fn packed_chunk(&self, chunk: usize) -> (&[T], usize) { + BitPackedV2Data::packed_chunk::(self, chunk) + } + + #[inline] + fn unpacked_chunks<'a, T: BitPackedV2Iter>( + &'a self, + scratch: &'a mut [MaybeUninit; FL_CHUNK_SIZE], + ) -> VortexResult> { + BitPackedV2Data::unpacked_chunks::( + self, + self.as_ref().dtype(), + self.as_ref().len(), + scratch, + ) + } +} + +impl> BitPackedV2ArrayExt for T {} + +#[cfg(test)] +mod test { + use std::sync::LazyLock; + + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + 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::BitPackedV2Data; + use crate::bitpacking_v2::array::BitPackedV2ArrayExt; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[test] + fn test_encode() { + let mut ctx = SESSION.create_execution_ctx(); + let values = [ + Some(1u64), + None, + Some(1), + None, + Some(1), + None, + Some(u64::MAX), + ]; + let uncompressed = PrimitiveArray::from_option_iter(values); + let packed = BitPackedV2Data::encode(&uncompressed.into_array(), 1, &mut ctx).unwrap(); + let expected = PrimitiveArray::from_option_iter(values); + let packed_primitive = packed + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); + assert_arrays_eq!(packed_primitive, expected, &mut ctx); + } + + #[test] + fn test_encode_too_wide() { + let mut ctx = SESSION.create_execution_ctx(); + let values = [Some(1u8), None, Some(1), None, Some(1), None]; + let uncompressed = PrimitiveArray::from_option_iter(values); + let _packed = BitPackedV2Data::encode(&uncompressed.clone().into_array(), 8, &mut ctx) + .expect_err("Cannot pack value into the same width"); + let _packed = BitPackedV2Data::encode(&uncompressed.into_array(), 9, &mut ctx) + .expect_err("Cannot pack value into larger width"); + } + + #[test] + fn signed_with_patches() { + let mut ctx = SESSION.create_execution_ctx(); + let values: Buffer = (0i32..=512).collect(); + let parray = values.clone().into_array(); + + let packed_with_patches = BitPackedV2Data::encode(&parray, 9, &mut ctx).unwrap(); + assert!(packed_with_patches.patches().is_some()); + let packed_primitive = packed_with_patches + .as_array() + .clone() + .execute::(&mut ctx) + .unwrap(); + assert_arrays_eq!( + packed_primitive, + PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable), + &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_v2/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs new file mode 100644 index 00000000000..fdc8f17b338 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs @@ -0,0 +1,436 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem; +use std::mem::MaybeUninit; +use std::ops::Range; + +use fastlanes::BitPacking; +use lending_iterator::gat; +use lending_iterator::prelude::Item; +#[gat(Item)] +use lending_iterator::prelude::LendingIterator; +use vortex_array::dtype::PhysicalPType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::BitPackedV2Data; +use crate::FL_CHUNK_SIZE; +use crate::bitpacking_v2::array::ChunkWidths; +use crate::bitpacking_v2::array::chunk_packed_bytes; + +const CHUNK_SIZE: usize = FL_CHUNK_SIZE; + +pub use crate::unpack_iter::BitPacked as BitPackedV2; +pub use crate::unpack_iter::BitPackingStrategy; +pub use crate::unpack_iter::UnpackStrategy; + +/// 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 +/// ``` +/// use std::mem::MaybeUninit; +/// +/// use lending_iterator::gat; +/// use lending_iterator::prelude::Item; +/// #[gat(Item)] +/// use lending_iterator::prelude::LendingIterator; +/// use vortex_array::IntoArray; +/// use vortex_array::VortexSessionExecute; +/// use vortex_buffer::buffer; +/// use vortex_fastlanes::BitPackedV2Data; +/// use vortex_fastlanes::BitPackedV2ArrayExt; +/// use vortex_fastlanes::FL_CHUNK_SIZE; +/// use vortex_fastlanes::bitpacking_v2::unpack_iter::BitUnpackedChunks; +/// +/// let mut ctx = vortex_array::array_session().create_execution_ctx(); +/// let array = BitPackedV2Data::encode(&buffer![2, 3, 4, 5].into_array(), 2, &mut ctx).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 +/// } +/// +/// if let Some(trailer) = unpacked_chunks.trailer() { +/// // handle partial trailing chunk +/// } +/// ``` +pub struct UnpackedChunks<'a, T: PhysicalPType, S: UnpackStrategy> { + strategy: S, + widths: &'a ChunkWidths, + offset: usize, + len: usize, + num_chunks: usize, + // 0 indicates full chunk of CHUNK_SIZE + last_chunk_length: usize, + packed: &'a [T::Physical], + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], +} + +pub type BitUnpackedChunks<'a, T> = UnpackedChunks<'a, T, BitPackingStrategy>; + +impl<'a, T: BitPackedV2> BitUnpackedChunks<'a, T> { + pub fn try_new( + array: &'a BitPackedV2Data, + len: usize, + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], + ) -> VortexResult { + Self::try_new_with_strategy( + BitPackingStrategy, + array.packed_slice::(), + array.chunk_widths(), + array.offset() as usize, + len, + scratch, + ) + } + + pub fn full_chunks(&mut self) -> BitUnpackIterator<'_, T> { + 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.num_chunks - last_chunk_is_sliced, + first_chunk_is_sliced, + ) + } +} + +impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { + pub fn try_new_with_strategy( + strategy: S, + packed: &'a [T::Physical], + widths: &'a ChunkWidths, + offset: usize, + len: usize, + scratch: &'a mut [MaybeUninit; CHUNK_SIZE], + ) -> VortexResult { + let (num_chunks, last_chunk_length) = + validate_packed::(packed.len(), widths, offset, len)?; + Ok(Self { + strategy, + widths, + offset, + len, + num_chunks, + last_chunk_length, + packed, + scratch, + }) + } + + #[allow(clippy::inline_always)] + #[inline(always)] + 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, bit_width) = self.chunk(0); + let dst: &mut [MaybeUninit] = self.scratch; + let dst: &mut [T::Physical] = unsafe { mem::transmute(dst) }; + + let header_end_slice = if self.num_chunks == 1 { + self.len + } else { + CHUNK_SIZE - self.offset + }; + // SAFETY: + // 1. chunk holds exactly one packed block at bit_width. + // 2. buffer is exactly CHUNK_SIZE. + unsafe { + self.strategy.unpack_chunk(bit_width, chunk, dst); + mem::transmute(&mut self.scratch[self.offset..][..header_end_slice]) + } + }) + } + + /// Decode all chunks (initial, full, and trailer) directly into the output range. + pub fn decode_into(&mut self, output: &mut [MaybeUninit]) { + debug_assert_eq!(output.len(), self.len); + let mut local_idx = 0; + + if let Some(initial) = self.initial() { + local_idx = initial.len(); + + // TODO(connor): use maybe_uninit_write_slice when it gets stabilized. + // SAFETY: &[T] and &[MaybeUninit] have the same layout. + let init_initial: &[MaybeUninit] = unsafe { mem::transmute(initial) }; + output[..local_idx].copy_from_slice(init_initial); + } + + local_idx = self.decode_full_chunks_into_at(output, local_idx); + + if let Some(trailer) = self.trailer() { + // TODO(connor): use maybe_uninit_write_slice when it gets stabilized. + // SAFETY: &[T] and &[MaybeUninit] have the same layout. + let init_trailer: &[MaybeUninit] = unsafe { mem::transmute(trailer) }; + output[local_idx..][..init_trailer.len()].copy_from_slice(init_trailer); + local_idx += init_trailer.len(); + } + + debug_assert_eq!(local_idx, self.len); + } + + /// Decode all chunks (initial, full, and trailer), mapping each unpacked value through f. + pub(crate) fn decode_map_into( + &mut self, + output: &mut [MaybeUninit], + mut f: impl FnMut(T) -> U, + ) { + debug_assert_eq!(output.len(), self.len); + + self.for_each_unpacked_chunk(|chunk, range| { + write_map(chunk, &mut output[range], &mut f); + }); + } + + /// Walk every unpacked chunk in array order, reusing the internal scratch buffer. + pub(crate) fn for_each_unpacked_chunk(&mut self, mut f: F) + where + F: FnMut(&mut [T], Range), + { + let mut local_idx = 0; + + if let Some(initial) = self.initial() { + let chunk_len = initial.len(); + f(initial, local_idx..local_idx + chunk_len); + local_idx += chunk_len; + } + + if self.num_chunks > 1 { + 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(bit_width as usize, chunk, dst); + let unpacked: &mut [T] = mem::transmute(&mut self.scratch[..]); + f(unpacked, local_idx..local_idx + CHUNK_SIZE); + } + local_idx += CHUNK_SIZE; + } + } + + if let Some(trailer) = self.trailer() { + let chunk_len = trailer.len(); + f(trailer, local_idx..local_idx + chunk_len); + local_idx += chunk_len; + } + + debug_assert_eq!(local_idx, self.len); + } + + /// Unpack full chunks into output range starting at the given index. + fn decode_full_chunks_into_at( + &mut self, + output: &mut [MaybeUninit], + start_idx: usize, + ) -> usize { + if self.num_chunks == 1 { + return start_idx; + } + + let mut local_idx = start_idx; + + 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(bit_width as usize, chunk, dst); + } + local_idx += CHUNK_SIZE; + } + local_idx + } + + fn full_chunks_range(&self) -> Range { + (self.first_chunk_is_sliced() as usize) + ..(self.num_chunks - self.last_chunk_is_sliced() as usize) + } + + /// 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, 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 holds exactly one packed block at bit_width. + // 2. buffer is exactly CHUNK_SIZE. + unsafe { + self.strategy.unpack_chunk(bit_width, chunk, dst); + mem::transmute(&mut self.scratch[..self.last_chunk_length]) + } + }) + } + + fn last_chunk_is_sliced(&self) -> bool { + self.last_chunk_length != 0 + } + + fn first_chunk_is_sliced(&self) -> bool { + self.offset != 0 + } +} + +/// 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], + widths: &ChunkWidths, + offset: usize, + len: usize, + mut f: F, +) -> VortexResult<()> +where + T: PhysicalPType, + F: FnMut(&[T::Physical], usize, Range), +{ + validate_packed::(packed.len(), widths, offset, len)?; + let padded_len = offset + len; + 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, + widths: &ChunkWidths, + offset: usize, + len: usize, +) -> VortexResult<(usize, usize)> { + vortex_ensure!( + offset < CHUNK_SIZE, + "Invalid bit-packed offset {offset}, expected < {CHUNK_SIZE}" + ); + let num_chunks = (offset + len).div_ceil(CHUNK_SIZE); + vortex_ensure!( + 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)) +} + +/// Iterator over full chunks of bitpacked array that yields unpacked chunks one at a time +pub struct BitUnpackIterator<'a, T: BitPackedV2 + 'a> { + packed: &'a [T::Physical], + widths: &'a ChunkWidths, + buffer: &'a mut [MaybeUninit; CHUNK_SIZE], + num_chunks: usize, + idx: usize, + /// Word offset of chunk `idx` within `packed`. + start: usize, +} + +impl<'a, T: BitPackedV2> BitUnpackIterator<'a, T> { + pub fn new( + packed: &'a [T::Physical], + widths: &'a ChunkWidths, + buffer: &'a mut [MaybeUninit; CHUNK_SIZE], + num_chunks: usize, + first_chunk_is_sliced: bool, + ) -> Self { + let idx = if first_chunk_is_sliced { 1 } else { 0 }; + Self { + packed, + widths, + buffer, + num_chunks, + idx, + start: widths.byte_offset(idx) / size_of::(), + } + } +} + +#[gat] +impl<'a, T: BitPackedV2 + 'a> LendingIterator for BitUnpackIterator<'a, T> { + type Item<'next> + where + Self: 'next, + = &'next mut [T; CHUNK_SIZE]; + + fn next(&'_ mut self) -> Option> { + if self.idx >= self.num_chunks { + return None; + } + + 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(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) }) + } +} + +fn write_map(src: &[T], dst: &mut [MaybeUninit], f: &mut impl FnMut(T) -> U) { + for (dst, &src) in dst.iter_mut().zip(src.iter()) { + dst.write(f(src)); + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/chunk_widths_tests.rs b/encodings/fastlanes/src/bitpacking_v2/chunk_widths_tests.rs new file mode 100644 index 00000000000..ec7a1007970 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/chunk_widths_tests.rs @@ -0,0 +1,428 @@ +// 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 rstest::rstest; +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::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::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use crate::BitPackedV2; +use crate::BitPackedV2Array; +use crate::BitPackedV2ArrayExt; +use crate::ChunkWidths; +use crate::FL_CHUNK_SIZE; +use crate::FoR; +use crate::bitpacking_v2::bitpack_compress::bitpack_encode_with_widths; +use crate::bitpacking_v2::bitpack_compress::bitpack_to_best_chunk_widths; +use crate::bitpacking_v2::bitpack_compress::bitpack_to_best_chunk_widths_multipass; + +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(()) +} + +#[test] +fn serde_roundtrip_keeps_widths() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?; + let array = packed.as_array(); + + let serialization = SESSION.array_serialize(array)?.unwrap(); + let children = array.children(); + let buffers = array + .buffers() + .into_iter() + .map(vortex_array::buffer::BufferHandle::new_host) + .collect::>(); + let deserialized = BitPackedV2Array::try_from_parts(ArrayVTable::deserialize( + &BitPackedV2, + array.dtype(), + array.len(), + &serialization.metadata, + &buffers, + &children, + &SESSION, + )?) + .map_err(|_| vortex_err!("expected fastlanes.bitpacked"))?; + + assert_eq!(deserialized.chunk_widths(), packed.chunk_widths()); + assert_arrays_eq!(deserialized, primitive(&values), &mut ctx); + Ok(()) +} + +#[rstest] +#[case::varied(encode(&varied(100)).unwrap())] +#[case::varied_exact(encode(&varied(0)).unwrap())] +fn conformance(#[case] array: BitPackedV2Array) { + 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_v2/compute/between.rs b/encodings/fastlanes/src/bitpacking_v2/compute/between.rs new file mode 100644 index 00000000000..f8ca755e182 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/between.rs @@ -0,0 +1,259 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Block-streaming between kernel for [`BitPackedV2Array`] against constant bounds. +//! +//! Reuses the same single-block scratch buffer as the compare kernel and folds a +//! `lower op_l v op_u upper` predicate per element, so the full primitive never +//! materialises. +//! +//! [`BitPackedV2Array`]: crate::BitPackedV2Array + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::between::BetweenKernel; +use vortex_array::scalar_fn::fns::between::BetweenOptions; +use vortex_array::scalar_fn::fns::between::StrictComparison; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::bitpacking_v2::compute::stream_predicate::stream_predicate; +use crate::bitpacking_v2::unpack_iter::BitPackedV2 as BitPackedV2Iter; + +impl BetweenKernel for BitPackedV2 { + fn between( + array: ArrayView<'_, Self>, + lower: &ArrayRef, + upper: &ArrayRef, + options: &BetweenOptions, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Only accelerate constant-bounds between; vary-by-row bounds fall through to the + // default `compare + and` pipeline. + let (Some(lower_const), Some(upper_const)) = (lower.as_constant(), upper.as_constant()) + else { + return Ok(None); + }; + let (Some(lower_prim), Some(upper_prim)) = ( + lower_const.as_primitive_opt(), + upper_const.as_primitive_opt(), + ) else { + return Ok(None); + }; + + let nullability = + array.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability(); + let arr_ptype = array.dtype().as_ptype(); + if lower_prim.ptype() != arr_ptype || upper_prim.ptype() != arr_ptype { + return Ok(None); + } + + let result = match_each_integer_ptype!(arr_ptype, |T| { + let lo: T = lower_prim + .typed_value::() + .vortex_expect("the between short circuit strips a null lower bound"); + let up: T = upper_prim + .typed_value::() + .vortex_expect("the between short circuit strips a null upper bound"); + between_constant_typed::(array, lo, up, options, nullability, ctx)? + }); + Ok(Some(result)) + } +} + +fn between_constant_typed( + array: ArrayView<'_, BitPackedV2>, + lower: T, + upper: T, + options: &BetweenOptions, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType + Copy + BitPackedV2Iter, +{ + // 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`. + match (options.lower_strict, options.upper_strict) { + (StrictComparison::Strict, StrictComparison::Strict) => between_impl( + array, + lower, + NativePType::is_lt, + upper, + NativePType::is_lt, + nullability, + ctx, + ), + (StrictComparison::Strict, StrictComparison::NonStrict) => between_impl( + array, + lower, + NativePType::is_lt, + upper, + NativePType::is_le, + nullability, + ctx, + ), + (StrictComparison::NonStrict, StrictComparison::Strict) => between_impl( + array, + lower, + NativePType::is_le, + upper, + NativePType::is_lt, + nullability, + ctx, + ), + (StrictComparison::NonStrict, StrictComparison::NonStrict) => between_impl( + array, + lower, + NativePType::is_le, + upper, + NativePType::is_le, + nullability, + ctx, + ), + } +} + +fn between_impl( + array: ArrayView<'_, BitPackedV2>, + lower: T, + lower_fn: Lo, + upper: T, + upper_fn: Up, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType + Copy + BitPackedV2Iter, + Lo: Fn(T, T) -> bool, + Up: Fn(T, T) -> bool, +{ + stream_predicate::( + array, + nullability, + |v| lower_fn(lower, v) & upper_fn(v, upper), + ctx, + ) +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builtins::ArrayBuiltins; + use vortex_array::scalar_fn::fns::between::BetweenOptions; + use vortex_array::scalar_fn::fns::between::StrictComparison; + use vortex_array::validity::Validity; + use vortex_buffer::BufferMut; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::BitPackedV2ArrayExt; + use crate::BitPackedV2Data; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + fn opts(lower: StrictComparison, upper: StrictComparison) -> BetweenOptions { + BetweenOptions { + lower_strict: lower, + upper_strict: upper, + } + } + + #[rstest] + #[case(StrictComparison::NonStrict, StrictComparison::NonStrict)] + #[case(StrictComparison::Strict, StrictComparison::NonStrict)] + #[case(StrictComparison::NonStrict, StrictComparison::Strict)] + #[case(StrictComparison::Strict, StrictComparison::Strict)] + fn multi_chunk_against_primitive_baseline( + #[case] lower_strict: StrictComparison, + #[case] upper_strict: StrictComparison, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: BufferMut = (0..3000u32).map(|i| i % 257).collect(); + let prim = PrimitiveArray::new(values.freeze(), Validity::NonNullable); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), 9, &mut ctx)?; + + let lower = ConstantArray::new(40u32, prim.len()).into_array(); + let upper = ConstantArray::new(200u32, prim.len()).into_array(); + let options = opts(lower_strict, upper_strict); + + let expected = prim + .into_array() + .between(lower.clone(), upper.clone(), options.clone())? + .execute::(&mut ctx)?; + let actual = packed + .into_array() + .between(lower, upper, options)? + .execute::(&mut ctx)?; + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn signed_with_patches_against_primitive_baseline() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..1500) + .map(|i| if i % 73 == 0 { 100_000 + i } else { i % 100 }) + .collect(); + let prim = PrimitiveArray::from_iter(values); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), 7, &mut ctx)?; + assert!(packed.patches().is_some(), "test setup expects patches"); + + let lower = ConstantArray::new(20i32, prim.len()).into_array(); + let upper = ConstantArray::new(80i32, prim.len()).into_array(); + let options = opts(StrictComparison::NonStrict, StrictComparison::NonStrict); + + let expected = prim + .into_array() + .between(lower.clone(), upper.clone(), options.clone())? + .execute::(&mut ctx)?; + let actual = packed + .into_array() + .between(lower, upper, options)? + .execute::(&mut ctx)?; + + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + #[test] + fn nullable_propagates_validity() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let prim = + PrimitiveArray::from_option_iter([Some(1u32), None, Some(3), Some(4), None, Some(6)]); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), 3, &mut ctx)?; + + let lower = ConstantArray::new(2u32, packed.len()).into_array(); + let upper = ConstantArray::new(5u32, packed.len()).into_array(); + let options = opts(StrictComparison::NonStrict, StrictComparison::NonStrict); + + let actual = packed + .into_array() + .between(lower.clone(), upper.clone(), options.clone())? + .execute::(&mut ctx)?; + let expected = prim + .into_array() + .between(lower, upper, options)? + .execute::(&mut ctx)?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/cast.rs b/encodings/fastlanes/src/bitpacking_v2/compute/cast.rs new file mode 100644 index 00000000000..bc5211a0983 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/cast.rs @@ -0,0 +1,271 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::builders::PrimitiveBuilder; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::cast::CastKernel; +use vortex_array::scalar_fn::fns::cast::CastReduce; +use vortex_array::validity::Validity; +use vortex_error::VortexResult; + +use crate::bitpacking_v2::BitPackedV2; +use crate::bitpacking_v2::array::BitPackedV2ArrayExt; +use crate::bitpacking_v2::array::bitpack_decompress::unpack_map_into_builder; + +/// Returns `true` if casting `src` to `tgt` is a widening integer cast for which every value a +/// bit-packed array can hold is guaranteed to be representable in `tgt` (so no per-value bounds +/// check is needed). This holds when `tgt` is strictly wider and either the source is unsigned +/// (always non-negative, fits in any wider type) or the target is also signed (sign-extension). +fn is_widening_int_cast(src: PType, tgt: PType) -> bool { + src.is_int() + && tgt.is_int() + && tgt.byte_width() > src.byte_width() + && (src.is_unsigned_int() || tgt.is_signed_int()) +} + +fn build_with_validity( + array: ArrayView<'_, BitPackedV2>, + dtype: &DType, + new_validity: Validity, +) -> VortexResult { + Ok(BitPackedV2::try_new( + array.packed().clone(), + dtype.as_ptype(), + new_validity, + array + .patches() + .map(|patches| patches.map_values(|values| values.cast(dtype.clone()))) + .transpose()?, + array.chunk_widths().clone(), + array.len(), + array.offset(), + )? + .into_array()) +} + +impl CastReduce for BitPackedV2 { + fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { + if !array.dtype().eq_ignore_nullability(dtype) { + return Ok(None); + } + let Some(new_validity) = array + .validity()? + .trivially_cast_nullability(dtype.nullability(), array.len())? + else { + return Ok(None); + }; + build_with_validity(array, dtype, new_validity).map(Some) + } +} + +impl CastKernel for BitPackedV2 { + fn cast( + array: ArrayView<'_, Self>, + dtype: &DType, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Nullability-only change: keep the values bit-packed, just adjust validity. + if array.dtype().eq_ignore_nullability(dtype) { + let new_validity = + array + .validity()? + .cast_nullability(dtype.nullability(), array.len(), ctx)?; + return build_with_validity(array, dtype, new_validity).map(Some); + } + + // Widening integer cast: unpack each FastLanes chunk into a cache-resident scratch buffer + // and cast-copy straight into the wide output, avoiding a full-length intermediate buffer + // and the generic cast kernel's bounds-check scan (unnecessary when widening). + let DType::Primitive(tgt, tgt_nullability) = dtype else { + return Ok(None); + }; + let (tgt, tgt_nullability) = (*tgt, *tgt_nullability); + let src = array.dtype().as_ptype(); + if !is_widening_int_cast(src, tgt) { + return Ok(None); + } + + // Surface the standard error if a nullable source with nulls is cast to a non-nullable + // type; on success the per-value validity is handled inside the unpack below. + array + .validity()? + .cast_nullability(tgt_nullability, array.len(), ctx)?; + + let result = match_each_integer_ptype!(tgt, |T| { + let mut builder = PrimitiveBuilder::::with_capacity(tgt_nullability, array.len()); + match_each_integer_ptype!(src, |F| { + unpack_map_into_builder::(array, &mut builder, ctx, |v: F| v.as_())?; + }); + builder.finish_into_primitive().into_array() + }); + Ok(Some(result)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builtins::ArrayBuiltins; + use vortex_array::compute::conformance::cast::test_cast_conformance; + use vortex_array::dtype::DType; + use vortex_array::dtype::NativePType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::match_each_integer_ptype; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::BitPackedV2Array; + use crate::BitPackedV2Data; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedV2Array { + BitPackedV2Data::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() + } + + #[test] + fn test_cast_bitpacked_u8_to_u32() { + let packed = bp(&buffer![10u8, 20, 30, 40, 50, 60].into_array(), 6); + + let casted = packed + .into_array() + .cast(DType::Primitive(PType::U32, Nullability::NonNullable)) + .unwrap(); + assert_eq!( + casted.dtype(), + &DType::Primitive(PType::U32, Nullability::NonNullable) + ); + + assert_arrays_eq!( + casted, + PrimitiveArray::from_iter([10u32, 20, 30, 40, 50, 60]), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn test_cast_bitpacked_nullable() { + let values = PrimitiveArray::from_option_iter([Some(5u16), None, Some(10), Some(15), None]); + let packed = bp(&values.into_array(), 4); + + let casted = packed + .into_array() + .cast(DType::Primitive(PType::U32, Nullability::Nullable)) + .unwrap(); + assert_eq!( + casted.dtype(), + &DType::Primitive(PType::U32, Nullability::Nullable) + ); + } + + /// End-to-end check that the real engine path `array.cast(target).execute()` routes through the + /// bit-packed widening pushdown and matches a plain primitive cast over the same values, across + /// every supported integer pair, several chunk-boundary lengths, and a sliced (offset > 0) case. + #[test] + fn test_cast_bitpacked_widening_via_execute() -> VortexResult<()> { + fn values(len: usize) -> PrimitiveArray { + PrimitiveArray::from_iter((0..len).map(|i| { + let value = if i % 17 == 0 { 31 } else { i % 8 }; + ::from_usize(value) + .expect("test values fit every integer ptype") + })) + } + + fn supported(src: PType, tgt: PType) -> bool { + src.is_int() + && tgt.is_int() + && tgt.byte_width() > src.byte_width() + && (src.is_unsigned_int() || tgt.is_signed_int()) + } + + let ptypes = [ + PType::I8, + PType::I16, + PType::I32, + PType::I64, + PType::U8, + PType::U16, + PType::U32, + PType::U64, + ]; + // Lengths exercise empty, sub-chunk, exact chunk, chunk+1, and multi-chunk-with-trailer. + let lengths = [0, 1, 7, 1023, 1024, 1025, 2051]; + + for src in ptypes { + for tgt in ptypes { + if !supported(src, tgt) { + continue; + } + + for len in lengths { + let source = match_each_integer_ptype!(src, |S| { values::(len) }); + let source_ref = source.into_array(); + let target = DType::Primitive(tgt, Nullability::NonNullable); + let mut ctx = SESSION.create_execution_ctx(); + + // Reference: plain primitive cast of the same values. + let reference = source_ref + .clone() + .cast(target.clone())? + .execute::(&mut ctx)?; + + // Candidate: bit-pack, then cast through the real engine. This dispatches to + // `BitPackedV2`'s `CastKernel` widening pushdown. + let packed = bp(&source_ref, 3).into_array(); + let casted = packed + .cast(target.clone())? + .execute::(&mut ctx)?; + assert_arrays_eq!(casted, reference, &mut ctx); + + // Also exercise the sliced/offset path (offset > 0, trailer present). + if len >= 4 { + let lo = len / 4; + let hi = len - len / 4; + let sliced = bp(&source_ref, 3).into_array().slice(lo..hi)?; + let casted = sliced + .cast(target.clone())? + .execute::(&mut ctx)?; + let reference = source_ref + .clone() + .slice(lo..hi)? + .cast(target.clone())? + .execute::(&mut ctx)?; + assert_arrays_eq!(casted, reference, &mut ctx); + } + } + } + } + + Ok(()) + } + + #[rstest] + #[case(bp(&buffer![0u8, 10, 20, 30, 40, 50, 60, 63].into_array(), 6))] + #[case(bp(&buffer![0u16, 100, 200, 300, 400, 500].into_array(), 9))] + #[case(bp(&buffer![0u32, 1000, 2000, 3000, 4000].into_array(), 12))] + #[case(bp(&PrimitiveArray::from_option_iter([Some(1u32), None, Some(7), Some(15), None]).into_array(), 4))] + fn test_cast_bitpacked_conformance(#[case] array: BitPackedV2Array) { + test_cast_conformance(&array.into_array(), &mut SESSION.create_execution_ctx()); + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/compare.rs b/encodings/fastlanes/src/bitpacking_v2/compute/compare.rs new file mode 100644 index 00000000000..8b616955254 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/compare.rs @@ -0,0 +1,325 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Block-streaming compare kernel for [`BitPackedV2Array`] against a constant. +//! +//! Avoids materialising the full primitive: the array is walked one 1024-element FastLanes +//! block at a time through a reusable scratch buffer, and a per-element bool is folded into +//! a [`BitBuffer`]. Patches are re-applied at the end by overwriting bits at the patched +//! indices with `predicate(patch_value)`. +//! +//! [`BitPackedV2Array`]: crate::BitPackedV2Array +//! [`BitBuffer`]: vortex_buffer::BitBuffer + +use fastlanes::BitPacking; +use fastlanes::BitPackingCompare; +use fastlanes::FastLanesComparable; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PhysicalPType; +use vortex_array::match_each_integer_ptype; +use vortex_array::scalar_fn::fns::binary::CompareKernel; +use vortex_array::scalar_fn::fns::operators::CompareOperator; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::bitpacking_v2::compute::compare_fused::stream_compare_fused; +use crate::bitpacking_v2::unpack_iter::BitPackedV2 as BitPackedV2Iter; + +impl CompareKernel for BitPackedV2 { + fn compare( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: CompareOperator, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Only accelerate compare-against-constant. + let Some(constant) = rhs.as_constant() else { + return Ok(None); + }; + let Some(constant_prim) = constant.as_primitive_opt() else { + return Ok(None); + }; + + // The adaptor strips a null-constant RHS, and binary expressions require both argument + // types to match before dispatch. + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + let lhs_ptype = lhs.dtype().as_ptype(); + if constant_prim.ptype() != lhs_ptype { + return Ok(None); + } + + let result = match_each_integer_ptype!(lhs_ptype, |T| { + let rhs: T = constant_prim + .typed_value::() + .vortex_expect("compare adaptor strips null constants"); + compare_constant_typed::(lhs, rhs, operator, nullability, ctx)? + }); + Ok(Some(result)) + } +} + +/// Compare every value against the constant via the fused FastLanes `unpack_cmp` kernel. +/// +/// `NativePType::is_eq` / `is_lt` etc. provide total comparison (matching the primitive between +/// kernel's dispatch shape). `NotEq` has no direct method, so use `!is_eq`. +fn compare_constant_typed( + lhs: ArrayView<'_, BitPackedV2>, + rhs: T, + operator: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType + + BitPackedV2Iter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, +{ + match operator { + CompareOperator::Eq => { + stream_compare_fused::(lhs, rhs, nullability, |a, b| a.is_eq(b), ctx) + } + CompareOperator::NotEq => { + stream_compare_fused::(lhs, rhs, nullability, |a, b| !a.is_eq(b), ctx) + } + CompareOperator::Lt => { + stream_compare_fused::(lhs, rhs, nullability, |a, b| a.is_lt(b), ctx) + } + CompareOperator::Lte => { + stream_compare_fused::(lhs, rhs, nullability, |a, b| a.is_le(b), ctx) + } + CompareOperator::Gt => { + stream_compare_fused::(lhs, rhs, nullability, |a, b| a.is_gt(b), ctx) + } + CompareOperator::Gte => { + stream_compare_fused::(lhs, rhs, nullability, |a, b| a.is_ge(b), ctx) + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::slice::SliceKernel; + use vortex_array::assert_arrays_eq; + use vortex_array::builtins::ArrayBuiltins; + use vortex_array::scalar_fn::fns::binary::CompareKernel; + use vortex_array::scalar_fn::fns::operators::CompareOperator; + use vortex_array::scalar_fn::fns::operators::Operator; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::BitPackedV2; + use crate::BitPackedV2ArrayExt; + use crate::BitPackedV2Data; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + /// All six operators on a small in-range input. + #[rstest] + #[case(Operator::Eq, vec![false, false, false, true, false, false, true])] + #[case(Operator::NotEq, vec![true, true, true, false, true, true, false])] + #[case(Operator::Lt, vec![true, true, true, false, false, false, false])] + #[case(Operator::Lte, vec![true, true, true, true, false, false, true])] + #[case(Operator::Gt, vec![false, false, false, false, true, true, false])] + #[case(Operator::Gte, vec![false, false, false, true, true, true, true])] + fn small(#[case] op: Operator, #[case] expected: Vec) { + let mut ctx = SESSION.create_execution_ctx(); + let values = PrimitiveArray::from_iter([0u32, 1, 2, 3, 4, 5, 3]); + let packed = BitPackedV2Data::encode(&values.into_array(), 3, &mut ctx).unwrap(); + let rhs = ConstantArray::new(3u32, packed.len()).into_array(); + let result = packed + .into_array() + .binary(rhs, op) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + } + + /// Sweep every native int type across several bit-widths. 2048 elements spans two + /// FastLanes blocks, exercising the per-type monomorphised inner loop. The kernel is + /// invoked *directly* and asserted `Some`, proving the streaming path engages (rather + /// than silently falling back to the Arrow compare), and its output is checked against + /// the Primitive fallback. + macro_rules! sweep { + ($name:ident, $T:ty, $($bw:expr),+) => { + #[test] + fn $name() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + for bw in [$($bw),+] { + let cap: u128 = 1u128 << bw; + let values: Vec<$T> = (0..2048u128).map(|i| (i % cap) as $T).collect(); + let prim = PrimitiveArray::from_iter(values); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), bw, &mut ctx)?; + let rhs_val = (cap.min(2048) / 2) as $T; + let rhs = ConstantArray::new(rhs_val, prim.len()).into_array(); + for op in [CompareOperator::Eq, CompareOperator::Lt, CompareOperator::Gte] { + let got = ::compare( + packed.as_view(), &rhs, op, &mut ctx, + )? + .expect("streaming compare kernel must engage") + .execute::(&mut ctx)?; + let want = prim + .clone() + .into_array() + .binary(rhs.clone(), Operator::from(op))? + .execute::(&mut ctx)?; + assert_arrays_eq!(got, want, &mut SESSION.create_execution_ctx()); + } + } + Ok(()) + } + }; + } + + sweep!(sweep_u8, u8, 1, 4, 7); + sweep!(sweep_u16, u16, 1, 8, 15); + sweep!(sweep_u32, u32, 1, 16, 31); + sweep!(sweep_u64, u64, 1, 32, 63); + sweep!(sweep_i8, i8, 1, 4, 7); + sweep!(sweep_i16, i16, 1, 8, 15); + sweep!(sweep_i32, i32, 1, 16, 31); + sweep!(sweep_i64, i64, 1, 32, 63); + + /// Inline-patch path: encode signed i32 values that exceed the bit-width range so they + /// end up in `Patches`. The streaming kernel must splice the patches in before the + /// predicate runs. + #[test] + fn signed_with_patches_matches_primitive() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..1500) + .map(|i| if i % 73 == 0 { 100_000 + i } else { i % 100 }) + .collect(); + let prim = PrimitiveArray::from_iter(values); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), 7, &mut ctx)?; + assert!(packed.patches().is_some(), "test setup expects patches"); + let rhs = ConstantArray::new(50i32, prim.len()).into_array(); + let expected = prim + .into_array() + .binary(rhs.clone(), Operator::Eq)? + .execute::(&mut ctx)?; + let actual = packed + .into_array() + .binary(rhs, Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } + + /// Sliced inputs: a non-zero block offset (and a length spanning several blocks) must still go + /// through the fused kernel and agree with the primitive fallback. Sweeps slice starts that + /// land both inside the first block and past it, with lengths that end mid-block and on a block + /// boundary. + #[rstest] + #[case(1, 4000)] // start mid-first-block, multi-block length + #[case(1023, 2)] // start at the last row of the first block + #[case(1024, 1024)] // start exactly on a block boundary, exactly one block long + #[case(1500, 1000)] // start mid-second-block + #[case(3, 1021)] // ends exactly on the first block boundary + fn sliced_matches_primitive( + #[case] start: usize, + #[case] slice_len: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..5000u32).map(|i| i % 128).collect(); + let prim = PrimitiveArray::from_iter(values); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), 7, &mut ctx)?; + + let sliced = packed.into_array().slice(start..start + slice_len)?; + let rhs = ConstantArray::new(50u32, slice_len).into_array(); + for op in [ + CompareOperator::Eq, + CompareOperator::Lt, + CompareOperator::Gte, + ] { + let got = ::compare( + sliced.as_::(), + &rhs, + op, + &mut ctx, + )? + .expect("fused compare kernel must engage for sliced arrays") + .execute::(&mut ctx)?; + let want = prim + .clone() + .into_array() + .slice(start..start + slice_len)? + .binary(rhs.clone(), Operator::from(op))? + .execute::(&mut ctx)?; + assert_arrays_eq!(got, want, &mut ctx); + } + Ok(()) + } + + /// Sliced *and* patched: combine a non-zero offset with out-of-range values that land in + /// `Patches`, exercising the `offset + (global - p_off)` patch-position math. + #[test] + fn sliced_with_patches_matches_primitive() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values: Vec = (0..4096) + .map(|i| if i % 91 == 0 { 100_000 + i } else { i % 100 }) + .collect(); + let prim = PrimitiveArray::from_iter(values); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), 7, &mut ctx)?; + assert!(packed.patches().is_some(), "test setup expects patches"); + + let (start, end) = (700usize, 3500usize); + // `ArrayRef::slice` leaves a lazy `SliceArray` over a patched `BitPackedV2` (the + // `SliceReduce` path bails when patches are present), so go through the `SliceKernel`, + // which reads the buffers and produces a sliced `BitPackedV2` with sliced patches. + let sliced = ::slice(packed.as_view(), start..end, &mut ctx)? + .expect("slice kernel produces a sliced bitpacked array"); + let rhs = ConstantArray::new(50i32, end - start).into_array(); + let got = ::compare( + sliced.as_::(), + &rhs, + CompareOperator::Eq, + &mut ctx, + )? + .expect("fused compare kernel must engage for sliced arrays with patches") + .execute::(&mut ctx)?; + let want = prim + .into_array() + .slice(start..end)? + .binary(rhs, Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!(got, want, &mut ctx); + Ok(()) + } + + /// Nullable input — the result must carry the array's validity. + #[test] + fn nullable_propagates_validity() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let prim = PrimitiveArray::from_option_iter([Some(1u32), None, Some(3), Some(4), None]); + let packed = BitPackedV2Data::encode(&prim.clone().into_array(), 3, &mut ctx)?; + let rhs = ConstantArray::new(3u32, packed.len()).into_array(); + let actual = packed + .into_array() + .binary(rhs.clone(), Operator::Eq)? + .execute::(&mut ctx)?; + let expected = prim + .into_array() + .binary(rhs, Operator::Eq)? + .execute::(&mut ctx)?; + assert_arrays_eq!(actual, expected, &mut ctx); + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/compare_fused.rs b/encodings/fastlanes/src/bitpacking_v2/compute/compare_fused.rs new file mode 100644 index 00000000000..58344d6bd9f --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/compare_fused.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fused compare kernel for [`BitPackedV2Array`] against a constant. +//! +//! Where [`super::stream_predicate`] unpacks a full 1024-element FastLanes block into a scratch +//! buffer and *then* folds a predicate over it, this path hands the comparison down into the +//! FastLanes [`BitPackingCompare::unchecked_unpack_cmp`] kernel, which compares each value against +//! the constant *as it is unpacked*, accumulating the boolean results straight into a 1024-bit +//! mask (`[u64; 16]`) in transposed FastLanes lane order - one register-resident word per lane, no +//! `[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::bitpacking_v2::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 +//! output buffer is over-allocated to whole 1024-bit blocks, so every block - the sliced first +//! block, the body, and the trailing partial - transposes straight into a 64-bit-word-aligned +//! slot with no per-block temporary and only one shared scratch `[u64; 16]`. The leading `offset` +//! garbage rows are represented as the final [`BitBuffer`] bit offset, which naturally handles +//! sub-byte slices without copy-aligning. Inline patches are spliced in afterwards by overwriting +//! the bits at the patched indices with `cmp(patch_value, rhs)`. +//! +//! [`BitPackedV2Array`]: crate::BitPackedV2Array +//! [`BitBuffer`]: vortex_buffer::BitBuffer + +use fastlanes::BitPacking; +use fastlanes::BitPackingCompare; +use fastlanes::FastLanesComparable; +use fastlanes::transpose_bits; +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PhysicalPType; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use super::stream_predicate::stream_predicate; +use crate::BitPackedV2; +use crate::BitPackedV2ArrayExt; +use crate::bitpacking_v2::unpack_iter::BitPackedV2 as BitPackedV2Iter; +use crate::bitpacking_v2::unpack_iter::for_each_packed_chunk; + +const CHUNK_SIZE: usize = 1024; +const U64_BITS: usize = u64::BITS as usize; +/// `u64` words spanning one FastLanes block (1024 bits / 64). +const WORDS_PER_CHUNK: usize = CHUNK_SIZE / U64_BITS; + +/// Compare the unpacked values of a [`BitPackedV2Array`] against `rhs` using the fused FastLanes +/// `unpack_cmp` kernel, producing a [`BoolArray`]. +/// +/// `cmp(value, rhs)` defines the predicate; it must be the total-order comparison matching the +/// requested operator (e.g. `|a, b| a.is_lt(b)`). +/// +/// [`BitPackedV2Array`]: crate::BitPackedV2Array +pub(super) fn stream_compare_fused( + array: ArrayView<'_, BitPackedV2>, + rhs: T, + nullability: Nullability, + cmp: F, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: NativePType + + BitPackedV2Iter + + FastLanesComparable::Physical>, + ::Physical: BitPacking + NativePType + BitPackingCompare, + F: Fn(T, T) -> bool + Copy, +{ + let len = array.len(); + 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 || widths.max_width() == 0 { + return stream_predicate::(array, nullability, move |v| cmp(v, rhs), ctx); + } + + // Over-allocate to whole 1024-bit blocks in padded coordinates so every block - including the + // trailing partial - has room for a full untranspose at a 64-bit-word-aligned offset. + let num_chunks = (offset + len).div_ceil(CHUNK_SIZE); + let mut words: BufferMut = BufferMut::zeroed(num_chunks * WORDS_PER_CHUNK); + + { + let words = words.as_mut_slice(); + let mut lane_major = [0u64; WORDS_PER_CHUNK]; + for_each_packed_chunk::( + array.packed_slice::<::Physical>(), + widths, + offset, + len, + |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. + unsafe { + <::Physical as BitPackingCompare>::unchecked_unpack_cmp::< + T, + _, + >(bit_width, packed_chunk, &mut lane_major, cmp, rhs); + } + transpose_bits::<::Physical>(&lane_major, out); + }, + )?; + } + + let mut bits = BitBufferMut::from_buffer(words.into_byte_buffer(), offset, len); + + // Patched indices hold placeholder packed values, so their fused result is meaningless; + // overwrite each with the comparison against the real patch value. + // TODO(joe): apply patches per `packed_chunked`. + if let Some(p) = array.patches() { + let p_idx = p.indices().clone().execute::(ctx)?; + // TODO(joe): push down cmp?? + let p_val = p.values().clone().execute::(ctx)?; + let p_off = p.offset(); + match_each_unsigned_integer_ptype!(p_idx.ptype(), |I| { + let indices = p_idx.as_slice::(); + let values = p_val.as_slice::(); + for (&global, &value) in indices.iter().zip(values) { + let global: usize = global.as_(); + let idx = global - p_off; + bits.set_to(idx, cmp(value, rhs)) + } + }); + } + + let validity = array.validity()?.union_nullability(nullability); + Ok(BoolArray::new(bits.freeze(), validity).into_array()) +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/filter.rs b/encodings/fastlanes/src/bitpacking_v2/compute/filter.rs new file mode 100644 index 00000000000..d74cb21a34f --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/filter.rs @@ -0,0 +1,351 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::MaybeUninit; + +use fastlanes::BitPacking; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::filter::FilterKernel; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::UnsignedPType; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_mask::MaskValuesRef; + +use super::chunked_indices; +use super::take::UNPACK_CHUNK_THRESHOLD; +use crate::BitPackedV2; +use crate::BitPackedV2ArrayExt; +use crate::BitPackedV2Data; + +/// The threshold over which it is faster to fully unpack the entire [`BitPackedV2Array`](crate::BitPackedV2Array) and then +/// filter the result than to unpack only specific bitpacked values into the output buffer. +pub const fn unpack_then_filter_threshold(ptype: PType) -> f64 { + // TODO(connor): Where did these numbers come from? Add a public link after validating them. + // These numbers probably don't work for in-place filtering either. + match ptype.byte_width() { + 1 => 0.03, + 2 => 0.03, + 4 => 0.075, + _ => 0.09, + // >8 bytes may have a higher threshold. These numbers are derived from a GCP c2-standard-4 + // with a "Cascade Lake" CPU. + } +} + +/// Kernel to execute filtering directly on a bit-packed array. +impl FilterKernel for BitPackedV2 { + fn filter( + array: ArrayView<'_, Self>, + mask: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let values = match mask { + Mask::AllTrue(_) | Mask::AllFalse(_) => { + return Ok(None); + } + Mask::Values(values) => values, + }; + + // If the density is high enough, then we would rather decompress the whole array and then apply + // a filter over decompressing values one by one. + if values.density() > unpack_then_filter_threshold(array.dtype().as_ptype()) { + return Ok(None); + } + + // Filter and patch using the correct unsigned type for FastLanes, then cast to signed if needed. + let primitive = + match_each_unsigned_integer_ptype!(array.dtype().as_ptype().to_unsigned(), |U| { + let (buffer, validity) = filter_primitive_without_patches::(array, values)?; + // reinterpret_cast for signed types. + let primitive = PrimitiveArray::new(buffer, validity); + if array.dtype().as_ptype().is_signed_int() { + PrimitiveArray::from_buffer_handle( + primitive.buffer_handle().clone(), + array.dtype().as_ptype(), + primitive.validity()?, + ) + } else { + primitive + } + }); + + let patches = array + .patches() + .map(|patches| patches.filter(&Mask::Values(MaskValuesRef::clone(values)), ctx)) + .transpose()? + .flatten(); + + if let Some(patches) = patches { + let mut prim_array = primitive; + prim_array = prim_array.patch(&patches, ctx)?; + return Ok(Some(prim_array.into_array())); + } + + Ok(Some(primitive.into_array())) + } +} + +/// Specialized filter kernel for primitive bit-packed arrays. +/// +/// Because the FastLanes bit-packing kernels are only implemented for unsigned types, the provided +/// `U` should be promoted to the unsigned variant for any target bit width. +/// For example, if the array is bit-packed `i16`, this function should be called with `U = u16`. +/// +/// This function fully decompresses the array for all but the most selective masks because the +/// FastLanes decompression is so fast and the bookkeepping necessary to decompress individual +/// elements is relatively slow. +/// +/// Returns a tuple of (values buffer, validity mask). +fn filter_primitive_without_patches( + array: ArrayView<'_, BitPackedV2>, + selection: &MaskValuesRef, +) -> VortexResult<(Buffer, Validity)> { + let values = filter_with_indices(array.data(), selection.indices()); + let validity = array + .validity()? + .filter(&Mask::Values(MaskValuesRef::clone(selection)))?; + + Ok((values.freeze(), validity)) +} + +fn filter_with_indices( + array: &BitPackedV2Data, + indices: &[usize], +) -> BufferMut { + let offset = array.offset() 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]; + + // Group the indices by the FastLanes chunk they belong to. + chunked_indices( + indices.iter().copied(), + offset, + |chunk_idx, indices_within_chunk| { + let (packed, bit_width) = array.packed_chunk::(chunk_idx); + + if indices_within_chunk.len() == 1024 { + // Unpack the entire chunk. + unsafe { + let values_len = values.len(); + values.set_len(values_len + 1024); + BitPacking::unchecked_unpack( + bit_width, + packed, + &mut values.as_mut_slice()[values_len..], + ); + } + } else if indices_within_chunk.len() > UNPACK_CHUNK_THRESHOLD { + // Unpack into a temporary chunk and then copy the values. + unsafe { + let dst: &mut [MaybeUninit] = &mut unpacked; + let dst: &mut [T] = std::mem::transmute(dst); + BitPacking::unchecked_unpack(bit_width, packed, dst); + } + values.extend_trusted( + indices_within_chunk + .iter() + .map(|&idx| unsafe { unpacked.get_unchecked(idx).assume_init() }), + ); + } else { + // Otherwise, unpack each element individually. + values.extend_trusted(indices_within_chunk.iter().map(|&idx| unsafe { + BitPacking::unchecked_unpack_single(bit_width, packed, idx) + })); + } + }, + ); + + values +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use vortex_array::IntoArray as _; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::compute::conformance::filter::test_filter_conformance; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_mask::Mask; + use vortex_session::VortexSession; + + use crate::BitPackedV2Data; + use crate::bitpacking_v2::array::BitPackedV2ArrayExt; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[test] + fn take_indices() { + let mut ctx = SESSION.create_execution_ctx(); + // Create a u8 array modulo 63. + let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); + + let mask = Mask::from_indices(bitpacked.len(), vec![0, 125, 2047, 2049, 2151, 2790]); + + let primitive_result = bitpacked.filter(mask).unwrap(); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]), + &mut ctx + ); + } + + #[test] + fn take_sliced_indices() { + let mut ctx = SESSION.create_execution_ctx(); + // Create a u8 array modulo 63. + let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); + let sliced = bitpacked.slice(128..2050).unwrap(); + + let mask = Mask::from_indices(sliced.len(), vec![1919, 1921]); + + let primitive_result = sliced.filter(mask).unwrap(); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([31u8, 33]), + &mut ctx + ); + } + + #[test] + fn filter_bitpacked() { + let mut ctx = SESSION.create_execution_ctx(); + let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); + let filtered = bitpacked.filter(Mask::from_indices(4096, 0..1024)).unwrap(); + let filtered_prim = filtered.execute::(&mut ctx).unwrap(); + assert_arrays_eq!( + filtered_prim, + PrimitiveArray::from_iter((0..1024).map(|i| (i % 63) as u8)), + &mut ctx + ); + } + + #[test] + fn filter_bitpacked_signed() { + let mut ctx = SESSION.create_execution_ctx(); + let values: Buffer = (0..500).collect(); + let unpacked = PrimitiveArray::new(values.clone(), Validity::NonNullable); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 9, &mut ctx).unwrap(); + let filtered = bitpacked + .filter(Mask::from_indices(values.len(), 0..250)) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + + assert_arrays_eq!( + filtered, + PrimitiveArray::from_iter(values[0..250].iter().copied()), + &mut ctx + ); + } + + #[test] + fn test_filter_bitpacked_conformance() { + let mut ctx = SESSION.create_execution_ctx(); + // Test with u8 values + let unpacked = buffer![1u8, 2, 3, 4, 5].into_array(); + let bitpacked = BitPackedV2Data::encode(&unpacked, 3, &mut ctx).unwrap(); + test_filter_conformance(&bitpacked.into_array(), &mut ctx); + + // Test with u32 values + let unpacked = buffer![100u32, 200, 300, 400, 500].into_array(); + let bitpacked = BitPackedV2Data::encode(&unpacked, 9, &mut ctx).unwrap(); + test_filter_conformance(&bitpacked.into_array(), &mut ctx); + + // Test with nullable values + let unpacked = PrimitiveArray::from_option_iter([Some(1u16), None, Some(3), Some(4), None]); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 3, &mut ctx).unwrap(); + test_filter_conformance(&bitpacked.into_array(), &mut ctx); + } + + /// Regression test for signed integers with patches. + /// + /// When filtering signed integers that have patches (exceptions), the patches + /// are stored with the signed type but FastLanes uses unsigned types internally. + /// This test ensures that the type handling is correct. + #[test] + fn filter_bitpacked_signed_with_patches() { + let mut ctx = SESSION.create_execution_ctx(); + // Create signed integer values where some exceed the bit width (causing patches). + // Values 0-127 fit in 7 bits, but 1000 and 2000 do not. + let values: Vec = vec![0, 10, 1000, 20, 30, 2000, 40, 50, 60, 70]; + let unpacked = PrimitiveArray::from_iter(values.clone()); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 7, &mut ctx).unwrap(); + assert!( + bitpacked.patches().is_some(), + "Expected patches for values exceeding bit width" + ); + + // Filter to include some patched and some non-patched values. + let filtered = bitpacked + .filter(Mask::from_indices(values.len(), vec![0, 2, 5, 9])) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + + assert_arrays_eq!( + filtered, + PrimitiveArray::from_iter([0i32, 1000, 2000, 70]), + &mut ctx + ); + } + + /// Regression test for signed integers with patches using low selectivity. + /// + /// This test uses a low selectivity filter which takes a different code path + /// that doesn't fully decompress the array first. + #[test] + fn filter_bitpacked_signed_with_patches_low_selectivity() { + let mut ctx = SESSION.create_execution_ctx(); + // Create a larger array with signed integers and some patches. + let values: Vec = (0..1000) + .map(|i| { + if i % 100 == 0 { + 10000 + i // These will be patches (exceed 7 bits) + } else { + i % 128 // These fit in 7 bits + } + }) + .collect(); + let unpacked = PrimitiveArray::from_iter(values.clone()); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 7, &mut ctx).unwrap(); + assert!( + bitpacked.patches().is_some(), + "Expected patches for values exceeding bit width" + ); + + // Use low selectivity (only select 2% of values) to avoid full decompression. + let indices: Vec = (0..20).collect(); + let filtered = bitpacked + .filter(Mask::from_indices(values.len(), indices)) + .unwrap() + .execute::(&mut ctx) + .unwrap(); + + let expected: Vec = values[0..20].to_vec(); + assert_arrays_eq!(filtered, PrimitiveArray::from_iter(expected), &mut ctx); + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/is_constant.rs b/encodings/fastlanes/src/bitpacking_v2/compute/is_constant.rs new file mode 100644 index 00000000000..e0d734491c6 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/is_constant.rs @@ -0,0 +1,207 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem::MaybeUninit; +use std::ops::Range; + +use itertools::Itertools; +use lending_iterator::LendingIterator; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::fns::is_constant::IsConstant; +use vortex_array::aggregate_fn::fns::is_constant::primitive::IS_CONST_LANE_WIDTH; +use vortex_array::aggregate_fn::fns::is_constant::primitive::compute_is_constant; +use vortex_array::aggregate_fn::kernels::DynAggregateKernel; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::IntegerPType; +use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::BitPackedV2ArrayExt; +use crate::bitpacking_v2::unpack_iter::BitPackedV2 as BitPackedV2Unpack; + +/// BitPackedV2-specific is_constant kernel with SIMD support. +#[derive(Debug)] +pub(crate) struct BitPackedV2IsConstantKernel; + +impl DynAggregateKernel for BitPackedV2IsConstantKernel { + fn aggregate( + &self, + aggregate_fn: &AggregateFnRef, + batch: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !aggregate_fn.is::() { + return Ok(None); + } + + let Some(array) = batch.as_opt::() else { + return Ok(None); + }; + + let result = match_each_integer_ptype!(array.dtype().as_ptype(), |P| { + bitpacked_is_constant::() }>(array, ctx)? + }); + + Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) + } +} + +fn bitpacked_is_constant( + array: ArrayView<'_, BitPackedV2>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut scratch = [const { MaybeUninit::::uninit() }; 1024]; + let mut bit_unpack_iterator = array.unpacked_chunks::(&mut scratch)?; + let patches = array + .patches() + .map(|p| -> VortexResult<_> { + let values = p.values().clone().execute::(ctx)?; + let indices = p.indices().clone().execute::(ctx)?; + let offset = p.offset(); + Ok((indices, values, offset)) + }) + .transpose()?; + + let mut header_constant_value = None; + let mut current_idx = 0; + if let Some(header) = bit_unpack_iterator.initial() { + if let Some((indices, patches, offset)) = &patches { + apply_patches( + header, + current_idx..header.len(), + indices, + patches.as_slice::(), + *offset, + ) + } + + if !compute_is_constant::<_, WIDTH>(header) { + return Ok(false); + } + header_constant_value = Some(header[0]); + current_idx = header.len(); + } + + let mut first_chunk_value = None; + { + let mut chunks_iter = bit_unpack_iterator.full_chunks(); + while let Some(chunk) = chunks_iter.next() { + if let Some((indices, patches, offset)) = &patches { + let chunk_len = chunk.len(); + apply_patches( + chunk, + current_idx..current_idx + chunk_len, + indices, + patches.as_slice::(), + *offset, + ) + } + + if !compute_is_constant::<_, WIDTH>(chunk) { + return Ok(false); + } + + if let Some(chunk_value) = first_chunk_value { + if chunk_value != chunk[0] { + return Ok(false); + } + } else { + if let Some(header_value) = header_constant_value + && header_value != chunk[0] + { + return Ok(false); + } + first_chunk_value = Some(chunk[0]); + } + + current_idx += chunk.len(); + } + } + + if let Some(trailer) = bit_unpack_iterator.trailer() { + if let Some((indices, patches, offset)) = &patches { + let chunk_len = trailer.len(); + apply_patches( + trailer, + current_idx..current_idx + chunk_len, + indices, + patches.as_slice::(), + *offset, + ) + } + + if !compute_is_constant::<_, WIDTH>(trailer) { + return Ok(false); + } + + if let Some(previous_const_value) = header_constant_value.or(first_chunk_value) + && previous_const_value != trailer[0] + { + return Ok(false); + } + } + + Ok(true) +} + +fn apply_patches( + values: &mut [T], + values_range: Range, + patch_indices: &PrimitiveArray, + patch_values: &[T], + indices_offset: usize, +) { + match_each_unsigned_integer_ptype!(patch_indices.ptype(), |I| { + apply_patches_idx_typed( + values, + values_range, + patch_indices.as_slice::(), + patch_values, + indices_offset, + ) + }); +} + +fn apply_patches_idx_typed( + values: &mut [T], + values_range: Range, + patch_indices: &[I], + patch_values: &[T], + indices_offset: usize, +) { + for (i, &v) in patch_indices + .iter() + .map(|i| i.as_() - indices_offset) + .zip_eq(patch_values) + .skip_while(|(i, _)| i < &values_range.start) + .take_while(|(i, _)| i < &values_range.end) + { + values[i - values_range.start] = v + } +} + +#[cfg(test)] +mod tests { + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::fns::is_constant::is_constant; + use vortex_array::array_session; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::BitPackedV2Data; + + #[test] + fn is_constant_with_patches() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let array = BitPackedV2Data::encode(&buffer![4; 1025].into_array(), 2, &mut ctx)?; + assert!(is_constant(&array.into_array(), &mut ctx)?); + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/mod.rs b/encodings/fastlanes/src/bitpacking_v2/compute/mod.rs new file mode 100644 index 00000000000..b5b6f9e412a --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/mod.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod between; +mod cast; +mod compare; +mod compare_fused; +mod filter; +pub(crate) mod is_constant; +mod slice; +mod stream_predicate; +mod take; + +// TODO(connor): This is duplicated in `encodings/fastlanes/src/bitpacking/kernels/mod.rs`. +fn chunked_indices( + mut indices: impl Iterator, + offset: usize, + mut chunk_fn: F, +) { + let mut indices_within_chunk: Vec = Vec::with_capacity(1024); + + let Some(first_idx) = indices.next() else { + return; + }; + + let mut current_chunk_idx = (first_idx + offset) / 1024; + indices_within_chunk.push((first_idx + offset) % 1024); + for idx in indices { + let new_chunk_idx = (idx + offset) / 1024; + + if new_chunk_idx != current_chunk_idx { + chunk_fn(current_chunk_idx, &indices_within_chunk); + indices_within_chunk.clear(); + } + + current_chunk_idx = new_chunk_idx; + indices_within_chunk.push((idx + offset) % 1024); + } + + if !indices_within_chunk.is_empty() { + chunk_fn(current_chunk_idx, &indices_within_chunk); + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; + use vortex_array::compute::conformance::consistency::test_array_consistency; + + use crate::BitPackedV2Array; + use crate::bitpacking_v2::bitpack_compress::bitpack_encode; + use crate::bitpacking_v2::compute::chunked_indices; + + fn bp(array: &PrimitiveArray, bit_width: u8) -> BitPackedV2Array { + bitpack_encode( + array, + bit_width, + &mut array_session().create_execution_ctx(), + ) + .unwrap() + } + + #[test] + fn chunk_indices_repeated() { + let mut called = false; + chunked_indices([0; 1025].into_iter(), 0, |chunk_idx, idxs| { + assert_eq!(chunk_idx, 0); + assert_eq!(idxs, [0; 1025]); + called = true; + }); + assert!(called); + } + + #[rstest] + // Basic integer arrays that can be bitpacked + #[case::u8_small(bp(&PrimitiveArray::from_iter([1u8, 2, 3, 4, 5]), 3))] + #[case::u16_array(bp(&PrimitiveArray::from_iter([10u16, 20, 30, 40, 50]), 6))] + #[case::u32_array(bp(&PrimitiveArray::from_iter([100u32, 200, 300, 400, 500]), 9))] + // Arrays with nulls + #[case::nullable_u8(bp(&PrimitiveArray::from_option_iter([Some(1u8), None, Some(3), Some(4), None]), 3))] + #[case::nullable_u32(bp(&PrimitiveArray::from_option_iter([Some(100u32), None, Some(300), Some(400), None]), 9))] + // Edge cases + #[case::single_element(bp(&PrimitiveArray::from_iter([42u32]), 6))] + #[case::all_zeros(bp(&PrimitiveArray::from_iter([0u16; 100]), 1))] + // Large arrays (multiple chunks - fastlanes uses 1024-element chunks) + #[case::large_u16(bp(&PrimitiveArray::from_iter((0..2048).map(|i| (i % 256) as u16)), 8))] + #[case::large_u32(bp(&PrimitiveArray::from_iter((0..3000).map(|i| (i % 1024) as u32)), 10))] + #[case::large_u8_many_chunks(bp(&PrimitiveArray::from_iter((0..5120).map(|i| (i % 128) as u8)), 7))] // 5 chunks + #[case::large_nullable(bp(&PrimitiveArray::from_option_iter((0..2500).map(|i| if i % 10 == 0 { None } else { Some((i % 512) as u16) })), 9))] + // Arrays with specific bit patterns + #[case::max_value_for_bits(bp(&PrimitiveArray::from_iter([7u8, 7, 7, 7, 7]), 3))] // max value for 3 bits + #[case::alternating_bits(bp(&PrimitiveArray::from_iter([0u16, 255, 0, 255, 0, 255]), 8))] + + fn test_bitpacked_consistency(#[case] array: BitPackedV2Array) { + let ctx = &mut array_session().create_execution_ctx(); + test_array_consistency(&array.into_array(), ctx); + } + + #[rstest] + #[case::u8_basic(bp(&PrimitiveArray::from_iter([1u8, 2, 3, 4, 5]), 3))] + #[case::u16_basic(bp(&PrimitiveArray::from_iter([10u16, 20, 30, 40, 50]), 6))] + #[case::u32_basic(bp(&PrimitiveArray::from_iter([100u32, 200, 300, 400, 500]), 9))] + #[case::u64_basic(bp(&PrimitiveArray::from_iter([1000u64, 2000, 3000, 4000, 5000]), 13))] + #[case::i32_basic(bp(&PrimitiveArray::from_iter([10i32, 20, 30, 40, 50]), 7))] + #[case::large_u32(bp(&PrimitiveArray::from_iter((0..100).map(|i| i as u32)), 7))] + fn test_bitpacked_binary_numeric(#[case] array: BitPackedV2Array) { + test_binary_numeric_array( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/slice.rs b/encodings/fastlanes/src/bitpacking_v2/compute/slice.rs new file mode 100644 index 00000000000..199fe54edbf --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/slice.rs @@ -0,0 +1,107 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::cmp::max; +use std::ops::Range; + +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::slice::SliceKernel; +use vortex_array::arrays::slice::SliceReduce; +use vortex_array::patches::Patches; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::bitpacking_v2::array::BitPackedV2ArrayExt; + +impl SliceReduce for BitPackedV2 { + fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { + // We cannot access buffers (to slice the patches). + if array.patches().is_some() { + return Ok(None); + } + + Ok(Some(slice_bitpacked(array, range, None)?)) + } +} + +impl SliceKernel for BitPackedV2 { + fn slice( + array: ArrayView<'_, Self>, + range: Range, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let patches = array + .patches() + .map(|p| p.slice(range.clone())) + .transpose()? + .flatten(); + + Ok(Some(slice_bitpacked(array, range, patches)?)) + } +} + +fn slice_bitpacked( + array: ArrayView<'_, BitPackedV2>, + range: Range, + patches: Option, +) -> VortexResult { + let offset_start = range.start + array.offset() as usize; + let offset_stop = range.end + array.offset() as usize; + let offset = offset_start % 1024; + let block_start = max(0, offset_start - offset); + let block_stop = offset_stop.div_ceil(1024) * 1024; + + let 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(BitPackedV2::try_new( + array.packed().slice(encoded_start..encoded_stop), + array.dtype().as_ptype(), + array.validity()?.slice(range.clone())?, + patches, + widths.slice(chunk_start..chunk_stop), + range.len(), + offset as u16, + )? + .into_array()) +} + +#[cfg(test)] +mod tests { + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::SliceArray; + use vortex_error::VortexResult; + + use crate::BitPackedV2; + use crate::bitpacking_v2::bitpack_compress::bitpack_encode; + + #[test] + fn test_reduce_parent_returns_bitpacked_slice() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter(0u32..2048); + let bitpacked = bitpack_encode(&values, 11, &mut ctx)?; + + let slice_array = SliceArray::new(bitpacked.clone().into_array(), 500..1500); + + let bitpacked_ref = bitpacked.into_array(); + let reduced = bitpacked_ref + .reduce_parent(&slice_array.into_array(), 0)? + .expect("expected slice kernel to execute"); + + assert!(reduced.is::()); + let reduced_bp = reduced.as_::(); + assert_eq!(reduced_bp.offset(), 500); + assert_eq!(reduced.len(), 1000); + + Ok(()) + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/stream_predicate.rs b/encodings/fastlanes/src/bitpacking_v2/compute/stream_predicate.rs new file mode 100644 index 00000000000..65eb1587e8b --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/stream_predicate.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Streaming, cache-reusable predicate evaluation over a [`BitPackedV2Array`]. +//! +//! Walks the encoded array one 1024-element FastLanes block at a time through a single +//! reusable scratch buffer, splices any [`Patches`] into the unpacked block +//! in place via a sorted-index cursor, then folds a `Fn(T) -> bool` predicate over the +//! block. The fold matches the canonical [`BitBuffer::collect_bool`] shape +//! (pack 64 bools into a `u64` in a tight auto-vectorisable inner loop) and writes the +//! resulting words straight into the output bit buffer, so the materialised primitive +//! never appears anywhere. +//! +//! [`BitPackedV2Array`]: crate::BitPackedV2Array +//! [`BitBuffer::collect_bool`]: vortex_buffer::BitBuffer::collect_bool +//! [`Patches`]: vortex_array::patches::Patches + +use std::mem::MaybeUninit; + +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; +use vortex_buffer::pack_bools_into_words; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::BitPackedV2ArrayExt; +use crate::FL_CHUNK_SIZE; +use crate::bitpacking_v2::unpack_iter::BitPackedV2 as BitPackedV2Iter; + +/// Stream `predicate` over the unpacked values of a [`BitPackedV2Array`](crate::BitPackedV2Array), one FastLanes +/// block at a time, producing a [`BoolArray`]. +pub(super) fn stream_predicate( + array: ArrayView<'_, BitPackedV2>, + nullability: Nullability, + predicate: P, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: BitPackedV2Iter + NativePType + Copy, + P: Fn(T) -> bool, +{ + let len = array.len(); + let mut words: BufferMut = BufferMut::zeroed(len.div_ceil(u64::BITS as usize)); + + if len > 0 { + let mut scratch = [const { MaybeUninit::::uninit() }; FL_CHUNK_SIZE]; + let mut chunks = array.unpacked_chunks::(&mut scratch)?; + let words = words.as_mut_slice(); + + if let Some(p) = array.patches() { + let p_idx_arr = p.indices().clone().execute::(ctx)?; + let p_val_arr = p.values().clone().execute::(ctx)?; + let p_off = p.offset(); + match_each_unsigned_integer_ptype!(p_idx_arr.ptype(), |I| { + let p_idx = p_idx_arr.as_slice::(); + let p_val = p_val_arr.as_slice::(); + let mut p_cur: usize = 0; + chunks.for_each_unpacked_chunk(|block, range| { + p_cur = splice_patches::(block, range.start, p_cur, p_idx, p_val, p_off); + pack_bools_into_words(words, range.start, block.len(), |i| predicate(block[i])); + }); + }); + } else { + chunks.for_each_unpacked_chunk(|block, range| { + pack_bools_into_words(words, range.start, block.len(), |i| predicate(block[i])); + }); + } + } + + let bits = BitBufferMut::from_buffer(words.into_byte_buffer(), 0, len); + let validity = array.validity()?.union_nullability(nullability); + Ok(BoolArray::new(bits.freeze(), validity).into_array()) +} + +/// Overwrite the unpacked block in place with any patches falling in +/// `[chunk_start, chunk_start + block.len())`, starting from `cursor` and returning the +/// advanced cursor. Sorted indices mean the cursor only moves forward across the walk. +#[inline] +fn splice_patches( + block: &mut [T], + chunk_start: usize, + mut cursor: usize, + indices: &[I], + values: &[T], + patch_offset: usize, +) -> usize +where + T: Copy, + I: AsPrimitive, +{ + let end = chunk_start + block.len(); + while cursor < indices.len() { + let global: usize = indices[cursor].as_(); + let local = global - patch_offset; + if local >= end { + break; + } + debug_assert!(local >= chunk_start); + block[local - chunk_start] = values[cursor]; + cursor += 1; + } + cursor +} diff --git a/encodings/fastlanes/src/bitpacking_v2/compute/take.rs b/encodings/fastlanes/src/bitpacking_v2/compute/take.rs new file mode 100644 index 00000000000..9392846a1a9 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/compute/take.rs @@ -0,0 +1,328 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::mem; +use std::mem::MaybeUninit; + +use fastlanes::BitPacking; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::dict::TakeExecute; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; +use vortex_array::match_each_unsigned_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect as _; +use vortex_error::VortexResult; + +use super::chunked_indices; +use crate::BitPackedV2; +use crate::BitPackedV2ArrayExt; +use crate::bitpacking_v2::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 +/// all 1024 elements takes ~8.8x as long as unpacking a single element on an M2 Macbook Air. +/// see +pub(super) const UNPACK_CHUNK_THRESHOLD: usize = 8; + +impl TakeExecute for BitPackedV2 { + fn take( + array: ArrayView<'_, Self>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // If the indices are large enough, it's faster to flatten and take the primitive array. + if indices.len() * UNPACK_CHUNK_THRESHOLD > array.len() { + let prim = array.array().clone().execute::(ctx)?; + return prim.into_array().take(indices.clone()).map(Some); + } + + // NOTE: we use the unsigned PType because all values in the BitPackedV2Array must + // be non-negative (pre-condition of creating the BitPackedV2Array). + let ptype: PType = PType::try_from(array.dtype())?; + let validity = array.validity()?; + let taken_validity = validity.take(indices)?; + + let indices = indices.clone().execute::(ctx)?; + let taken = match_each_unsigned_integer_ptype!(ptype.to_unsigned(), |T| { + match_each_integer_ptype!(indices.ptype(), |I| { + take_primitive::(array, &indices, taken_validity, ctx)? + }) + }); + let taken = if ptype.is_signed_int() { + PrimitiveArray::from_buffer_handle( + taken.buffer_handle().clone(), + ptype, + taken.validity()?, + ) + } else { + taken + }; + Ok(Some(taken.into_array())) + } +} + +fn take_primitive( + array: ArrayView<'_, BitPackedV2>, + indices: &PrimitiveArray, + taken_validity: Validity, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if indices.is_empty() { + return Ok(PrimitiveArray::new(Buffer::::empty(), taken_validity)); + } + + let offset = array.offset() as usize; + + // Group indices by 1024-element chunk, *without* allocating on the heap + let indices_iter = indices.as_slice::().iter().map(|i| { + i.to_usize() + .vortex_expect("index must be expressible as usize") + }); + + let mut output = BufferMut::::with_capacity(indices.len()); + let mut unpacked = [const { MaybeUninit::uninit() }; 1024]; + + chunked_indices(indices_iter, offset, |chunk_idx, indices_within_chunk| { + let (packed, bit_width) = array.packed_chunk::(chunk_idx); + + let mut have_unpacked = false; + let (offset_chunks, remainder) = indices_within_chunk.as_chunks::(); + + // this loop only runs if we have at least UNPACK_CHUNK_THRESHOLD offsets + for offset_chunk in offset_chunks { + if !have_unpacked { + unsafe { + let dst: &mut [MaybeUninit] = &mut unpacked; + let dst: &mut [T] = mem::transmute(dst); + BitPacking::unchecked_unpack(bit_width, packed, dst); + } + have_unpacked = true; + } + + for &index in offset_chunk { + output.push(unsafe { unpacked[index].assume_init() }); + } + } + + // if we have a remainder (i.e., < UNPACK_CHUNK_THRESHOLD leftover offsets), we need to handle it + if !remainder.is_empty() { + if have_unpacked { + // we already bulk unpacked this chunk, so we can just push the remaining elements + for &index in remainder { + output.push(unsafe { unpacked[index].assume_init() }); + } + } else { + // we had fewer than UNPACK_CHUNK_THRESHOLD offsets in the first place, + // so we need to unpack each one individually + for &index in remainder { + output.push(unsafe { + bitpack_decompress::unpack_single_primitive::(packed, bit_width, index) + }); + } + } + } + }); + + let unpatched_taken = if array.dtype().as_ptype().is_signed_int() { + let primitive = PrimitiveArray::new(output, taken_validity); + PrimitiveArray::from_buffer_handle( + primitive.buffer_handle().clone(), + array.dtype().as_ptype(), + primitive.validity()?, + ) + } else { + PrimitiveArray::new(output, taken_validity) + }; + if let Some(patches) = array.patches() + && let Some(patches) = patches.take(&indices.clone().into_array(), ctx)? + { + return unpatched_taken.patch(&patches, ctx); + } + + Ok(unpatched_taken) +} + +#[cfg(test)] +#[expect(clippy::cast_possible_truncation)] +mod test { + use std::sync::LazyLock; + + use rand::RngExt; + use rand::distr::Uniform; + use rand::rng; + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::compute::conformance::take::test_take_conformance; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_session::VortexSession; + + use crate::BitPackedV2Array; + use crate::BitPackedV2Data; + use crate::bitpacking_v2::array::BitPackedV2ArrayExt; + use crate::bitpacking_v2::compute::take::take_primitive; + + static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + crate::initialize(&session); + session + }); + + #[test] + fn take_indices() { + let mut ctx = SESSION.create_execution_ctx(); + let indices = buffer![0, 125, 2047, 2049, 2151, 2790].into_array(); + + // Create a u8 array modulo 63. + let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); + + let primitive_result = bitpacked.take(indices).unwrap(); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([0u8, 62, 31, 33, 9, 18]), + &mut ctx + ); + } + + #[test] + fn take_with_patches() { + let mut ctx = SESSION.create_execution_ctx(); + let unpacked = Buffer::from_iter(0u32..1024).into_array(); + let bitpacked = BitPackedV2Data::encode(&unpacked, 2, &mut ctx).unwrap(); + + let indices = buffer![0, 2, 4, 6].into_array(); + + let primitive_result = bitpacked.take(indices).unwrap(); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([0u32, 2, 4, 6]), + &mut ctx + ); + } + + #[test] + fn take_sliced_indices() { + let mut ctx = SESSION.create_execution_ctx(); + let indices = buffer![1919, 1921].into_array(); + + // Create a u8 array modulo 63. + let unpacked = PrimitiveArray::from_iter((0..4096).map(|i| (i % 63) as u8)); + let bitpacked = BitPackedV2Data::encode(&unpacked.into_array(), 6, &mut ctx).unwrap(); + let sliced = bitpacked.slice(128..2050).unwrap(); + + let primitive_result = sliced.take(indices).unwrap(); + assert_arrays_eq!( + primitive_result, + PrimitiveArray::from_iter([31u8, 33]), + &mut ctx + ); + } + + #[test] + #[cfg_attr(miri, ignore)] // This test is too slow on miri + fn take_random_indices() { + let mut ctx = SESSION.create_execution_ctx(); + let num_patches: usize = 128; + let values = (0..u16::MAX as u32 + num_patches as u32).collect::>(); + let uncompressed = PrimitiveArray::new(values.clone(), Validity::NonNullable); + let packed = BitPackedV2Data::encode(&uncompressed.into_array(), 16, &mut ctx).unwrap(); + assert!(packed.patches().is_some()); + + let rng = rng(); + let range = Uniform::new(0, values.len()).unwrap(); + let random_indices = + PrimitiveArray::from_iter(rng.sample_iter(range).take(10_000).map(|i| i as u32)); + let taken = packed.take(random_indices.clone().into_array()).unwrap(); + + // sanity check + random_indices + .as_slice::() + .iter() + .enumerate() + .for_each(|(ti, i)| { + assert_eq!( + u32::try_from(&packed.execute_scalar(*i as usize, &mut ctx).unwrap()).unwrap(), + values[*i as usize] + ); + assert_eq!( + u32::try_from(&taken.execute_scalar(ti, &mut ctx).unwrap()).unwrap(), + values[*i as usize] + ); + }); + } + + #[test] + #[cfg_attr(miri, ignore)] + fn take_signed_with_patches() { + let mut ctx = SESSION.create_execution_ctx(); + let start = + BitPackedV2Data::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx) + .unwrap(); + + let taken_primitive = take_primitive::( + start.as_view(), + &PrimitiveArray::from_iter([0u64, 1, 2, 3]), + Validity::NonNullable, + &mut ctx, + ) + .unwrap(); + assert_arrays_eq!( + taken_primitive, + PrimitiveArray::from_iter([1i32, 2, 3, 4]), + &mut ctx + ); + } + + #[test] + fn take_nullable_with_nullables() { + let mut ctx = SESSION.create_execution_ctx(); + let start = + BitPackedV2Data::encode(&buffer![1i32, 2i32, 3i32, 4i32].into_array(), 1, &mut ctx) + .unwrap(); + + let taken_primitive = start + .take( + PrimitiveArray::from_option_iter([Some(0u64), Some(1), None, Some(3)]).into_array(), + ) + .unwrap(); + assert_arrays_eq!( + taken_primitive, + PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, Some(4)]), + &mut ctx + ); + let taken_primitive_prim = taken_primitive.execute::(&mut ctx).unwrap(); + assert_eq!(taken_primitive_prim.invalid_count(&mut ctx).unwrap(), 1); + } + + fn bp(array: vortex_array::ArrayRef, bit_width: u8) -> BitPackedV2Array { + BitPackedV2Data::encode(&array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() + } + + #[rstest] + #[case(bp(PrimitiveArray::from_iter((0..100).map(|i| (i % 63) as u8)).into_array(), 6))] + #[case(bp(PrimitiveArray::from_iter((0..256).map(|i| i as u32)).into_array(), 8))] + #[case(bp(buffer![1i32, 2, 3, 4, 5, 6, 7, 8].into_array(), 3))] + #[case(bp( + PrimitiveArray::from_option_iter([Some(10u16), None, Some(20), Some(30), None]).into_array(), + 5 + ))] + #[case(bp(buffer![42u32].into_array(), 6))] + #[case(bp(PrimitiveArray::from_iter((0..1024).map(|i| i as u32)).into_array(), 8))] + fn test_take_bitpacked_conformance(#[case] bitpacked: BitPackedV2Array) { + test_take_conformance(&bitpacked.into_array(), &mut SESSION.create_execution_ctx()); + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/mod.rs b/encodings/fastlanes/src/bitpacking_v2/mod.rs new file mode 100644 index 00000000000..de0ad52f90b --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/mod.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +mod array; +pub use array::BitPackedV2ArrayExt; +pub use array::BitPackedV2ArraySlotsExt; +pub use array::BitPackedV2Data; +pub use array::BitPackedV2DataParts; +pub use array::BitPackedV2Slots; +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 vtable; + +pub use vtable::BitPackedV2; +pub use vtable::BitPackedV2Array; + +pub(crate) fn initialize(session: &vortex_session::VortexSession) { + vtable::initialize(session); +} diff --git a/encodings/fastlanes/src/bitpacking_v2/vtable/kernels.rs b/encodings/fastlanes/src/bitpacking_v2/vtable/kernels.rs new file mode 100644 index 00000000000..7a75fbed298 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/kernels.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayVTable; +use vortex_array::arrays::Dict; +use vortex_array::arrays::Filter; +use vortex_array::arrays::Slice; +use vortex_array::arrays::dict::TakeExecuteAdaptor; +use vortex_array::arrays::filter::FilterExecuteAdaptor; +use vortex_array::arrays::slice::SliceExecuteAdaptor; +use vortex_array::optimizer::kernels::ArrayKernelsExt; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::fns::between::Between; +use vortex_array::scalar_fn::fns::between::BetweenExecuteAdaptor; +use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::binary::CompareExecuteAdaptor; +use vortex_array::scalar_fn::fns::cast::Cast; +use vortex_array::scalar_fn::fns::cast::CastExecuteAdaptor; +use vortex_session::VortexSession; + +use crate::BitPackedV2; + +pub(crate) fn initialize(session: &VortexSession) { + let kernels = session.kernels(); + kernels.register_execute_parent_kernel( + Between.id(), + BitPackedV2, + BetweenExecuteAdaptor(BitPackedV2), + ); + kernels.register_execute_parent_kernel(Cast.id(), BitPackedV2, CastExecuteAdaptor(BitPackedV2)); + kernels.register_execute_parent_kernel( + Binary.id(), + BitPackedV2, + CompareExecuteAdaptor(BitPackedV2), + ); + kernels.register_execute_parent_kernel( + Filter.id(), + BitPackedV2, + FilterExecuteAdaptor(BitPackedV2), + ); + kernels.register_execute_parent_kernel( + Slice.id(), + BitPackedV2, + SliceExecuteAdaptor(BitPackedV2), + ); + kernels.register_execute_parent_kernel(Dict.id(), BitPackedV2, TakeExecuteAdaptor(BitPackedV2)); +} diff --git a/encodings/fastlanes/src/bitpacking_v2/vtable/mod.rs b/encodings/fastlanes/src/bitpacking_v2/vtable/mod.rs new file mode 100644 index 00000000000..6b3f2be44d1 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/mod.rs @@ -0,0 +1,385 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::hash::Hash; +use std::hash::Hasher; + +use prost::Message; +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::IntoArray; +use vortex_array::buffer::BufferHandle; +use vortex_array::builders::ArrayBuilder; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::match_each_integer_ptype; +use vortex_array::patches::Patches; +use vortex_array::patches::PatchesData; +use vortex_array::patches::PatchesMetadata; +use vortex_array::require_patches; +use vortex_array::require_validity; +use vortex_array::serde::ArrayChildren; +use vortex_array::validity::Validity; +use vortex_array::vtable::VTable; +use vortex_array::vtable::child_to_validity; +use vortex_array::vtable::validity_to_child; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::BitPackedV2ArrayExt; +use crate::BitPackedV2Data; +use crate::BitPackedV2DataParts; +use crate::ChunkWidths; +use crate::bitpacking_v2::array::BitPackedV2Slots; +use crate::bitpacking_v2::array::BitPackedV2SlotsView; +use crate::bitpacking_v2::array::PATCH_SLOTS; +use crate::bitpacking_v2::bitpack_decompress::unpack_array; +use crate::bitpacking_v2::bitpack_decompress::unpack_into_primitive_builder; +use crate::bitpacking_v2::vtable::rules::RULES; +mod kernels; +mod operations; +mod rules; +mod validity; + +/// A [`BitPackedV2`]-encoded Vortex array. +pub type BitPackedV2Array = Array; + +pub(crate) fn initialize(session: &VortexSession) { + kernels::initialize(session); +} + +#[derive(Clone, prost::Message)] +pub struct BitPackedV2Metadata { + #[prost(uint32, tag = "1")] + pub(crate) bit_width: u32, + #[prost(uint32, tag = "2")] + pub(crate) offset: u32, // must be <1024 + #[prost(message, optional, tag = "3")] + pub(crate) patches: Option, + /// One width per 1024-element chunk. Empty only in files written before per-chunk widths, + /// where every chunk is packed at `bit_width`. + #[prost(bytes = "vec", tag = "4")] + pub(crate) bit_widths: Vec, +} + +impl ArrayHash for BitPackedV2Data { + fn array_hash(&self, state: &mut H, accuracy: EqMode) { + self.offset.hash(state); + self.widths.hash(state); + self.packed.array_hash(state, accuracy); + self.patches_data.hash(state); + } +} + +impl ArrayEq for BitPackedV2Data { + fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { + self.offset == other.offset + && self.widths == other.widths + && self.packed.array_eq(&other.packed, accuracy) + && self.patches_data == other.patches_data + } +} + +impl VTable for BitPackedV2 { + type TypedArrayData = BitPackedV2Data; + + type OperationsVTable = Self; + type ValidityVTable = Self; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("fastlanes.bitpacked_v2"); + *ID + } + + fn validate( + &self, + data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let bp_slots = BitPackedV2SlotsView::from_slots(slots); + + let validity = child_to_validity(bp_slots.validity_child, dtype.nullability()); + let patches = + PatchesData::patches_from_slots(data.patches_data.as_ref(), len, slots, PATCH_SLOTS); + BitPackedV2Data::validate( + &data.packed, + dtype.as_ptype(), + &validity, + patches.as_ref(), + &data.widths, + len, + data.offset, + ) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 1 + } + + fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + match idx { + 0 => array.packed().clone(), + _ => vortex_panic!("BitPackedV2Array buffer index {idx} out of bounds"), + } + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + match idx { + 0 => Some("packed".to_string()), + _ => None, + } + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + let mut data = array.data().clone(); + data.packed = buffers[0].clone(); + Ok( + ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) + .with_slots(array.slots().iter().cloned().collect()), + ) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + Ok(Some( + BitPackedV2Metadata { + bit_width: array.bit_width() as u32, + offset: array.offset() as u32, + patches: array + .patches() + .map(|p| p.to_metadata(array.len(), array.dtype())) + .transpose()?, + bit_widths: array.chunk_widths().as_slice().to_vec(), + } + .encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = BitPackedV2Metadata::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 slots = { + let mut s = ArraySlots::with_capacity(4); + PatchesData::push_slots(&mut s, patches.as_ref()); + s.push(validity_to_child(&validity, len)); + s + }; + let bit_width = u8::try_from(metadata.bit_width).map_err(|_| { + vortex_err!( + "BitPackedV2Metadata bit_width {} does not fit in u8", + metadata.bit_width + ) + })?; + // Files written with a single width carry no per-chunk widths: expand it. + let widths = if metadata.bit_widths.is_empty() { + ChunkWidths::uniform(bit_width, (len + metadata.offset as usize).div_ceil(1024)) + } else { + ChunkWidths::new(Buffer::from(metadata.bit_widths)) + }; + let data = BitPackedV2Data::try_new( + packed, + patches, + widths, + u16::try_from(metadata.offset).map_err(|_| { + vortex_err!( + "BitPackedV2Metadata offset {} does not fit in u16", + metadata.offset + ) + })?, + )?; + Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + } + + fn append_to_builder( + array: ArrayView<'_, Self>, + builder: &mut dyn ArrayBuilder, + ctx: &mut ExecutionCtx, + ) -> VortexResult<()> { + match_each_integer_ptype!(array.dtype().as_ptype(), |T| { + unpack_into_primitive_builder::( + array, + builder + .as_any_mut() + .downcast_mut() + .vortex_expect("bit packed array must canonicalize into a primitive array"), + ctx, + ) + }) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + BitPackedV2Slots::NAMES[idx].to_string() + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + require_patches!( + array, + BitPackedV2Slots::PATCH_INDICES, + BitPackedV2Slots::PATCH_VALUES, + BitPackedV2Slots::PATCH_CHUNK_OFFSETS + ); + require_validity!(array, BitPackedV2Slots::VALIDITY_CHILD); + + Ok(ExecutionResult::done( + unpack_array(array.as_view(), ctx)?.into_array(), + )) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + RULES.evaluate(array, parent, child_idx) + } +} + +#[derive(Clone, Debug)] +pub struct BitPackedV2; + +impl BitPackedV2 { + /// 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, + widths: ChunkWidths, + len: usize, + offset: u16, + ) -> VortexResult { + let dtype = DType::Primitive(ptype, validity.nullability()); + let slots = { + let mut s = ArraySlots::with_capacity(4); + PatchesData::push_slots(&mut s, patches.as_ref()); + s.push(validity_to_child(&validity, len)); + s + }; + let data = BitPackedV2Data::try_new(packed, patches, widths, offset)?; + Array::try_from_parts(ArrayParts::new(BitPackedV2, dtype, len, data).with_slots(slots)) + } + + pub fn into_parts(array: BitPackedV2Array) -> BitPackedV2DataParts { + let len = array.len(); + let patches = array.patches(); + let validity = array.validity().vortex_expect("BitPackedV2 validity"); + let data = array.into_data(); + BitPackedV2DataParts { + offset: data.offset, + widths: data.widths, + len, + packed: data.packed, + patches, + validity, + } + } + + /// Encode an array into a bitpacked representation with the given bit width. + pub fn encode( + array: &ArrayRef, + bit_width: u8, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + BitPackedV2Data::encode(array, bit_width, ctx) + } +} + +#[cfg(test)] +mod tests { + use prost::Message; + use vortex_array::test_harness::check_metadata; + + use super::BitPackedV2Metadata; + + #[cfg_attr(miri, ignore)] + #[test] + fn test_bitpacked_v2_metadata() { + let metadata = BitPackedV2Metadata { + bit_width: 24, + offset: 1023, + patches: None, + bit_widths: vec![3, 24, 0], + } + .encode_to_vec(); + check_metadata("bitpacked_v2.metadata", &metadata); + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs b/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs new file mode 100644 index 00000000000..674cb50e017 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs @@ -0,0 +1,293 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayView; +use vortex_array::ExecutionCtx; +use vortex_array::scalar::Scalar; +use vortex_array::vtable::OperationsVTable; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::bitpacking_v2::array::BitPackedV2ArrayExt; +use crate::bitpacking_v2::bitpack_decompress; +impl OperationsVTable for BitPackedV2 { + fn scalar_at( + array: ArrayView<'_, BitPackedV2>, + index: usize, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok( + if let Some(patches) = array.patches() + && let Some(patch) = patches.get_patched(index)? + { + patch + } else { + bitpack_decompress::unpack_single(array, index) + }, + ) + } +} + +#[cfg(test)] +mod test { + use std::ops::Range; + + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::SliceArray; + use vortex_array::assert_arrays_eq; + use vortex_array::assert_nth_scalar; + use vortex_array::buffer::BufferHandle; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::patches::Patches; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; + use vortex_buffer::Alignment; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBuffer; + use vortex_buffer::buffer; + + use crate::BitPackedV2; + use crate::BitPackedV2Array; + use crate::BitPackedV2Data; + use crate::ChunkWidths; + use crate::bitpacking_v2::array::BitPackedV2ArrayExt; + use crate::test::SESSION; + + fn bp(array: &ArrayRef, bit_width: u8) -> BitPackedV2Array { + BitPackedV2Data::encode(array, bit_width, &mut SESSION.create_execution_ctx()).unwrap() + } + + fn slice_via_reduce(array: &BitPackedV2Array, range: Range) -> BitPackedV2Array { + let array_ref = array.clone().into_array(); + let slice_array = SliceArray::new(array_ref.clone(), range); + let sliced = array_ref + .reduce_parent(&slice_array.into_array(), 0) + .expect("execute_parent failed") + .expect("expected slice kernel to execute"); + sliced.as_::().into_owned() + } + + #[test] + pub fn slice_block() { + let arr = bp( + &PrimitiveArray::from_iter((0u32..2048).map(|v| v % 64)).into_array(), + 6, + ); + let sliced = slice_via_reduce(&arr, 1024..2048); + assert_nth_scalar!(sliced, 0, 1024u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 1023, + 2047u32 % 64, + &mut SESSION.create_execution_ctx() + ); + assert_eq!(sliced.offset(), 0); + assert_eq!(sliced.len(), 1024); + } + + #[test] + pub fn slice_within_block() { + let arr = bp( + &PrimitiveArray::from_iter((0u32..2048).map(|v| v % 64)).into_array(), + 6, + ); + let sliced = slice_via_reduce(&arr, 512..1434); + assert_nth_scalar!(sliced, 0, 512u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 921, + 1433u32 % 64, + &mut SESSION.create_execution_ctx() + ); + assert_eq!(sliced.offset(), 512); + assert_eq!(sliced.len(), 922); + } + + #[test] + fn slice_within_block_u8s() { + let packed = bp( + &PrimitiveArray::from_iter((0..10_000).map(|i| (i % 63) as u8)).into_array(), + 7, + ); + + let compressed = packed.slice(768..9999).unwrap(); + assert_nth_scalar!( + compressed, + 0, + (768 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + compressed, + compressed.len() - 1, + (9998 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn slice_block_boundary_u8s() { + let packed = bp( + &PrimitiveArray::from_iter((0..10_000).map(|i| (i % 63) as u8)).into_array(), + 7, + ); + + let compressed = packed.slice(7168..9216).unwrap(); + assert_nth_scalar!( + compressed, + 0, + (7168 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + compressed, + compressed.len() - 1, + (9215 % 63) as u8, + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn double_slice_within_block() { + let arr = bp( + &PrimitiveArray::from_iter((0u32..2048).map(|v| v % 64)).into_array(), + 6, + ); + let sliced = slice_via_reduce(&arr, 512..1434); + assert_nth_scalar!(sliced, 0, 512u32 % 64, &mut SESSION.create_execution_ctx()); + assert_nth_scalar!( + sliced, + 921, + 1433u32 % 64, + &mut SESSION.create_execution_ctx() + ); + assert_eq!(sliced.offset(), 512); + assert_eq!(sliced.len(), 922); + let doubly_sliced = slice_via_reduce(&sliced, 127..911); + assert_nth_scalar!( + doubly_sliced, + 0, + (512u32 + 127) % 64, + &mut SESSION.create_execution_ctx() + ); + assert_nth_scalar!( + doubly_sliced, + 783, + (512u32 + 910) % 64, + &mut SESSION.create_execution_ctx() + ); + assert_eq!(doubly_sliced.offset(), 639); + assert_eq!(doubly_sliced.len(), 784); + } + + #[test] + fn slice_empty_patches() { + let mut ctx = SESSION.create_execution_ctx(); + // We create an array that has 1 element that does not fit in the 6-bit range. + let array = BitPackedV2Data::encode(&buffer![0u32..=64].into_array(), 6, &mut ctx).unwrap(); + + assert!(array.patches().is_some()); + + let patch_indices = array.patches().unwrap().indices().clone(); + assert_eq!(patch_indices.len(), 1); + + // Slicing with patches requires the execute path (not reduce) since patches.slice() + // reads buffers. The slice range 0..64 excludes the patch at index 64, so the + // resulting array should have no patches. + let array_ref = array.into_array(); + let slice_array = SliceArray::new(array_ref, 0..64); + let mut ctx = SESSION.create_execution_ctx(); + let sliced_bp = slice_array + .into_array() + .execute::(&mut ctx) + .expect("slice execution failed") + .as_::() + .into_owned(); + assert!(sliced_bp.patches().is_none()); + } + + #[test] + fn take_after_slice() { + // Check that our take implementation respects the offsets applied after slicing. + + let array = bp( + &PrimitiveArray::from_iter((63u32..).take(3072)).into_array(), + 6, + ); + + // Slice the array. + // The resulting array will still have 3 1024-element chunks. + let sliced = array.slice(922..2061).unwrap(); + + // Take one element from each chunk. + // Chunk 1: physical indices 922-1023, logical indices 0-101 + // Chunk 2: physical indices 1024-2047, logical indices 102-1125 + // Chunk 3: physical indices 2048-2060, logical indices 1126-1138 + + let taken = sliced + .take(buffer![101i64, 1125, 1138].into_array()) + .unwrap(); + assert_eq!(taken.len(), 3); + } + + #[test] + fn scalar_at_invalid_patches() { + let packed_array = BitPackedV2::try_new( + BufferHandle::new_host(ByteBuffer::copy_from_aligned( + [0u8; 128], + Alignment::of::(), + )), + PType::U32, + Validity::AllInvalid, + Some( + Patches::new( + 8, + 0, + buffer![1u32].into_array(), + PrimitiveArray::new(buffer![999u32], Validity::AllValid).into_array(), + None, + ) + .unwrap(), + ), + ChunkWidths::uniform(1, 1), + 8, + 0, + ) + .unwrap() + .into_array(); + assert_eq!( + packed_array + .execute_scalar(1, &mut SESSION.create_execution_ctx()) + .unwrap(), + Scalar::null(DType::Primitive(PType::U32, Nullability::Nullable)) + ); + } + + #[test] + fn scalar_at() { + let mut ctx = SESSION.create_execution_ctx(); + let values = (0u32..257).collect::>(); + let uncompressed = values.clone().into_array(); + let packed = BitPackedV2Data::encode(&uncompressed, 8, &mut ctx).unwrap(); + assert!(packed.patches().is_some()); + + let patches = packed.patches().unwrap().indices().clone(); + assert_eq!( + usize::try_from( + &patches + .execute_scalar(0, &mut SESSION.create_execution_ctx()) + .unwrap() + ) + .unwrap(), + 256 + ); + + let expected = PrimitiveArray::from_iter(values.iter().copied()); + assert_arrays_eq!(packed, expected, &mut ctx); + } +} diff --git a/encodings/fastlanes/src/bitpacking_v2/vtable/rules.rs b/encodings/fastlanes/src/bitpacking_v2/vtable/rules.rs new file mode 100644 index 00000000000..e301ada1e09 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/rules.rs @@ -0,0 +1,13 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::arrays::slice::SliceReduceAdaptor; +use vortex_array::optimizer::rules::ParentRuleSet; +use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; + +use crate::BitPackedV2; + +pub(crate) const RULES: ParentRuleSet = ParentRuleSet::new(&[ + ParentRuleSet::lift(&CastReduceAdaptor(BitPackedV2)), + ParentRuleSet::lift(&SliceReduceAdaptor(BitPackedV2)), +]); diff --git a/encodings/fastlanes/src/bitpacking_v2/vtable/validity.rs b/encodings/fastlanes/src/bitpacking_v2/vtable/validity.rs new file mode 100644 index 00000000000..9591439f51f --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/validity.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayView; +use vortex_array::validity::Validity; +use vortex_array::vtable::ValidityVTable; +use vortex_error::VortexResult; + +use crate::BitPackedV2; +use crate::BitPackedV2ArrayExt; + +impl ValidityVTable for BitPackedV2 { + fn validity(array: ArrayView<'_, BitPackedV2>) -> VortexResult { + Ok(BitPackedV2ArrayExt::validity(&array)) + } +} diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index 43d83c6fc7f..7fb0feb527d 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -27,6 +27,14 @@ //! but are not fully binary compatible. See the underlying [fastlanes](https://github.com/spiraldb/fastlanes) crate for more details. pub use bitpacking::*; +pub use bitpacking_v2::BitPackedV2; +pub use bitpacking_v2::BitPackedV2Array; +pub use bitpacking_v2::BitPackedV2ArrayExt; +pub use bitpacking_v2::BitPackedV2ArraySlotsExt; +pub use bitpacking_v2::BitPackedV2Data; +pub use bitpacking_v2::BitPackedV2DataParts; +pub use bitpacking_v2::BitPackedV2Slots; +pub use bitpacking_v2::ChunkWidths; pub use delta::*; pub use r#for::*; pub use rle::*; @@ -41,6 +49,7 @@ use vortex_error::VortexResult; pub mod bit_transpose; mod bitpacking; +pub mod bitpacking_v2; mod delta; mod r#for; mod rle; @@ -66,6 +75,7 @@ pub(crate) const fn untranspose_idx(idx: usize) -> usize { } use bitpacking::compute::is_constant::BitPackedIsConstantKernel; +use bitpacking_v2::compute::is_constant::BitPackedV2IsConstantKernel; use r#for::compute::is_constant::FoRIsConstantKernel; use r#for::compute::is_sorted::FoRIsSortedKernel; use vortex_array::ArrayVTable; @@ -86,11 +96,13 @@ pub fn initialize(session: &VortexSession) { } else { session.arrays().register(BitPacked); } + session.arrays().register(BitPackedV2); session.arrays().register(Delta); session.arrays().register(FoR); session.arrays().register(RLE); session.arrays().register(TransposedBool); bitpacking::initialize(session); + bitpacking_v2::initialize(session); r#for::initialize(session); rle::initialize(session); @@ -100,6 +112,11 @@ pub fn initialize(session: &VortexSession) { Some(IsConstant.id()), &BitPackedIsConstantKernel, ); + session.aggregate_fns().register_aggregate_kernel( + BitPackedV2.id(), + Some(IsConstant.id()), + &BitPackedV2IsConstantKernel, + ); session.aggregate_fns().register_aggregate_kernel( FoR.id(), Some(IsConstant.id()),