From 3e5dc30b2a64dc3457a7c624a2bd3658ddae23ac Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 18 Sep 2026 15:19:45 -0400 Subject: [PATCH] feat: select bitpacked widths independently per chunk Signed-off-by: "Matt Katz" --- .../src/bitpacking/array/bitpack_compress.rs | 125 ++++++- .../src/bitpacking/chunk_widths_tests.rs | 305 +++++++++++++++++- encodings/fastlanes/src/bitpacking/mod.rs | 5 +- 3 files changed, 416 insertions(+), 19 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs index c5f361dcdb4..d10a70d32da 100644 --- a/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs +++ b/encodings/fastlanes/src/bitpacking/array/bitpack_compress.rs @@ -35,7 +35,25 @@ use crate::bitpack_decompress::count_exceptions; use crate::bitpacking::array::ChunkWidths; use crate::bitpacking::array::chunk_packed_bytes; -/// Encode with caller-supplied chunk widths, gathering exceptions for values that do not fit. +/// Choose a cost-model width for each chunk, then pack values and gather exceptions. +pub fn bitpack_to_best_chunk_widths( + array: &PrimitiveArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let plan = chunk_width_plan(array.as_view(), ctx)?; + bitpack_encode_planned(array, plan, ctx) +} + +/// 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, @@ -58,7 +76,7 @@ pub fn bitpack_encode_with_widths( /// Bit-pack `array` at the single best global width chosen by [`find_best_bit_width`]. /// /// Every chunk shares that width, so the result serializes under the original -/// `fastlanes.bitpacked` format. +/// `fastlanes.bitpacked` format. See [`bitpack_to_best_chunk_widths`] for per-chunk widths. pub fn bitpack_to_best_bit_width( array: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -171,6 +189,76 @@ struct ChunkWidthPlan { 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, @@ -383,6 +471,26 @@ where } } +/// The width minimising one chunk's cost: its packed block plus the exceptions left behind. +/// +/// A chunk always occupies a whole `128 * width` byte block, so a partial trailing chunk is +/// charged for its padding. +fn best_chunk_width(bit_width_freq: &[usize], bytes_per_exception: usize) -> u8 { + let len: usize = bit_width_freq.iter().sum(); + let mut num_packed = 0; + let mut best_cost = usize::MAX; + let mut best_width = 0; + for (bit_width, freq) in bit_width_freq.iter().enumerate() { + num_packed += *freq; + let cost = chunk_packed_bytes(bit_width as u8) + (len - num_packed) * bytes_per_exception; + if cost < best_cost { + best_cost = cost; + best_width = bit_width; + } + } + best_width as u8 +} + pub fn bit_width_histogram( array: ArrayView<'_, Primitive>, ctx: &mut ExecutionCtx, @@ -536,6 +644,19 @@ mod tests { 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(); diff --git a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs index 1600a6491bb..8c7d330e41c 100644 --- a/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs +++ b/encodings/fastlanes/src/bitpacking/chunk_widths_tests.rs @@ -15,20 +15,29 @@ 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::Constant; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::SliceArray; use vortex_array::arrays::slice::SliceKernel; use vortex_array::assert_arrays_eq; use vortex_array::buffer::BufferHandle; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::cast::test_cast_conformance; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::compute::conformance::take::test_take_conformance; use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::between::BetweenOptions; +use vortex_array::scalar_fn::fns::between::StrictComparison; +use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::serde::ArrayChildren; use vortex_array::serde::SerializeOptions; use vortex_array::serde::SerializedArray; @@ -40,6 +49,7 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::ReadContext; @@ -50,9 +60,11 @@ use crate::BitPackedArraySlotsExt; use crate::BitPackedPlugin; use crate::ChunkWidths; use crate::FL_CHUNK_SIZE; +use crate::FoR; use crate::bitpacked_v2_id; use crate::bitpacking::bitpack_compress::bitpack_encode_with_widths; use crate::bitpacking::bitpack_compress::bitpack_to_best_bit_width; +use crate::bitpacking::bitpack_compress::bitpack_to_best_chunk_widths; use crate::bitpacking::plugin::BitPackedMetadata; use crate::bitpacking::plugin::BitPackedV2Metadata; @@ -84,26 +96,264 @@ fn varied(len_tail: usize) -> Vec { fn encode(values: &[u32]) -> VortexResult { let mut ctx = SESSION.create_execution_ctx(); - let widths = ChunkWidths::new(Buffer::from_iter(values.chunks(FL_CHUNK_SIZE).map( - |chunk| { - chunk - .iter() - .map(|v| (u32::BITS - v.leading_zeros()) as u8) - .max() - .unwrap_or(0) - }, - ))); - bitpack_encode_with_widths( - &PrimitiveArray::from_iter(values.iter().copied()), - widths, - &mut 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(&mut SESSION.create_execution_ctx())?; + 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 + .chunk_widths(&mut SESSION.create_execution_ctx())? + .max_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(&mut SESSION.create_execution_ctx())? + .len(), + 3 + ); + assert_eq!( + packed + .chunk_widths(&mut SESSION.create_execution_ctx())? + .uniform_width(), + Some(7) + ); + assert_eq!( + packed + .chunk_widths(&mut SESSION.create_execution_ctx())? + .max_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(&mut SESSION.create_execution_ctx())?.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(&mut SESSION.create_execution_ctx())? + .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(); @@ -131,6 +381,20 @@ fn explicit_widths_including_full_width_chunk() -> VortexResult<()> { Ok(()) } +#[test] +fn for_fused_decode() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = varied(100); + let packed = encode(&values)?.into_array(); + let for_array = FoR::try_new(packed, Scalar::from(1000u32))?; + assert_arrays_eq!( + for_array, + PrimitiveArray::from_iter(values.iter().map(|v| v + 1000)), + &mut ctx + ); + Ok(()) +} + /// Serialize `array` through the session and read it back through the plugin registered for the /// serialized ID, as a file reader would. fn serde_roundtrip(array: &BitPackedArray) -> VortexResult<(ArrayId, Vec, ArrayRef)> { @@ -556,3 +820,16 @@ fn each_format_keeps_its_contract() -> VortexResult<()> { ); Ok(()) } + +#[rstest] +#[case::varied(encode(&varied(100)).unwrap())] +#[case::varied_exact(encode(&varied(0)).unwrap())] +fn conformance(#[case] array: BitPackedArray) { + let mut ctx = SESSION.create_execution_ctx(); + let array = array.into_array(); + test_array_consistency(&array, &mut ctx); + test_take_conformance(&array, &mut ctx); + test_filter_conformance(&array, &mut ctx); + test_cast_conformance(&array, &mut ctx); + test_binary_numeric_array(&array, &mut ctx); +} diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index 24864c0c3e1..61912717c2c 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -13,6 +13,8 @@ pub use array::bitpack_decompress; pub use array::chunk_packed_bytes; pub use array::unpack_iter; +#[cfg(test)] +mod chunk_widths_tests; pub(crate) mod compute; mod plugin; @@ -27,6 +29,3 @@ pub use vtable::BitPackedArray; pub(crate) fn initialize(session: &vortex_session::VortexSession) { vtable::initialize(session); } - -#[cfg(test)] -mod chunk_widths_tests;