From 7de5950c61e820042b8e4de4e2b42223ec32bd5e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 3 Sep 2026 10:00:09 -0400 Subject: [PATCH 1/2] feat(fastlanes): add BitPackedV2 with a bit width per FastLanes chunk A new fastlanes.bitpacked_v2 encoding that packs every 1024-element chunk at its own width. Widths are chosen per chunk from the chunk's histogram using its exact packed-block cost plus the exceptions left behind, and the histogram, width choice, exception gathering and packing all happen while the chunk is in cache. Values that do not fit their chunk's width become patches, as in BitPacked. Serde carries one width byte per chunk. This is the array only: encode, decode, scalar_at, validity and serde. Compute kernels follow in a separate change, so filter, take, compare and slice fall back to canonicalisation for now. BitPacked is unchanged and nothing writes v2 yet. Signed-off-by: Matt Katz --- encodings/fastlanes/Cargo.toml | 4 + .../fastlanes/benches/bitpack_chunk_widths.rs | 254 ++++++ .../goldenfiles/bitpacked_v2.metadata | Bin 0 -> 10 bytes .../bitpacking_v2/array/bitpack_compress.rs | 761 ++++++++++++++++++ .../bitpacking_v2/array/bitpack_decompress.rs | 742 +++++++++++++++++ .../fastlanes/src/bitpacking_v2/array/mod.rs | 595 ++++++++++++++ .../src/bitpacking_v2/array/unpack_iter.rs | 339 ++++++++ .../src/bitpacking_v2/chunk_widths_tests.rs | 428 ++++++++++ encodings/fastlanes/src/bitpacking_v2/mod.rs | 22 + .../fastlanes/src/bitpacking_v2/vtable/mod.rs | 370 +++++++++ .../src/bitpacking_v2/vtable/operations.rs | 140 ++++ .../src/bitpacking_v2/vtable/validity.rs | 16 + encodings/fastlanes/src/lib.rs | 10 + 13 files changed, 3681 insertions(+) create mode 100644 encodings/fastlanes/benches/bitpack_chunk_widths.rs create mode 100644 encodings/fastlanes/goldenfiles/bitpacked_v2.metadata create mode 100644 encodings/fastlanes/src/bitpacking_v2/array/bitpack_compress.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/array/mod.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/chunk_widths_tests.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/mod.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/vtable/mod.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/vtable/validity.rs diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..8170bb5ed12 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -70,3 +70,7 @@ harness = false name = "cast_bitpacked" harness = false required-features = ["_test-harness"] + +[[bench]] +name = "bitpack_chunk_widths" +harness = false diff --git a/encodings/fastlanes/benches/bitpack_chunk_widths.rs b/encodings/fastlanes/benches/bitpack_chunk_widths.rs new file mode 100644 index 00000000000..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 0000000000000000000000000000000000000000..868731b909c85c1a06e88c04f44e201918cbd86e GIT binary patch literal 10 Rcmd;(5ctoo#LO(g000TZ0e1iZ literal 0 HcmV?d00001 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..dc4a240b97b --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs @@ -0,0 +1,742 @@ +// 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); + }, + ) +} + +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_into_primitive_builder::( + bitpacked.as_view(), + &mut builder, + &mut SESSION.create_execution_ctx(), + )?; + + 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_into_primitive_builder::( + bitpacked.as_view(), + &mut builder, + &mut SESSION.create_execution_ctx(), + )?; + + 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_into_primitive_builder::( + bitpacked.as_view(), + &mut builder, + &mut SESSION.create_execution_ctx(), + )?; + + 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..9124a0d2f0e --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs @@ -0,0 +1,339 @@ +// 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); + } + + /// 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 + } +} + +/// 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) }) + } +} 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/mod.rs b/encodings/fastlanes/src/bitpacking_v2/mod.rs new file mode 100644 index 00000000000..675bbfd4b63 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/mod.rs @@ -0,0 +1,22 @@ +// 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; + +mod vtable; + +pub use vtable::BitPackedV2; +pub use vtable::BitPackedV2Array; 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..f2ae1879da6 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/mod.rs @@ -0,0 +1,370 @@ +// 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; +mod operations; +mod validity; + +/// A [`BitPackedV2`]-encoded Vortex array. +pub type BitPackedV2Array = Array; + +#[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(), + )) + } +} + +#[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..2eed8e04923 --- /dev/null +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs @@ -0,0 +1,140 @@ +// 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 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::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() + } + + #[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/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..34e1efc92e9 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; @@ -86,6 +95,7 @@ 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); From 34bd0f368a185e968b7e33d6afc86257ef7e9ffb Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 3 Sep 2026 10:04:52 -0400 Subject: [PATCH 2/2] feat(fastlanes): compute kernels for BitPackedV2 Port slice (reduce and execute), take, filter, fused compare, between, streaming predicate, cast and is_constant to BitPackedV2, indexing each chunk's packed block through the width table, and register them alongside the v1 kernels. Signed-off-by: Matt Katz --- .../bitpacking_v2/array/bitpack_decompress.rs | 35 +- .../src/bitpacking_v2/array/unpack_iter.rs | 97 +++++ .../src/bitpacking_v2/compute/between.rs | 259 +++++++++++++ .../src/bitpacking_v2/compute/cast.rs | 271 ++++++++++++++ .../src/bitpacking_v2/compute/compare.rs | 325 ++++++++++++++++ .../bitpacking_v2/compute/compare_fused.rs | 150 ++++++++ .../src/bitpacking_v2/compute/filter.rs | 351 ++++++++++++++++++ .../src/bitpacking_v2/compute/is_constant.rs | 207 +++++++++++ .../src/bitpacking_v2/compute/mod.rs | 117 ++++++ .../src/bitpacking_v2/compute/slice.rs | 107 ++++++ .../bitpacking_v2/compute/stream_predicate.rs | 113 ++++++ .../src/bitpacking_v2/compute/take.rs | 328 ++++++++++++++++ encodings/fastlanes/src/bitpacking_v2/mod.rs | 5 + .../src/bitpacking_v2/vtable/kernels.rs | 47 +++ .../fastlanes/src/bitpacking_v2/vtable/mod.rs | 15 + .../src/bitpacking_v2/vtable/operations.rs | 153 ++++++++ .../src/bitpacking_v2/vtable/rules.rs | 13 + encodings/fastlanes/src/lib.rs | 7 + 18 files changed, 2597 insertions(+), 3 deletions(-) create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/between.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/cast.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/compare.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/compare_fused.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/filter.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/is_constant.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/mod.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/slice.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/stream_predicate.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/compute/take.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/vtable/kernels.rs create mode 100644 encodings/fastlanes/src/bitpacking_v2/vtable/rules.rs diff --git a/encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs b/encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs index dc4a240b97b..b5006344f1a 100644 --- a/encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs +++ b/encodings/fastlanes/src/bitpacking_v2/array/bitpack_decompress.rs @@ -66,6 +66,32 @@ pub(crate) fn unpack_into_primitive_builder( ) } +/// 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, @@ -422,10 +448,11 @@ mod tests { let bitpacked = encode(&empty, 0); let mut builder = PrimitiveBuilder::::new(Nullability::NonNullable); - unpack_into_primitive_builder::( + unpack_map_into_builder( bitpacked.as_view(), &mut builder, &mut SESSION.create_execution_ctx(), + |v: u32| v, )?; let result = builder.finish_into_primitive(); @@ -450,10 +477,11 @@ mod tests { // Unpack into a new builder. let mut builder = PrimitiveBuilder::::with_capacity(Nullability::Nullable, 5); - unpack_into_primitive_builder::( + unpack_map_into_builder( bitpacked.as_view(), &mut builder, &mut SESSION.create_execution_ctx(), + |v: u32| v, )?; let result = builder.finish_into_primitive(); @@ -487,10 +515,11 @@ mod tests { // Unpack into a new builder. let mut builder = PrimitiveBuilder::::with_capacity(Nullability::NonNullable, 100); - unpack_into_primitive_builder::( + unpack_map_into_builder( bitpacked.as_view(), &mut builder, &mut SESSION.create_execution_ctx(), + |v: u32| v, )?; let result = builder.finish_into_primitive(); diff --git a/encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs b/encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs index 9124a0d2f0e..fdc8f17b338 100644 --- a/encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs +++ b/encodings/fastlanes/src/bitpacking_v2/array/unpack_iter.rs @@ -192,6 +192,60 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { 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, @@ -252,6 +306,43 @@ impl<'a, T: PhysicalPType, S: UnpackStrategy> UnpackedChunks<'a, T, S> { } } +/// 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( @@ -337,3 +428,9 @@ impl<'a, T: BitPackedV2 + 'a> LendingIterator for BitUnpackIterator<'a, T> { 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/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 index 675bbfd4b63..de0ad52f90b 100644 --- a/encodings/fastlanes/src/bitpacking_v2/mod.rs +++ b/encodings/fastlanes/src/bitpacking_v2/mod.rs @@ -15,8 +15,13 @@ 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 index f2ae1879da6..6b3f2be44d1 100644 --- a/encodings/fastlanes/src/bitpacking_v2/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/mod.rs @@ -51,12 +51,19 @@ 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")] @@ -296,6 +303,14 @@ impl VTable for BitPackedV2 { 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)] diff --git a/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs b/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs index 2eed8e04923..674cb50e017 100644 --- a/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs +++ b/encodings/fastlanes/src/bitpacking_v2/vtable/operations.rs @@ -30,11 +30,15 @@ impl OperationsVTable for BitPackedV2 { #[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; @@ -58,6 +62,155 @@ mod test { 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. 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/lib.rs b/encodings/fastlanes/src/lib.rs index 34e1efc92e9..7fb0feb527d 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -75,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; @@ -101,6 +102,7 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(RLE); session.arrays().register(TransposedBool); bitpacking::initialize(session); + bitpacking_v2::initialize(session); r#for::initialize(session); rle::initialize(session); @@ -110,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()),