From 3877465a7a7485f9ac4b7ea4240309a833ce2b53 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 16:23:13 -0400 Subject: [PATCH 01/16] Add decimal byte-part splitting and assembly helpers Introduce typed splitting and reassembly for i128 and i256 decimals, including sign extension and boundary tests. Route existing single-part canonicalization through the same assembly helper without changing its wire representation. Signed-off-by: "Matt Katz" --- .../src/decimal_byte_parts/limbs.rs | 439 ++++++++++++++++++ .../src/decimal_byte_parts/mod.rs | 33 +- 2 files changed, 450 insertions(+), 22 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs new file mode 100644 index 00000000000..dfa7dfc1b3c --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs @@ -0,0 +1,439 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splitting decimal values into 64-bit parts, and reassembling them. +//! +//! A `DecimalByteParts` array stores each value as a signed most significant part (MSP) +//! followed by `k` unsigned 64-bit lower parts ordered most significant first. The encoded +//! value is +//! +//! ```text +//! msp * 2^(64k) + Σ_{i, +} + +/// The decimal storage type that reassembling the given parts produces. +/// +/// # Errors +/// +/// Returns an error if `msp_ptype` is not a signed integer, or if there are more than +/// [`MAX_LOWER_PARTS`] lower parts. +pub(crate) fn assembled_values_type( + msp_ptype: PType, + lower_part_count: usize, +) -> VortexResult { + if lower_part_count > MAX_LOWER_PARTS { + vortex_bail!("at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}"); + } + if lower_part_count == 0 { + return DecimalType::try_from(msp_ptype); + } + let bits = msp_ptype.bit_width() + LOWER_PART_BITS * lower_part_count; + Ok(if bits <= 128 { + DecimalType::I128 + } else { + DecimalType::I256 + }) +} + +/// Split a canonical decimal array into a signed most significant part and unsigned 64-bit +/// lower parts. +/// +/// Values narrower than 128 bits are already a single signed part, so they are returned +/// with no lower parts. `i128` values split into an `i64` MSP and one lower part, `i256` +/// values into an `i64` MSP and three lower parts. +/// +/// # Errors +/// +/// Returns an error if the array's validity cannot be derived. +pub fn split_decimal(decimal: &DecimalArray) -> VortexResult { + let validity = decimal.validity()?; + Ok(match decimal.values_type() { + DecimalType::I8 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I16 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I32 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I64 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I128 => { + let (msp, lower) = split_i128(&decimal.buffer::()); + DecimalParts::new(msp, [lower], validity) + } + DecimalType::I256 => { + let (msp, lower) = split_i256(&decimal.buffer::()); + DecimalParts::new(msp, lower, validity) + } + }) +} + +/// Reassemble decimal byte parts into a canonical decimal array. +/// +/// The parts must already be canonical primitive arrays: a signed MSP, and `u64` lower +/// parts ordered most significant first. +/// +/// # Errors +/// +/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity +/// cannot be derived. +pub(crate) fn assemble_decimal( + msp: &PrimitiveArray, + lower_parts: &[PrimitiveArray], + decimal_dtype: DecimalDType, +) -> VortexResult { + let validity = msp.validity()?; + if lower_parts.is_empty() { + return Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { + // SAFETY: the buffer is typed by the array's own ptype, the decimal dtype is the + // array's, and the validity is taken from the same array. + unsafe { DecimalArray::new_unchecked(msp.to_buffer::

(), decimal_dtype, validity) } + })); + } + + // Slice every part to the MSP's length up front: the assembly loops then index slices the + // compiler knows are long enough, so the per-row bounds checks fall away. + let len = msp.len(); + let lower: Vec<&[u64]> = lower_parts + .iter() + .map(|part| { + vortex_ensure!( + part.dtype() == &LOWER_PART_DTYPE, + "lower part must be non-nullable u64" + ); + let part = part.as_slice::(); + vortex_ensure!( + part.len() >= len, + "lower part has len {}, expected at least {len}", + part.len() + ); + Ok(&part[..len]) + }) + .collect::>()?; + + // The part count is dispatched to a constant so every 64-bit word lands at a compile-time + // index. Leaving it dynamic costs 1.8x on the `i256` path — see `benches/decimal_assemble.rs`. + let values = match assembled_values_type(msp.ptype(), lower.len())? { + // A single lower part can never widen to an `i256`: the MSP is at most 64 bits, so + // 64 + 64 fits an `i128` and takes the branch below. + DecimalType::I256 => match lower.as_slice() { + [first, second] => assemble_i256(msp, [first, second]), + [first, second, third] => assemble_i256(msp, [first, second, third]), + _ => vortex_bail!("unsupported lower part count {}", lower.len()), + }, + _ => { + return Ok(DecimalArray::new( + assemble_i128(msp, lower[0]), + decimal_dtype, + validity, + )); + } + }; + Ok(DecimalArray::new(values, decimal_dtype, validity)) +} + +/// 64-bit words in an `i256`. +const VALUE_WORDS: usize = 4; + +/// The 64-bit words of an `i256`, ascending significance. +/// +/// An `i256` is exactly `{_0: u64, _1: u64, _2: u64, _3: i64}` — three unsigned words beneath +/// a single signed one — which is the same shape this encoding stores. That is why splitting +/// and reassembling are pure reinterpretation rather than arithmetic: no carry ever crosses a +/// word boundary, so each word can be compressed independently and put back verbatim. +/// +/// The sign lives in the most significant word alone. When the most significant part is +/// narrower than 64 bits, or sits below word 3, the words above it are its sign extension. +type ValueWords = [u64; VALUE_WORDS]; + +/// Reinterpret an `i256` as its 64-bit words. +#[inline] +const fn i256_to_words(value: i256) -> ValueWords { + let (low, high) = value.to_parts(); + #[expect( + clippy::cast_possible_truncation, + reason = "each cast takes the low 64 bits of a word pair by construction" + )] + [ + low as u64, + (low >> LOWER_PART_BITS) as u64, + high as u64, + (high >> LOWER_PART_BITS) as u64, + ] +} + +/// Reinterpret 64-bit words as an `i256`, with the most significant word carrying the sign. +#[inline] +const fn i256_from_words(words: ValueWords) -> i256 { + i256::from_parts( + (words[0] as u128) | ((words[1] as u128) << LOWER_PART_BITS), + ((words[2] as u128) | ((words[3] as u128) << LOWER_PART_BITS)) as i128, + ) +} + +/// The words of a value whose most significant part sits at `msp_word`, with every word above +/// it filled with the MSP's sign. +#[inline] +fn sign_extended_words(msp: i64, msp_word: usize) -> ValueWords { + let mut words = [if msp < 0 { u64::MAX } else { 0 }; VALUE_WORDS]; + words[msp_word] = msp.cast_unsigned(); + words +} + +impl DecimalParts { + /// Parts for a decimal already stored in a single signed integer. + fn flat(values: Buffer, validity: Validity) -> Self { + Self { + msp: PrimitiveArray::new(values, validity).into_array(), + lower_parts: Vec::new(), + } + } + + fn new( + msp: Buffer, + lower_parts: impl IntoIterator>, + validity: Validity, + ) -> Self { + Self { + msp: PrimitiveArray::new(msp, validity).into_array(), + lower_parts: lower_parts + .into_iter() + .map(|part| PrimitiveArray::new(part, Validity::NonNullable).into_array()) + .collect(), + } + } +} + +#[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "splitting a wide integer into 64-bit windows truncates by construction" +)] +fn split_i128(values: &Buffer) -> (Buffer, Buffer) { + let mut msp = BufferMut::::with_capacity(values.len()); + let mut lower = BufferMut::::with_capacity(values.len()); + for value in values.iter() { + msp.push((value >> LOWER_PART_BITS) as i64); + lower.push(*value as u64); + } + (msp.freeze(), lower.freeze()) +} + +/// The inverse of [`assemble_i256`] at `K == MAX_LOWER_PARTS`: word 3 becomes the signed MSP, +/// and words 2, 1, 0 become the lower parts, most significant first. +fn split_i256(values: &Buffer) -> (Buffer, [Buffer; MAX_LOWER_PARTS]) { + let mut msp = BufferMut::::with_capacity(values.len()); + let mut lower = std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| { + BufferMut::::with_capacity(values.len()) + }); + for value in values.iter() { + let words = i256_to_words(*value); + msp.push(words[MAX_LOWER_PARTS].cast_signed()); + for (part, word) in lower + .iter_mut() + .zip(words.iter().take(MAX_LOWER_PARTS).rev()) + { + part.push(*word); + } + } + (msp.freeze(), lower.map(BufferMut::freeze)) +} + +/// Only one lower part can share 128 bits with a signed MSP, so this shape is fixed. +#[expect( + clippy::useless_conversion, + reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" +)] +fn assemble_i128(msp: &PrimitiveArray, lower: &[u64]) -> Buffer { + // Store into a pre-sized buffer rather than pushing into a reserved one: at 16 bytes per + // row the bounds-checked `push` dominates, and dropping it is 1.6x — see + // `i128_row_write` against `i128_row_const` in `benches/decimal_assemble.rs`. The same + // shape does not pay off for `i256`, where zeroing 32 bytes per row costs more than the + // push it saves. + let mut out = BufferMut::::zeroed(msp.len()); + match_each_signed_integer_ptype!(msp.ptype(), |P| { + for ((slot, value), part) in out + .as_mut_slice() + .iter_mut() + .zip(msp.as_slice::

()) + .zip(lower) + { + *slot = (i128::from(i64::from(*value)) << LOWER_PART_BITS) | i128::from(*part); + } + }); + out.freeze() +} + +/// The lower parts fill the least significant 64-bit words, the MSP the word above them, and +/// the remaining high words are the MSP's sign extension. +/// +/// `K` is a constant so the word indices are compile-time constants and the placement loop +/// unrolls; the same loop with a runtime part count is 1.8x slower. +#[expect( + clippy::useless_conversion, + reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" +)] +fn assemble_i256(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer { + let mut out = BufferMut::::with_capacity(msp.len()); + match_each_signed_integer_ptype!(msp.ptype(), |P| { + for (row, value) in msp.as_slice::

().iter().enumerate() { + // The MSP occupies word `K`, the lower parts the `K` words beneath it most + // significant first, and anything above word `K` is the MSP's sign. + let mut words = sign_extended_words(i64::from(*value), K); + for (i, part) in lower.iter().enumerate() { + words[K - 1 - i] = part[row]; + } + out.push(i256_from_words(words)); + } + }); + out.freeze() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::DecimalArray; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::i256; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::*; + + fn round_trip(decimal: DecimalArray) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal)?; + let msp = parts.msp.execute::(&mut ctx)?; + let lower = parts + .lower_parts + .into_iter() + .map(|part| part.execute::(&mut ctx)) + .collect::>>()?; + assemble_decimal(&msp, &lower, decimal.decimal_dtype()) + } + + #[rstest] + #[case::zero(0)] + #[case::one(1)] + #[case::minus_one(-1)] + #[case::limb_boundary(1i128 << 64)] + #[case::just_below_limb_boundary((1i128 << 64) - 1)] + #[case::negative_limb_boundary(-(1i128 << 64))] + #[case::max(i128::MAX)] + #[case::min(i128::MIN)] + fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) + } + + #[rstest] + #[case::zero(i256::ZERO)] + #[case::one(i256::ONE)] + #[case::minus_one(i256::ZERO - i256::ONE)] + #[case::max(i256::MAX)] + #[case::min(i256::MIN)] + #[case::word_1(i256::from_parts(1u128 << 64, 0))] + #[case::word_2(i256::from_parts(0, 1))] + #[case::word_3(i256::from_parts(0, 1i128 << 64))] + #[case::mixed(i256::from_parts(u128::MAX, -3))] + fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) + } + + #[test] + fn test_split_narrow_decimal_has_no_lower_parts() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![1i32, 2, 3], + DecimalDType::new(9, 2), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal)?; + assert!(parts.lower_parts.is_empty()); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); + Ok(()) + } + + #[test] + fn test_split_i256_part_count_and_types() -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), + DecimalDType::new(76, 0), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal)?; + assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); + for part in &parts.lower_parts { + assert_eq!(part.dtype(), &LOWER_PART_DTYPE); + } + Ok(()) + } + + #[test] + fn test_assembled_values_type() -> VortexResult<()> { + assert_eq!(assembled_values_type(PType::I32, 0)?, DecimalType::I32); + assert_eq!(assembled_values_type(PType::I64, 1)?, DecimalType::I128); + assert_eq!(assembled_values_type(PType::I8, 1)?, DecimalType::I128); + assert_eq!(assembled_values_type(PType::I8, 2)?, DecimalType::I256); + assert_eq!(assembled_values_type(PType::I64, 3)?, DecimalType::I256); + assert!(assembled_values_type(PType::I64, 4).is_err()); + Ok(()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index d5b0024f5b7..c05252ca45d 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -9,6 +9,10 @@ use vortex_array::Array; use vortex_array::ArrayParts; use vortex_array::ArrayView; pub(crate) mod compute; +mod limbs; +pub use limbs::DecimalParts; +pub use limbs::MAX_LOWER_PARTS; +pub use limbs::split_decimal; mod rules; mod slice; @@ -22,13 +26,11 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::array_slots; -use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::PType; -use vortex_array::match_each_signed_integer_ptype; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -46,6 +48,7 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::decimal_byte_parts::limbs::assemble_decimal; use crate::decimal_byte_parts::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. @@ -266,26 +269,12 @@ fn to_canonical_decimal( array: &DecimalBytePartsArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): support parts len != 1 - let prim = array.msp().clone().execute::(ctx)?; - // Depending on the decimal type and the min/max of the primitive array we can choose - // the correct buffer size - - Ok(match_each_signed_integer_ptype!(prim.ptype(), |P| { - // SAFETY: The primitive array's buffer is already validated with correct type. - // The decimal dtype matches the array's dtype, and validity is preserved. - unsafe { - DecimalArray::new_unchecked( - prim.to_buffer::

(), - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - prim.validity()?, - ) - } - .into_array() - })) + let msp = array.msp().clone().execute::(ctx)?; + let decimal_dtype = *array + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype"); + Ok(assemble_decimal(&msp, &[], decimal_dtype)?.into_array()) } impl OperationsVTable for DecimalByteParts { From 966977b9a25252d403df68f6cc621a4b48a6445f Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 17:42:24 -0400 Subject: [PATCH 02/16] Zero null words when splitting wide decimal storage Resolve validity once for wide storage and populate only valid rows in zero-initialized part buffers, so arbitrary null payloads do not inflate lower-part compression. Preserve the all-valid loops and narrow zero-copy path. Cover sliced, empty, nullable, and wider-than-precision storage. Signed-off-by: "Matt Katz" --- .../{limbs.rs => limbs/mod.rs} | 199 ++++++------------ .../src/decimal_byte_parts/limbs/tests.rs | 156 ++++++++++++++ 2 files changed, 225 insertions(+), 130 deletions(-) rename encodings/decimal-byte-parts/src/decimal_byte_parts/{limbs.rs => limbs/mod.rs} (71%) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs similarity index 71% rename from encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs rename to encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs index dfa7dfc1b3c..92d7a38d9a4 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -16,6 +16,7 @@ //! 64-bit window of the magnitude. use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; @@ -33,6 +34,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_mask::Mask; /// The maximum number of lower parts an encoded decimal can carry. /// @@ -86,11 +88,13 @@ pub(crate) fn assembled_values_type( /// Values narrower than 128 bits are already a single signed part, so they are returned /// with no lower parts. `i128` values split into an `i64` MSP and one lower part, `i256` /// values into an `i64` MSP and three lower parts. +/// Lower parts are non-nullable, with zeroes at null positions so arbitrary null-slot bytes +/// do not affect their compression. The MSP retains the decimal's validity. /// /// # Errors /// -/// Returns an error if the array's validity cannot be derived. -pub fn split_decimal(decimal: &DecimalArray) -> VortexResult { +/// Returns an error if the array's validity cannot be derived or executed. +pub fn split_decimal(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { let validity = decimal.validity()?; Ok(match decimal.values_type() { DecimalType::I8 => DecimalParts::flat(decimal.buffer::(), validity), @@ -98,11 +102,13 @@ pub fn split_decimal(decimal: &DecimalArray) -> VortexResult { DecimalType::I32 => DecimalParts::flat(decimal.buffer::(), validity), DecimalType::I64 => DecimalParts::flat(decimal.buffer::(), validity), DecimalType::I128 => { - let (msp, lower) = split_i128(&decimal.buffer::()); + let mask = validity.execute_mask(decimal.len(), ctx)?; + let (msp, lower) = split_i128(&decimal.buffer::(), &mask); DecimalParts::new(msp, [lower], validity) } DecimalType::I256 => { - let (msp, lower) = split_i256(&decimal.buffer::()); + let mask = validity.execute_mask(decimal.len(), ctx)?; + let (msp, lower) = split_i256(&decimal.buffer::(), &mask); DecimalParts::new(msp, lower, validity) } }) @@ -249,32 +255,71 @@ impl DecimalParts { clippy::cast_sign_loss, reason = "splitting a wide integer into 64-bit windows truncates by construction" )] -fn split_i128(values: &Buffer) -> (Buffer, Buffer) { - let mut msp = BufferMut::::with_capacity(values.len()); - let mut lower = BufferMut::::with_capacity(values.len()); - for value in values.iter() { - msp.push((value >> LOWER_PART_BITS) as i64); - lower.push(*value as u64); +fn split_i128(values: &Buffer, validity: &Mask) -> (Buffer, Buffer) { + if validity.all_true() { + let mut msp = BufferMut::::with_capacity(values.len()); + let mut lower = BufferMut::::with_capacity(values.len()); + for value in values.iter() { + msp.push((value >> LOWER_PART_BITS) as i64); + lower.push(*value as u64); + } + return (msp.freeze(), lower.freeze()); + } + + let mut msp = BufferMut::::zeroed(values.len()); + let mut lower = BufferMut::::zeroed(values.len()); + if let Mask::Values(valid) = validity { + let msp = msp.as_mut_slice(); + let lower = lower.as_mut_slice(); + valid.bit_buffer().for_each_set_index(|i| { + let value = values[i]; + msp[i] = (value >> LOWER_PART_BITS) as i64; + lower[i] = value as u64; + }); } (msp.freeze(), lower.freeze()) } /// The inverse of [`assemble_i256`] at `K == MAX_LOWER_PARTS`: word 3 becomes the signed MSP, /// and words 2, 1, 0 become the lower parts, most significant first. -fn split_i256(values: &Buffer) -> (Buffer, [Buffer; MAX_LOWER_PARTS]) { - let mut msp = BufferMut::::with_capacity(values.len()); - let mut lower = std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| { - BufferMut::::with_capacity(values.len()) - }); - for value in values.iter() { - let words = i256_to_words(*value); - msp.push(words[MAX_LOWER_PARTS].cast_signed()); - for (part, word) in lower - .iter_mut() - .zip(words.iter().take(MAX_LOWER_PARTS).rev()) - { - part.push(*word); +fn split_i256( + values: &Buffer, + validity: &Mask, +) -> (Buffer, [Buffer; MAX_LOWER_PARTS]) { + if validity.all_true() { + let mut msp = BufferMut::::with_capacity(values.len()); + let mut lower = std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| { + BufferMut::::with_capacity(values.len()) + }); + for value in values.iter() { + let words = i256_to_words(*value); + msp.push(words[MAX_LOWER_PARTS].cast_signed()); + for (part, word) in lower + .iter_mut() + .zip(words.iter().take(MAX_LOWER_PARTS).rev()) + { + part.push(*word); + } } + return (msp.freeze(), lower.map(BufferMut::freeze)); + } + + let mut msp = BufferMut::::zeroed(values.len()); + let mut lower = + std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| BufferMut::::zeroed(values.len())); + if let Mask::Values(valid) = validity { + let msp = msp.as_mut_slice(); + let mut lower = lower.each_mut().map(BufferMut::as_mut_slice); + valid.bit_buffer().for_each_set_index(|i| { + let words = i256_to_words(values[i]); + msp[i] = words[MAX_LOWER_PARTS].cast_signed(); + for (part, word) in lower + .iter_mut() + .zip(words.iter().take(MAX_LOWER_PARTS).rev()) + { + part[i] = *word; + } + }); } (msp.freeze(), lower.map(BufferMut::freeze)) } @@ -330,110 +375,4 @@ fn assemble_i256(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Bu } #[cfg(test)] -mod tests { - use rstest::rstest; - use vortex_array::VortexSessionExecute; - use vortex_array::array_session; - use vortex_array::arrays::DecimalArray; - use vortex_array::dtype::DecimalDType; - use vortex_array::dtype::i256; - use vortex_array::validity::Validity; - use vortex_buffer::Buffer; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - - use super::*; - - fn round_trip(decimal: DecimalArray) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); - let parts = split_decimal(&decimal)?; - let msp = parts.msp.execute::(&mut ctx)?; - let lower = parts - .lower_parts - .into_iter() - .map(|part| part.execute::(&mut ctx)) - .collect::>>()?; - assemble_decimal(&msp, &lower, decimal.decimal_dtype()) - } - - #[rstest] - #[case::zero(0)] - #[case::one(1)] - #[case::minus_one(-1)] - #[case::limb_boundary(1i128 << 64)] - #[case::just_below_limb_boundary((1i128 << 64) - 1)] - #[case::negative_limb_boundary(-(1i128 << 64))] - #[case::max(i128::MAX)] - #[case::min(i128::MIN)] - fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { - let decimal = DecimalArray::new( - Buffer::from(vec![value]), - DecimalDType::new(38, 2), - Validity::NonNullable, - ); - let round_tripped = round_trip(decimal)?; - assert_eq!(round_tripped.buffer::().as_slice(), &[value]); - Ok(()) - } - - #[rstest] - #[case::zero(i256::ZERO)] - #[case::one(i256::ONE)] - #[case::minus_one(i256::ZERO - i256::ONE)] - #[case::max(i256::MAX)] - #[case::min(i256::MIN)] - #[case::word_1(i256::from_parts(1u128 << 64, 0))] - #[case::word_2(i256::from_parts(0, 1))] - #[case::word_3(i256::from_parts(0, 1i128 << 64))] - #[case::mixed(i256::from_parts(u128::MAX, -3))] - fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { - let decimal = DecimalArray::new( - Buffer::from(vec![value]), - DecimalDType::new(76, 2), - Validity::NonNullable, - ); - let round_tripped = round_trip(decimal)?; - assert_eq!(round_tripped.buffer::().as_slice(), &[value]); - Ok(()) - } - - #[test] - fn test_split_narrow_decimal_has_no_lower_parts() -> VortexResult<()> { - let decimal = DecimalArray::new( - buffer![1i32, 2, 3], - DecimalDType::new(9, 2), - Validity::NonNullable, - ); - let parts = split_decimal(&decimal)?; - assert!(parts.lower_parts.is_empty()); - assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); - Ok(()) - } - - #[test] - fn test_split_i256_part_count_and_types() -> VortexResult<()> { - let decimal = DecimalArray::new( - Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), - DecimalDType::new(76, 0), - Validity::NonNullable, - ); - let parts = split_decimal(&decimal)?; - assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); - assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); - for part in &parts.lower_parts { - assert_eq!(part.dtype(), &LOWER_PART_DTYPE); - } - Ok(()) - } - - #[test] - fn test_assembled_values_type() -> VortexResult<()> { - assert_eq!(assembled_values_type(PType::I32, 0)?, DecimalType::I32); - assert_eq!(assembled_values_type(PType::I64, 1)?, DecimalType::I128); - assert_eq!(assembled_values_type(PType::I8, 1)?, DecimalType::I128); - assert_eq!(assembled_values_type(PType::I8, 2)?, DecimalType::I256); - assert_eq!(assembled_values_type(PType::I64, 3)?, DecimalType::I256); - assert!(assembled_values_type(PType::I64, 4).is_err()); - Ok(()) - } -} +mod tests; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs new file mode 100644 index 00000000000..2e9d39f458c --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::*; + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::all_valid(Validity::AllValid)] +#[case::all_null(Validity::AllInvalid)] +#[case::mixed(Validity::from_iter((0..263).map(|i| i % 3 != 1)))] +fn test_split_zeroes_null_words( + #[case] validity: Validity, + #[values(false, true)] wide_256: bool, + #[values(0, 1, 257)] len: usize, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = if wide_256 { + DecimalArray::new( + buffer![i256::from_i128(-1); 263], + DecimalDType::new(76, 2), + validity, + ) + } else { + DecimalArray::new(buffer![-1i128; 263], DecimalDType::new(38, 2), validity) + }; + let decimal = decimal + .slice(3..len + 3)? + .execute::(&mut ctx)?; + let expected = PrimitiveArray::new( + decimal + .validity()? + .execute_mask(len, &mut ctx)? + .iter() + .map(|valid| if valid { u64::MAX } else { 0 }) + .collect::>(), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + for lower in parts.lower_parts { + assert_arrays_eq!(expected.clone(), lower, &mut ctx); + } + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + Ok(()) +} + +fn round_trip(decimal: DecimalArray) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal, &mut ctx)?; + let msp = parts.msp.execute::(&mut ctx)?; + let lower = parts + .lower_parts + .into_iter() + .map(|part| part.execute::(&mut ctx)) + .collect::>>()?; + assemble_decimal(&msp, &lower, decimal.decimal_dtype()) +} + +#[rstest] +#[case::zero(0)] +#[case::one(1)] +#[case::minus_one(-1)] +#[case::limb_boundary(1i128 << 64)] +#[case::just_below_limb_boundary((1i128 << 64) - 1)] +#[case::negative_limb_boundary(-(1i128 << 64))] +#[case::max(i128::MAX)] +#[case::min(i128::MIN)] +fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) +} + +#[rstest] +#[case::zero(i256::ZERO)] +#[case::one(i256::ONE)] +#[case::minus_one(i256::ZERO - i256::ONE)] +#[case::max(i256::MAX)] +#[case::min(i256::MIN)] +#[case::word_1(i256::from_parts(1u128 << 64, 0))] +#[case::word_2(i256::from_parts(0, 1))] +#[case::word_3(i256::from_parts(0, 1i128 << 64))] +#[case::mixed(i256::from_parts(u128::MAX, -3))] +fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) +} + +#[rstest] +fn test_split_narrow_decimal_has_no_lower_parts( + #[values(Validity::NonNullable, Validity::AllInvalid, Validity::from_iter([true, false, true]))] + validity: Validity, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.lower_parts.is_empty()); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); + let msp = parts.msp.execute::(&mut ctx)?; + assert_eq!( + msp.as_slice::().as_ptr(), + decimal.buffer::().as_ptr() + ); + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + Ok(()) +} + +#[test] +fn test_split_i256_part_count_and_types() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), + DecimalDType::new(76, 0), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); + for part in &parts.lower_parts { + assert_eq!(part.dtype(), &LOWER_PART_DTYPE); + } + Ok(()) +} + +#[test] +fn test_assembled_values_type() -> VortexResult<()> { + assert_eq!(assembled_values_type(PType::I32, 0)?, DecimalType::I32); + assert_eq!(assembled_values_type(PType::I64, 1)?, DecimalType::I128); + assert_eq!(assembled_values_type(PType::I8, 1)?, DecimalType::I128); + assert_eq!(assembled_values_type(PType::I8, 2)?, DecimalType::I256); + assert_eq!(assembled_values_type(PType::I64, 3)?, DecimalType::I256); + assert!(assembled_values_type(PType::I64, 4).is_err()); + Ok(()) +} From 643f4b1b9132a89d668ce8b2849fe8d1a6a43ead Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 22:05:33 -0400 Subject: [PATCH 03/16] Simplify decimal part splitting and assembly Keep i256 words in most-significant-first order and assemble the two 128-bit halves directly. Dispatch on lower-part count, validate signed MSPs and equal child lengths, and cover word order and sign extension. Signed-off-by: "Matt Katz" --- .../src/decimal_byte_parts/limbs/mod.rs | 374 ++++++++---------- .../src/decimal_byte_parts/limbs/tests.rs | 84 +++- 2 files changed, 241 insertions(+), 217 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs index 92d7a38d9a4..a22749fe79d 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Splitting decimal values into 64-bit parts, and reassembling them. +//! Splitting decimal values into 64-bit parts and reassembling them. //! //! A `DecimalByteParts` array stores each value as a signed most significant part (MSP) //! followed by `k` unsigned 64-bit lower parts ordered most significant first. The encoded @@ -11,9 +11,8 @@ //! msp * 2^(64k) + Σ_{i, } -/// The decimal storage type that reassembling the given parts produces. -/// -/// # Errors -/// -/// Returns an error if `msp_ptype` is not a signed integer, or if there are more than -/// [`MAX_LOWER_PARTS`] lower parts. -pub(crate) fn assembled_values_type( - msp_ptype: PType, - lower_part_count: usize, -) -> VortexResult { - if lower_part_count > MAX_LOWER_PARTS { - vortex_bail!("at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}"); +impl DecimalParts { + /// Construct decimal parts from an MSP with no lower parts. + fn from_msp(values: Buffer, validity: Validity) -> Self { + Self { + msp: PrimitiveArray::new(values, validity).into_array(), + lower_parts: Vec::new(), + } } - if lower_part_count == 0 { - return DecimalType::try_from(msp_ptype); + + fn new( + msp: Buffer, + lower_parts: impl IntoIterator>, + validity: Validity, + ) -> Self { + Self { + msp: PrimitiveArray::new(msp, validity).into_array(), + lower_parts: lower_parts + .into_iter() + .map(|part| PrimitiveArray::new(part, Validity::NonNullable).into_array()) + .collect(), + } } - let bits = msp_ptype.bit_width() + LOWER_PART_BITS * lower_part_count; - Ok(if bits <= 128 { - DecimalType::I128 - } else { - DecimalType::I256 - }) } -/// Split a canonical decimal array into a signed most significant part and unsigned 64-bit -/// lower parts. +/// Split a canonical decimal array into a signed most significant part (MSP) and unsigned 64-bit +/// lower parts. The MSP is at most 64 bits. /// /// Values narrower than 128 bits are already a single signed part, so they are returned -/// with no lower parts. `i128` values split into an `i64` MSP and one lower part, `i256` -/// values into an `i64` MSP and three lower parts. -/// Lower parts are non-nullable, with zeroes at null positions so arbitrary null-slot bytes -/// do not affect their compression. The MSP retains the decimal's validity. +/// with no lower parts. `i128` values split into an `i64` MSP and one lower part. `i256` +/// values split into an `i64` MSP and three lower parts. +/// +/// The MSP retains the decimal's validity while lower parts are non-nullable. Lower parts +/// are constructed with zeroes at null positions instead of invalid bytes. /// /// # Errors /// @@ -97,10 +96,10 @@ pub(crate) fn assembled_values_type( pub fn split_decimal(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { let validity = decimal.validity()?; Ok(match decimal.values_type() { - DecimalType::I8 => DecimalParts::flat(decimal.buffer::(), validity), - DecimalType::I16 => DecimalParts::flat(decimal.buffer::(), validity), - DecimalType::I32 => DecimalParts::flat(decimal.buffer::(), validity), - DecimalType::I64 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I8 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I16 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I32 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I64 => DecimalParts::from_msp(decimal.buffer::(), validity), DecimalType::I128 => { let mask = validity.execute_mask(decimal.len(), ctx)?; let (msp, lower) = split_i128(&decimal.buffer::(), &mask); @@ -114,142 +113,9 @@ pub fn split_decimal(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> VortexRe }) } -/// Reassemble decimal byte parts into a canonical decimal array. -/// -/// The parts must already be canonical primitive arrays: a signed MSP, and `u64` lower -/// parts ordered most significant first. +/// Split each `i128` into an `i64` MSP and an `u64` lower part. /// -/// # Errors -/// -/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity -/// cannot be derived. -pub(crate) fn assemble_decimal( - msp: &PrimitiveArray, - lower_parts: &[PrimitiveArray], - decimal_dtype: DecimalDType, -) -> VortexResult { - let validity = msp.validity()?; - if lower_parts.is_empty() { - return Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { - // SAFETY: the buffer is typed by the array's own ptype, the decimal dtype is the - // array's, and the validity is taken from the same array. - unsafe { DecimalArray::new_unchecked(msp.to_buffer::

(), decimal_dtype, validity) } - })); - } - - // Slice every part to the MSP's length up front: the assembly loops then index slices the - // compiler knows are long enough, so the per-row bounds checks fall away. - let len = msp.len(); - let lower: Vec<&[u64]> = lower_parts - .iter() - .map(|part| { - vortex_ensure!( - part.dtype() == &LOWER_PART_DTYPE, - "lower part must be non-nullable u64" - ); - let part = part.as_slice::(); - vortex_ensure!( - part.len() >= len, - "lower part has len {}, expected at least {len}", - part.len() - ); - Ok(&part[..len]) - }) - .collect::>()?; - - // The part count is dispatched to a constant so every 64-bit word lands at a compile-time - // index. Leaving it dynamic costs 1.8x on the `i256` path — see `benches/decimal_assemble.rs`. - let values = match assembled_values_type(msp.ptype(), lower.len())? { - // A single lower part can never widen to an `i256`: the MSP is at most 64 bits, so - // 64 + 64 fits an `i128` and takes the branch below. - DecimalType::I256 => match lower.as_slice() { - [first, second] => assemble_i256(msp, [first, second]), - [first, second, third] => assemble_i256(msp, [first, second, third]), - _ => vortex_bail!("unsupported lower part count {}", lower.len()), - }, - _ => { - return Ok(DecimalArray::new( - assemble_i128(msp, lower[0]), - decimal_dtype, - validity, - )); - } - }; - Ok(DecimalArray::new(values, decimal_dtype, validity)) -} - -/// 64-bit words in an `i256`. -const VALUE_WORDS: usize = 4; - -/// The 64-bit words of an `i256`, ascending significance. -/// -/// An `i256` is exactly `{_0: u64, _1: u64, _2: u64, _3: i64}` — three unsigned words beneath -/// a single signed one — which is the same shape this encoding stores. That is why splitting -/// and reassembling are pure reinterpretation rather than arithmetic: no carry ever crosses a -/// word boundary, so each word can be compressed independently and put back verbatim. -/// -/// The sign lives in the most significant word alone. When the most significant part is -/// narrower than 64 bits, or sits below word 3, the words above it are its sign extension. -type ValueWords = [u64; VALUE_WORDS]; - -/// Reinterpret an `i256` as its 64-bit words. -#[inline] -const fn i256_to_words(value: i256) -> ValueWords { - let (low, high) = value.to_parts(); - #[expect( - clippy::cast_possible_truncation, - reason = "each cast takes the low 64 bits of a word pair by construction" - )] - [ - low as u64, - (low >> LOWER_PART_BITS) as u64, - high as u64, - (high >> LOWER_PART_BITS) as u64, - ] -} - -/// Reinterpret 64-bit words as an `i256`, with the most significant word carrying the sign. -#[inline] -const fn i256_from_words(words: ValueWords) -> i256 { - i256::from_parts( - (words[0] as u128) | ((words[1] as u128) << LOWER_PART_BITS), - ((words[2] as u128) | ((words[3] as u128) << LOWER_PART_BITS)) as i128, - ) -} - -/// The words of a value whose most significant part sits at `msp_word`, with every word above -/// it filled with the MSP's sign. -#[inline] -fn sign_extended_words(msp: i64, msp_word: usize) -> ValueWords { - let mut words = [if msp < 0 { u64::MAX } else { 0 }; VALUE_WORDS]; - words[msp_word] = msp.cast_unsigned(); - words -} - -impl DecimalParts { - /// Parts for a decimal already stored in a single signed integer. - fn flat(values: Buffer, validity: Validity) -> Self { - Self { - msp: PrimitiveArray::new(values, validity).into_array(), - lower_parts: Vec::new(), - } - } - - fn new( - msp: Buffer, - lower_parts: impl IntoIterator>, - validity: Validity, - ) -> Self { - Self { - msp: PrimitiveArray::new(msp, validity).into_array(), - lower_parts: lower_parts - .into_iter() - .map(|part| PrimitiveArray::new(part, Validity::NonNullable).into_array()) - .collect(), - } - } -} - +/// For each valid row, the original value is `msp * 2^64 + lower`. Invalid rows are zeroed. #[expect( clippy::cast_possible_truncation, clippy::cast_sign_loss, @@ -266,8 +132,11 @@ fn split_i128(values: &Buffer, validity: &Mask) -> (Buffer, Buffer::zeroed(values.len()); let mut lower = BufferMut::::zeroed(values.len()); + if let Mask::Values(valid) = validity { let msp = msp.as_mut_slice(); let lower = lower.as_mut_slice(); @@ -280,61 +149,141 @@ fn split_i128(values: &Buffer, validity: &Mask) -> (Buffer, Buffer, validity: &Mask, ) -> (Buffer, [Buffer; MAX_LOWER_PARTS]) { + // With no nulls, append every value without zeroing the output buffers first. if validity.all_true() { let mut msp = BufferMut::::with_capacity(values.len()); let mut lower = std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| { BufferMut::::with_capacity(values.len()) }); for value in values.iter() { - let words = i256_to_words(*value); - msp.push(words[MAX_LOWER_PARTS].cast_signed()); - for (part, word) in lower - .iter_mut() - .zip(words.iter().take(MAX_LOWER_PARTS).rev()) - { - part.push(*word); + let [msp_word, lower_words @ ..] = i256_to_words(*value); + msp.push(msp_word.cast_signed()); + for (part, word) in lower.iter_mut().zip(lower_words) { + part.push(word); } } return (msp.freeze(), lower.map(BufferMut::freeze)); } + // Lower parts are stored as non-nullable arrays, so use zeros at null positions instead + // of copying their garbage values. let mut msp = BufferMut::::zeroed(values.len()); let mut lower = std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| BufferMut::::zeroed(values.len())); + if let Mask::Values(valid) = validity { let msp = msp.as_mut_slice(); let mut lower = lower.each_mut().map(BufferMut::as_mut_slice); valid.bit_buffer().for_each_set_index(|i| { - let words = i256_to_words(values[i]); - msp[i] = words[MAX_LOWER_PARTS].cast_signed(); - for (part, word) in lower - .iter_mut() - .zip(words.iter().take(MAX_LOWER_PARTS).rev()) - { - part[i] = *word; + let [msp_word, lower_words @ ..] = i256_to_words(values[i]); + msp[i] = msp_word.cast_signed(); + for (part, word) in lower.iter_mut().zip(lower_words) { + part[i] = word; } }); } (msp.freeze(), lower.map(BufferMut::freeze)) } -/// Only one lower part can share 128 bits with a signed MSP, so this shape is fixed. +/// Split an `i256` into four `u64` words, most significant first. +#[inline] +const fn i256_to_words(value: i256) -> [u64; 4] { + let (low, high) = value.to_parts(); + #[expect( + clippy::cast_possible_truncation, + reason = "each cast takes the low 64 bits of a word pair by construction" + )] + [ + (high >> LOWER_PART_BITS) as u64, + high as u64, + (low >> LOWER_PART_BITS) as u64, + low as u64, + ] +} + +/// Reassemble primitive arrays that constitute decimal byte parts into a canonical decimal array. +/// +/// The MSP must be signed. There must be between zero and three (inclusive) `u64` lower parts, ordered +/// most significant first. The lower parts must be non-nullable. Every input array must have the same length. +/// +/// With no lower parts, the MSP buffer is reused as the decimal values. One lower part +/// assembles into `i128`. Two or three lower parts assemble into `i256`. +/// +/// # Errors +/// +/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity +/// cannot be derived. +pub(crate) fn assemble_decimal( + msp: &PrimitiveArray, + lower_parts: &[PrimitiveArray], + decimal_dtype: DecimalDType, +) -> VortexResult { + let validity = msp.validity()?; + vortex_ensure!(msp.dtype().as_ptype().is_signed_int()); + + if lower_parts.is_empty() { + return Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { + // SAFETY: the buffer is typed by the array's own ptype, the decimal dtype is the + // array's, and the validity is taken from the same array. + unsafe { DecimalArray::new_unchecked(msp.to_buffer::

(), decimal_dtype, validity) } + })); + } + + let len = msp.len(); + let lower: Vec<&[u64]> = lower_parts + .iter() + .map(|part| { + vortex_ensure!( + part.dtype() == &LOWER_PART_DTYPE, + "lower part must be non-nullable u64" + ); + let part = part.as_slice::(); + vortex_ensure!( + part.len() == len, + "lower part has len {}, expected {len}", + part.len() + ); + Ok(part) + }) + .collect::>()?; + + Ok(match lower.as_slice() { + [first] => DecimalArray::new(assemble_i128(msp, first), decimal_dtype, validity), + [first, second] => { + DecimalArray::new(assemble_i256(msp, [first, second]), decimal_dtype, validity) + } + [first, second, third] => DecimalArray::new( + assemble_i256(msp, [first, second, third]), + decimal_dtype, + validity, + ), + _ => vortex_bail!( + "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", + lower.len() + ), + }) +} + +/// Reassemble a signed MSP and one `u64` lower part into `i128` values. +/// +/// For each row, the result is `msp * 2^64 + lower`. #[expect( clippy::useless_conversion, reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" )] fn assemble_i128(msp: &PrimitiveArray, lower: &[u64]) -> Buffer { - // Store into a pre-sized buffer rather than pushing into a reserved one: at 16 bytes per - // row the bounds-checked `push` dominates, and dropping it is 1.6x — see - // `i128_row_write` against `i128_row_const` in `benches/decimal_assemble.rs`. The same - // shape does not pay off for `i256`, where zeroing 32 bytes per row costs more than the - // push it saves. let mut out = BufferMut::::zeroed(msp.len()); match_each_signed_integer_ptype!(msp.ptype(), |P| { for ((slot, value), part) in out @@ -343,17 +292,19 @@ fn assemble_i128(msp: &PrimitiveArray, lower: &[u64]) -> Buffer { .zip(msp.as_slice::

()) .zip(lower) { + // Sign-extend the MSP, then shift it into the high 64 bits. The unsigned + // lower part fills the low 64 bits. *slot = (i128::from(i64::from(*value)) << LOWER_PART_BITS) | i128::from(*part); } }); out.freeze() } -/// The lower parts fill the least significant 64-bit words, the MSP the word above them, and -/// the remaining high words are the MSP's sign extension. +/// Reassemble a signed MSP and two or three `u64` lower parts into `i256` values. /// -/// `K` is a constant so the word indices are compile-time constants and the placement loop -/// unrolls; the same loop with a runtime part count is 1.8x slower. +/// The last two lower parts form the unsigned low 128 bits. With two lower parts, the +/// signed high 128 bits are the MSP widened to `i128`. With three, the high half contains +/// the MSP followed by the first lower part. #[expect( clippy::useless_conversion, reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" @@ -362,13 +313,18 @@ fn assemble_i256(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Bu let mut out = BufferMut::::with_capacity(msp.len()); match_each_signed_integer_ptype!(msp.ptype(), |P| { for (row, value) in msp.as_slice::

().iter().enumerate() { - // The MSP occupies word `K`, the lower parts the `K` words beneath it most - // significant first, and anything above word `K` is the MSP's sign. - let mut words = sign_extended_words(i64::from(*value), K); - for (i, part) in lower.iter().enumerate() { - words[K - 1 - i] = part[row]; - } - out.push(i256_from_words(words)); + // The last two lower parts always form the unsigned low 128 bits. + let low = + (u128::from(lower[K - 2][row]) << LOWER_PART_BITS) | u128::from(lower[K - 1][row]); + let msp = i128::from(i64::from(*value)); + let high = if K == 2 { + // Widening the MSP supplies the remaining sign bits. + msp + } else { + // With three lower parts, the first one follows the MSP in the high half. + (msp << LOWER_PART_BITS) | i128::from(lower[0][row]) + }; + out.push(i256::from_parts(low, high)); } }); out.freeze() diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs index 2e9d39f458c..57856a4f2c4 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs @@ -144,13 +144,81 @@ fn test_split_i256_part_count_and_types() -> VortexResult<()> { Ok(()) } -#[test] -fn test_assembled_values_type() -> VortexResult<()> { - assert_eq!(assembled_values_type(PType::I32, 0)?, DecimalType::I32); - assert_eq!(assembled_values_type(PType::I64, 1)?, DecimalType::I128); - assert_eq!(assembled_values_type(PType::I8, 1)?, DecimalType::I128); - assert_eq!(assembled_values_type(PType::I8, 2)?, DecimalType::I256); - assert_eq!(assembled_values_type(PType::I64, 3)?, DecimalType::I256); - assert!(assembled_values_type(PType::I64, 4).is_err()); +#[rstest] +fn test_split_i256_part_order( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] validity: Validity, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + buffer![ + i256::from_parts((2u128 << 64) | 3, (1i128 << 64) | 4), + i256::ZERO, + i256::from_parts((6u128 << 64) | 7, (-2i128 << 64) | 5), + ], + DecimalDType::new(76, 0), + validity.clone(), + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_arrays_eq!( + PrimitiveArray::new(buffer![1i64, 0, -2], validity), + parts.msp, + &mut ctx + ); + assert_eq!(parts.lower_parts.len(), 3); + for (part, expected) in parts.lower_parts.into_iter().zip([ + buffer![4u64, 0, 5], + buffer![2u64, 0, 6], + buffer![3u64, 0, 7], + ]) { + assert_arrays_eq!( + PrimitiveArray::new(expected, Validity::NonNullable), + part, + &mut ctx + ); + } + Ok(()) +} + +#[rstest] +fn test_assemble_rejects_mismatched_lower_lengths( + #[values(1, 2, 3)] lower_count: usize, + #[values(0, 1, 3)] lower_len: usize, +) { + let msp = PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable); + let mut lower = vec![PrimitiveArray::new(buffer![0u64; 2], Validity::NonNullable); lower_count]; + lower[lower_count - 1] = PrimitiveArray::new(buffer![0u64; lower_len], Validity::NonNullable); + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + assert!(assemble_decimal(&msp, &lower, dtype).is_err()); +} + +#[rstest] +fn test_assemble_i256_part_order_and_sign_extension( + #[values(false, true)] narrow_msp: bool, + #[values(2, 3)] lower_count: usize, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![3i8, -3], Validity::NonNullable) + } else { + PrimitiveArray::new(buffer![3i64, -3], Validity::NonNullable) + }; + let lower = + [4u64, 1, 2].map(|word| PrimitiveArray::new(buffer![word; 2], Validity::NonNullable)); + let dtype = DecimalDType::new(76, 0); + let actual = assemble_decimal(&msp, &lower[3 - lower_count..], dtype)?; + let low = (1u128 << 64) | 2; + let expected = if lower_count == 2 { + buffer![i256::from_parts(low, 3), i256::from_parts(low, -3)] + } else { + buffer![ + i256::from_parts(low, (3i128 << 64) | 4), + i256::from_parts(low, (-3i128 << 64) | 4), + ] + }; + assert_arrays_eq!( + DecimalArray::new(expected, dtype, Validity::NonNullable), + actual, + &mut ctx + ); Ok(()) } From d93444d6224a6f049148e74ed942c157dd920bec Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 22:41:22 -0400 Subject: [PATCH 04/16] Move decimal array assembly integration to the array layer Keep the splitting and assembly helpers independent of the array changes in the next PR. Signed-off-by: "Matt Katz" --- .../src/decimal_byte_parts/mod.rs | 29 ++++++++++++++----- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index c05252ca45d..4ec2f03995b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -26,11 +26,13 @@ use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::array_slots; +use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::PType; +use vortex_array::match_each_signed_integer_ptype; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -48,7 +50,6 @@ use vortex_error::vortex_panic; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::decimal_byte_parts::limbs::assemble_decimal; use crate::decimal_byte_parts::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. @@ -269,12 +270,26 @@ fn to_canonical_decimal( array: &DecimalBytePartsArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - let msp = array.msp().clone().execute::(ctx)?; - let decimal_dtype = *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"); - Ok(assemble_decimal(&msp, &[], decimal_dtype)?.into_array()) + // TODO(joe): support parts len != 1 + let prim = array.msp().clone().execute::(ctx)?; + // Depending on the decimal type and the min/max of the primitive array we can choose + // the correct buffer size + + Ok(match_each_signed_integer_ptype!(prim.ptype(), |P| { + // SAFETY: The primitive array's buffer is already validated with correct type. + // The decimal dtype matches the array's dtype, and validity is preserved. + unsafe { + DecimalArray::new_unchecked( + prim.to_buffer::

(), + *array + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype"), + prim.validity()?, + ) + } + .into_array() + })) } impl OperationsVTable for DecimalByteParts { From 08ffec1373507929946c0c827fa872f6176d7edd Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 14:13:32 -0400 Subject: [PATCH 05/16] optimize wide split Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/limbs/mod.rs | 188 +++++++++--------- .../src/decimal_byte_parts/limbs/tests.rs | 4 +- 2 files changed, 96 insertions(+), 96 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs index a22749fe79d..10df9ceded7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -102,115 +102,118 @@ pub fn split_decimal(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> VortexRe DecimalType::I64 => DecimalParts::from_msp(decimal.buffer::(), validity), DecimalType::I128 => { let mask = validity.execute_mask(decimal.len(), ctx)?; - let (msp, lower) = split_i128(&decimal.buffer::(), &mask); - DecimalParts::new(msp, [lower], validity) + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i128_to_parts); + DecimalParts::new(msp, lower, validity) } DecimalType::I256 => { let mask = validity.execute_mask(decimal.len(), ctx)?; - let (msp, lower) = split_i256(&decimal.buffer::(), &mask); + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i256_to_parts); DecimalParts::new(msp, lower, validity) } }) } -/// Split each `i128` into an `i64` MSP and an `u64` lower part. +/// Split wide integers into a signed MSP and `N` unsigned lower parts. /// -/// For each valid row, the original value is `msp * 2^64 + lower`. Invalid rows are zeroed. -#[expect( - clippy::cast_possible_truncation, - clippy::cast_sign_loss, - reason = "splitting a wide integer into 64-bit windows truncates by construction" -)] -fn split_i128(values: &Buffer, validity: &Mask) -> (Buffer, Buffer) { - if validity.all_true() { - let mut msp = BufferMut::::with_capacity(values.len()); - let mut lower = BufferMut::::with_capacity(values.len()); - for value in values.iter() { - msp.push((value >> LOWER_PART_BITS) as i64); - lower.push(*value as u64); +/// `to_parts` returns the MSP and lower words in most-significant-first order. +/// It is specialized for each input type: `i128` has one lower word and `i256` +/// has three. Null rows get zeros in every output buffer. +fn split_wide( + values: &Buffer, + validity: &Mask, + to_parts: impl Fn(T) -> (i64, [u64; N]), +) -> (Buffer, [Buffer; N]) { + let len = values.len(); + let mut msp = BufferMut::::with_capacity(len); + let mut lower = std::array::from_fn::<_, N, _>(|_| BufferMut::::with_capacity(len)); + + // Zero out all parts if all null + if validity.all_false() { + msp.push_n(0, len); + for part in &mut lower { + part.push_n(0, len); } - return (msp.freeze(), lower.freeze()); + return (msp.freeze(), lower.map(BufferMut::freeze)); } - // Lower parts are stored as non-nullable arrays, so use zeros at null positions instead - // of copying their garbage values. - let mut msp = BufferMut::::zeroed(values.len()); - let mut lower = BufferMut::::zeroed(values.len()); - - if let Mask::Values(valid) = validity { - let msp = msp.as_mut_slice(); - let lower = lower.as_mut_slice(); - valid.bit_buffer().for_each_set_index(|i| { - let value = values[i]; - msp[i] = (value >> LOWER_PART_BITS) as i64; - lower[i] = value as u64; - }); - } - (msp.freeze(), lower.freeze()) -} + // Allocate without zeroing, then initialize every part of each row together. + let msp_out = &mut msp.spare_capacity_mut()[..len]; + let mut lower_out = lower + .each_mut() + .map(|part| &mut part.spare_capacity_mut()[..len]); -/// Split each `i256` into an `i64` MSP and three `u64` lower parts, ordered most significant -/// first. -/// -/// For each valid row, the original value is -/// -/// `msp * 2^192 + lower[0] * 2^128 + lower[1] * 2^64 + lower[2]`. -/// -/// Invalid rows are zeroed. -fn split_i256( - values: &Buffer, - validity: &Mask, -) -> (Buffer, [Buffer; MAX_LOWER_PARTS]) { - // With no nulls, append every value without zeroing the output buffers first. - if validity.all_true() { - let mut msp = BufferMut::::with_capacity(values.len()); - let mut lower = std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| { - BufferMut::::with_capacity(values.len()) - }); - for value in values.iter() { - let [msp_word, lower_words @ ..] = i256_to_words(*value); - msp.push(msp_word.cast_signed()); - for (part, word) in lower.iter_mut().zip(lower_words) { - part.push(word); + match validity { + Mask::AllTrue(_) => { + for row in 0..len { + let (high, words) = to_parts(values[row]); + msp_out[row].write(high); + for (part, word) in lower_out.iter_mut().zip(words) { + part[row].write(word); + } } } - return (msp.freeze(), lower.map(BufferMut::freeze)); + Mask::Values(validity) => { + // A shorter bitmap would leave output slots uninitialized before set_len. + assert_eq!( + validity.bit_buffer().len(), + len, + "values and validity must have the same length" + ); + for (chunk_index, ((chunk, bits), msp)) in values + .chunks(64) + .zip(validity.bit_buffer().chunks().iter_padded()) + .zip(msp_out.chunks_mut(64)) + .enumerate() + { + for (i, (&value, msp)) in chunk.iter().zip(msp).enumerate() { + let mask = 0u64.wrapping_sub((bits >> i) & 1); + let (high, words) = to_parts(value); + msp.write(high & mask.cast_signed()); + for (part, word) in lower_out.iter_mut().zip(words) { + part[chunk_index * 64 + i].write(word & mask); + } + } + } + } + Mask::AllFalse(_) => unreachable!("AllFalse case addressed above"), } - // Lower parts are stored as non-nullable arrays, so use zeros at null positions instead - // of copying their garbage values. - let mut msp = BufferMut::::zeroed(values.len()); - let mut lower = - std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| BufferMut::::zeroed(values.len())); - - if let Mask::Values(valid) = validity { - let msp = msp.as_mut_slice(); - let mut lower = lower.each_mut().map(BufferMut::as_mut_slice); - valid.bit_buffer().for_each_set_index(|i| { - let [msp_word, lower_words @ ..] = i256_to_words(values[i]); - msp[i] = msp_word.cast_signed(); - for (part, word) in lower.iter_mut().zip(lower_words) { - part[i] = word; - } - }); + // SAFETY: the input and all output slices have len elements. Both branches + // initialize every slot, including null rows and the final partial chunk. + // The bitmap length check prevents the masked iteration from ending early. + unsafe { + msp.set_len(len); + for part in &mut lower { + part.set_len(len); + } } (msp.freeze(), lower.map(BufferMut::freeze)) } -/// Split an `i256` into four `u64` words, most significant first. +/// Extract the high signed word and low unsigned word of an `i128`. #[inline] -const fn i256_to_words(value: i256) -> [u64; 4] { +#[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "each cast preserves a 64-bit window of the original two's complement bits" +)] +const fn i128_to_parts(value: i128) -> (i64, [u64; 1]) { + ((value >> LOWER_PART_BITS) as i64, [value as u64]) +} + +/// Extract the signed MSP and three unsigned lower words of an `i256`. +#[inline] +#[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "each cast preserves a 64-bit window of the original two's complement bits" +)] +const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_LOWER_PARTS]) { let (low, high) = value.to_parts(); - #[expect( - clippy::cast_possible_truncation, - reason = "each cast takes the low 64 bits of a word pair by construction" - )] - [ - (high >> LOWER_PART_BITS) as u64, - high as u64, - (low >> LOWER_PART_BITS) as u64, - low as u64, - ] + ( + (high >> LOWER_PART_BITS) as i64, + [high as u64, (low >> LOWER_PART_BITS) as u64, low as u64], + ) } /// Reassemble primitive arrays that constitute decimal byte parts into a canonical decimal array. @@ -284,18 +287,13 @@ pub(crate) fn assemble_decimal( reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" )] fn assemble_i128(msp: &PrimitiveArray, lower: &[u64]) -> Buffer { - let mut out = BufferMut::::zeroed(msp.len()); + let mut out = BufferMut::::with_capacity(msp.len()); match_each_signed_integer_ptype!(msp.ptype(), |P| { - for ((slot, value), part) in out - .as_mut_slice() - .iter_mut() - .zip(msp.as_slice::

()) - .zip(lower) - { + out.extend_trusted(msp.as_slice::

().iter().zip(lower).map(|(value, part)| { // Sign-extend the MSP, then shift it into the high 64 bits. The unsigned // lower part fills the low 64 bits. - *slot = (i128::from(i64::from(*value)) << LOWER_PART_BITS) | i128::from(*part); - } + (i128::from(i64::from(*value)) << LOWER_PART_BITS) | i128::from(*part) + })); }); out.freeze() } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs index 57856a4f2c4..3e3de06c44e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs @@ -20,10 +20,12 @@ use super::*; #[case::all_valid(Validity::AllValid)] #[case::all_null(Validity::AllInvalid)] #[case::mixed(Validity::from_iter((0..263).map(|i| i % 3 != 1)))] +#[case::sparse(Validity::from_iter((0..263).map(|i| i % 16 == 0)))] +#[case::null_prefix_and_suffix(Validity::from_iter((0..263).map(|i| (67..196).contains(&i))))] fn test_split_zeroes_null_words( #[case] validity: Validity, #[values(false, true)] wide_256: bool, - #[values(0, 1, 257)] len: usize, + #[values(0, 1, 63, 64, 65, 257)] len: usize, ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let decimal = if wide_256 { From 1f2143a8a896fe6bd73f40c34c26138b9b718765 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 14:24:00 -0400 Subject: [PATCH 06/16] Benchmark decimal splitting and assembly in the parts layer Measure split_decimal and assemble_decimal directly across storage widths and input sizes, with fixtures outside the timed calls. Cover all-valid, all-null, random-null, and clustered-null split inputs. Signed-off-by: "Matt Katz" --- Cargo.lock | 2 + encodings/decimal-byte-parts/Cargo.toml | 10 ++++ .../decimal-byte-parts/benches/common/mod.rs | 51 ++++++++++++++++ .../benches/dbp_assemble.rs | 48 +++++++++++++++ .../decimal-byte-parts/benches/dbp_split.rs | 59 +++++++++++++++++++ .../src/decimal_byte_parts/limbs/mod.rs | 2 +- .../src/decimal_byte_parts/mod.rs | 5 ++ 7 files changed, 176 insertions(+), 1 deletion(-) create mode 100644 encodings/decimal-byte-parts/benches/common/mod.rs create mode 100644 encodings/decimal-byte-parts/benches/dbp_assemble.rs create mode 100644 encodings/decimal-byte-parts/benches/dbp_split.rs diff --git a/Cargo.lock b/Cargo.lock index f7d3103c9d5..04baecf6504 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10921,8 +10921,10 @@ dependencies = [ name = "vortex-decimal-byte-parts" version = "0.1.0" dependencies = [ + "codspeed-divan-compat", "num-traits", "prost 0.14.4", + "rand 0.10.2", "rstest", "vortex-array", "vortex-buffer", diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 4934ec4fa27..9f2e387a4da 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -26,5 +26,15 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } + +[[bench]] +name = "dbp_assemble" +harness = false + +[[bench]] +name = "dbp_split" +harness = false diff --git a/encodings/decimal-byte-parts/benches/common/mod.rs b/encodings/decimal-byte-parts/benches/common/mod.rs new file mode 100644 index 00000000000..eed00b98c1a --- /dev/null +++ b/encodings/decimal-byte-parts/benches/common/mod.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared decimal inputs for splitting and assembly benchmarks. + +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::vortex_panic; + +pub(super) fn cases() -> Vec<(DecimalType, usize)> { + [DecimalType::I64, DecimalType::I128, DecimalType::I256] + .into_iter() + .flat_map(|values_type| [1_024, 8_192, 65_536].map(|len| (values_type, len))) + .collect() +} + +pub(super) fn decimal_array( + values_type: DecimalType, + len: usize, + validity: Validity, +) -> DecimalArray { + let mut rng = StdRng::seed_from_u64(42); + + macro_rules! decimal { + ($T:ty, $precision:literal) => {{ + let max = <$T>::pow(10, $precision) - 1; + let values: Buffer<$T> = (0..len).map(|_| rng.random_range(-max..=max)).collect(); + DecimalArray::new(values, DecimalDType::new($precision, 2), validity) + }}; + } + + match values_type { + DecimalType::I64 => decimal!(i64, 18), + DecimalType::I128 => decimal!(i128, 38), + DecimalType::I256 => { + // Keep the magnitude below 10^76 while exercising all four signed/unsigned words. + let values: Buffer = (0..len) + .map(|_| i256::from_parts(rng.random(), rng.random::() >> 4)) + .collect(); + DecimalArray::new(values, DecimalDType::new(76, 2), validity) + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } +} diff --git a/encodings/decimal-byte-parts/benches/dbp_assemble.rs b/encodings/decimal-byte-parts/benches/dbp_assemble.rs new file mode 100644 index 00000000000..74a66a133f5 --- /dev/null +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reassembling primitive decimal parts across storage widths and lengths. + +mod common; + +use divan::Bencher; +use divan::black_box; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DecimalType; +use vortex_array::validity::Validity; +use vortex_decimal_byte_parts::_benchmarking::assemble_decimal; +use vortex_decimal_byte_parts::split_decimal; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::common::cases; +use crate::common::decimal_array; + +fn main() { + divan::main(); +} + +#[divan::bench(args = cases())] +fn assemble(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let decimal = decimal_array(values_type, len, Validity::NonNullable); + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal, &mut ctx).vortex_expect("split benchmark input"); + let msp = parts + .msp + .execute::(&mut ctx) + .vortex_expect("execute benchmark MSP"); + let lower_parts = parts + .lower_parts + .into_iter() + .map(|part| part.execute::(&mut ctx)) + .collect::>>() + .vortex_expect("execute benchmark lower parts"); + let decimal_dtype = decimal.decimal_dtype(); + + bencher.bench(|| { + assemble_decimal(black_box(&msp), black_box(&lower_parts), decimal_dtype) + .vortex_expect("assemble decimal byte parts") + }); +} diff --git a/encodings/decimal-byte-parts/benches/dbp_split.rs b/encodings/decimal-byte-parts/benches/dbp_split.rs new file mode 100644 index 00000000000..8258c50931b --- /dev/null +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -0,0 +1,59 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splitting decimal arrays across storage widths, lengths, and validity paths. + +mod common; + +use divan::Bencher; +use divan::black_box; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::dtype::DecimalType; +use vortex_array::validity::Validity; +use vortex_decimal_byte_parts::split_decimal; +use vortex_error::VortexExpect; + +use crate::common::cases; +use crate::common::decimal_array; + +fn main() { + divan::main(); +} + +#[divan::bench(args = cases())] +fn all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Validity::AllValid); +} + +#[divan::bench(args = cases())] +fn all_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Validity::AllInvalid); +} + +#[divan::bench(args = cases())] +fn mixed_nulls(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let mut rng = StdRng::seed_from_u64(42); + let validity = Validity::from_iter((0..len).map(|_| rng.random_bool(0.5))); + bench_split(bencher, values_type, len, validity); +} + +#[divan::bench(args = cases())] +fn clustered_nulls(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + const CLUSTER_LEN: usize = 256; + let validity = Validity::from_iter((0..len).map(|i| (i / CLUSTER_LEN).is_multiple_of(2))); + bench_split(bencher, values_type, len, validity); +} + +fn bench_split(bencher: Bencher, values_type: DecimalType, len: usize, validity: Validity) { + let decimal = decimal_array(values_type, len, validity); + let session = array_session(); + bencher + .with_inputs(|| session.create_execution_ctx()) + .bench_refs(|ctx| { + split_decimal(black_box(&decimal), ctx).vortex_expect("split decimal array") + }); +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs index 10df9ceded7..a7f351e6e05 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -228,7 +228,7 @@ const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_LOWER_PARTS]) { /// /// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity /// cannot be derived. -pub(crate) fn assemble_decimal( +pub fn assemble_decimal( msp: &PrimitiveArray, lower_parts: &[PrimitiveArray], decimal_dtype: DecimalDType, diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 4ec2f03995b..a7f63bfa082 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -16,6 +16,11 @@ pub use limbs::split_decimal; mod rules; mod slice; +#[doc(hidden)] +pub mod _benchmarking { + pub use super::limbs::assemble_decimal; +} + use prost::Message as _; use vortex_array::ArrayEq; use vortex_array::ArrayHash; From 8ef08a6b3666595bee9bcf739ed00426a5fac77e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 16:27:02 -0400 Subject: [PATCH 07/16] Support multi-part decimal arrays and kernels Represent wide decimals with a signed high part and up to three unsigned low parts. Add validation, execution, kernel support, property tests, and assembly benchmarks while keeping serialization on the frozen format. Signed-off-by: "Matt Katz" --- .gitignore | 2 + Cargo.lock | 84 ++- Cargo.toml | 1 + encodings/decimal-byte-parts/Cargo.toml | 1 + .../src/decimal_byte_parts/compute/cast.rs | 17 +- .../src/decimal_byte_parts/compute/compare.rs | 47 ++ .../src/decimal_byte_parts/compute/filter.rs | 44 +- .../decimal_byte_parts/compute/is_constant.rs | 27 +- .../src/decimal_byte_parts/compute/kernel.rs | 8 - .../src/decimal_byte_parts/compute/mask.rs | 16 +- .../src/decimal_byte_parts/compute/mod.rs | 33 + .../src/decimal_byte_parts/compute/take.rs | 107 +++- .../src/decimal_byte_parts/limbs/mod.rs | 28 +- .../src/decimal_byte_parts/mod.rs | 606 +++++++++++++++--- .../src/decimal_byte_parts/rules.rs | 42 +- .../src/decimal_byte_parts/slice.rs | 14 +- .../src/decimal_byte_parts/testing.rs | 53 ++ encodings/decimal-byte-parts/tests/props.rs | 198 ++++++ vortex-btrblocks/src/trace_tests.rs | 5 +- .../kernel/encodings/decimal_byte_parts.rs | 7 + 20 files changed, 1139 insertions(+), 201 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs create mode 100644 encodings/decimal-byte-parts/tests/props.rs diff --git a/.gitignore b/.gitignore index f9613807332..6db14ce5f6a 100644 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,8 @@ coverage.xml *.cover *.py,cover .hypothesis/ +# hegeltest's example database, the Rust equivalent of .hypothesis/ +.hegel/ .pytest_cache/ cover/ diff --git a/Cargo.lock b/Cargo.lock index 04baecf6504..639e06bf078 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1120,7 +1120,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f21ff1fc630079352bae9b024f85519bf1f641cf7f326623f4c0b59f7ea834fd" dependencies = [ "compact_str", - "miniz_oxide", + "miniz_oxide 0.9.1", "thiserror 2.0.20", ] @@ -2257,6 +2257,25 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "dashu-base" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "993b95dc1b248e3f5747dcb017a41d6e75853a2e5ee4504f7d537c5b8dffdae4" + +[[package]] +name = "dashu-int" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49c05a0d5cb0b39fcc87c46432fdac24b90dce239857c7f6b798be4ffc3c42c6" +dependencies = [ + "cfg-if", + "dashu-base", + "num-modular", + "rustversion", + "static_assertions", +] + [[package]] name = "datafusion" version = "54.1.0" @@ -4098,7 +4117,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.9.1", "zlib-rs", ] @@ -4699,6 +4718,51 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hegeltest" +version = "0.28.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "100bcd6ef825f5b6a60e2f55c05bb626ebf254dd8a09d16e006c4bb7883e7f1c" +dependencies = [ + "crc32fast", + "dashu-int", + "hegeltest-c", + "hegeltest-macros", + "miniz_oxide 0.8.9", + "parking_lot", + "paste", + "rand 0.10.2", + "rustc-hash", + "tempfile", +] + +[[package]] +name = "hegeltest-c" +version = "0.30.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a672fd53360ca4122c1a145a85e8fef835508d7b40eb9de43498978e796c54b" +dependencies = [ + "dashu-int", + "hashbrown 0.17.1", + "libm", + "miniz_oxide 0.8.9", + "parking_lot", + "rand 0.10.2", + "rustc-hash", + "tempfile", +] + +[[package]] +name = "hegeltest-macros" +version = "0.28.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba792d78fa3740a7c1627085c34618b998b8aa0f63625721235234f525aad1aa" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "hermit-abi" version = "0.5.3" @@ -6555,6 +6619,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", +] + [[package]] name = "miniz_oxide" version = "0.9.1" @@ -6862,6 +6935,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-modular" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc41a1374056e9672221567958a66c16be12d0e2c1b408761e14d901c237d5e0" + [[package]] name = "num-rational" version = "0.4.2" @@ -10922,6 +11001,7 @@ name = "vortex-decimal-byte-parts" version = "0.1.0" dependencies = [ "codspeed-divan-compat", + "hegeltest", "num-traits", "prost 0.14.4", "rand 0.10.2", diff --git a/Cargo.toml b/Cargo.toml index 6a9afca5e25..1bf8177ec0e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -177,6 +177,7 @@ glob = "0.3.2" goldenfile = "1" half = { version = "2.7.1", features = ["std", "num-traits"] } hashbrown = "0.17.1" +hegeltest = "0.28.7" http = "1.5.0" humansize = "2.1.3" indicatif = "0.18.0" diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 9f2e387a4da..e9ea8569af1 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -27,6 +27,7 @@ vortex-session = { workspace = true } [dev-dependencies] divan = { workspace = true } +hegeltest = { workspace = true } rand = { workspace = true } rstest = { workspace = true } vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index 5ae1bf0101e..7b949fcd695 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use crate::DecimalByteParts; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::with_msp; impl CastReduce for DecimalByteParts { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { @@ -29,9 +30,7 @@ impl CastReduce for DecimalByteParts { .msp() .cast(array.msp().dtype().with_nullability(*target_nullability))?; - Ok(Some( - DecimalByteParts::try_new(new_msp, *target_decimal)?.into_array(), - )) + with_msp(array, new_msp, *target_decimal).map(|a| Some(a.into_array())) } } @@ -49,10 +48,14 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use crate::DecimalByteParts; use crate::DecimalBytePartsArray; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_of; + use crate::decimal_byte_parts::testing::i256_parts; #[test] fn test_cast_decimal_byte_parts_nullability() { @@ -117,6 +120,14 @@ mod tests { buffer![-100i32, -200, 300, -400, 500].into_array(), DecimalDType::new(10, 2), ).unwrap())] + #[case::one_lower_part(i128_parts( + vec![1i128 << 70, -(1i128 << 70), 5, (1i128 << 64) - 1, 0], + Validity::NonNullable, + ))] + #[case::three_lower_parts(i256_parts( + vec![i256_of(1, 0), i256_of(-1, 5), i256_of(0, u128::MAX)], + Validity::NonNullable, + ))] fn test_cast_decimal_byte_parts_conformance(#[case] array: DecimalBytePartsArray) { test_cast_conformance( &array.into_array(), diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index 3044bd6e605..fe4d69801a3 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -39,6 +39,12 @@ impl CompareKernel for DecimalByteParts { return Ok(None); }; + // The MSP alone only determines the ordering when it holds the whole value. With + // lower parts present, fall back to comparing the canonical decimal. + if !lhs.lower_parts().is_empty() { + return Ok(None); + } + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); let scalar_type = lhs.msp().dtype().with_nullability(nullability); @@ -158,10 +164,12 @@ mod tests { use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::i128_parts; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session(); @@ -220,6 +228,45 @@ mod tests { Ok(()) } + #[test] + fn compare_decimal_const_with_lower_parts() -> VortexResult<()> { + // The MSP-only pushdown is invalid once lower parts carry part of the value, so this + // must fall back to the canonical comparison rather than compare MSPs. + let values = vec![1i128 << 70, (1i128 << 70) + 1, 5, -(1i128 << 70)]; + let lhs = i128_parts(values.clone(), Validity::NonNullable).into_array(); + let decimal_dtype = *lhs + .dtype() + .as_decimal_opt() + .vortex_expect("decimal byte parts array"); + + let pivot = (1i128 << 70) + 1; + let rhs = ConstantArray::new( + Scalar::decimal( + DecimalValue::I128(pivot), + decimal_dtype, + Nullability::NonNullable, + ), + lhs.len(), + ) + .into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + for (operator, predicate) in [ + (Operator::Eq, (|v, p| v == p) as fn(i128, i128) -> bool), + (Operator::NotEq, |v, p| v != p), + (Operator::Lt, |v, p| v < p), + (Operator::Lte, |v, p| v <= p), + (Operator::Gt, |v, p| v > p), + (Operator::Gte, |v, p| v >= p), + ] { + let res = lhs.clone().binary(rhs.clone(), operator)?; + let expected = + BoolArray::from_iter(values.iter().map(|v| predicate(*v, pivot))).into_array(); + assert_arrays_eq!(res, expected, &mut ctx); + } + Ok(()) + } + #[test] fn compare_decimal_const_unconvertible_comparison() { let decimal_dtype = DecimalDType::new(40, 2); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index a47a6ed846b..e4fb03a5ca0 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs @@ -5,22 +5,15 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::filter::FilterReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::map_parts; + impl FilterReduce for DecimalByteParts { fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { - DecimalByteParts::try_new( - array.msp().filter(mask.clone())?, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - ) - .map(|d| Some(d.into_array())) + map_parts(array, |part| part.filter(mask.clone())).map(|d| Some(d.into_array())) } } @@ -32,9 +25,13 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::dtype::DecimalDType; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_of; + use crate::decimal_byte_parts::testing::i256_parts; #[test] fn test_filter_decimal_byte_parts() { @@ -59,4 +56,31 @@ mod test { &mut array_session().create_execution_ctx(), ); } + + #[test] + fn test_filter_decimal_byte_parts_with_lower_parts() { + let array = i128_parts( + vec![1i128 << 70, -(1i128 << 70), 5, (1i128 << 64) - 1, 0], + Validity::NonNullable, + ); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + + let array = i256_parts( + vec![ + i256_of(1, 0), + i256_of(-1, 5), + i256_of(0, u128::MAX), + i256_of(1 << 64, 7), + i256_of(0, 0), + ], + Validity::from_iter([true, false, true, true, false]), + ); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs index 065bc5e0051..3fe59111f6e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors 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; @@ -15,7 +16,9 @@ use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; /// DecimalByteParts-specific is_constant kernel. /// -/// Delegates to checking if the MSP (most significant part) is constant. +/// Delegates to checking that every part is constant: the MSP (most significant part) plus +/// each lower part. An all-null array is constant regardless of the bits its lower parts +/// hold in null slots. #[derive(Debug)] pub(crate) struct DecimalBytePartsIsConstantKernel; @@ -34,7 +37,27 @@ impl DynAggregateKernel for DecimalBytePartsIsConstantKernel { return Ok(None); }; - let result = is_constant(array.msp(), ctx)?; + let result = is_constant_parts(array, ctx)?; Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } + +fn is_constant_parts( + array: ArrayView<'_, DecimalByteParts>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !is_constant(array.msp(), ctx)? { + return Ok(false); + } + // Null slots hold undefined bits in the lower parts, so they cannot make a constant + // (all-null) array non-constant. + if array.array().all_invalid(ctx)? { + return Ok(true); + } + for part in array.lower_parts().iter() { + if !is_constant(part, ctx)? { + return Ok(false); + } + } + Ok(true) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 5e8d28e3526..cb71ba7880c 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -1,9 +1,6 @@ // 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::dict::TakeExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; @@ -19,9 +16,4 @@ pub(crate) fn initialize(session: &VortexSession) { DecimalByteParts, CompareExecuteAdaptor(DecimalByteParts), ); - kernels.register_execute_parent_kernel( - Dict.id(), - DecimalByteParts, - TakeExecuteAdaptor(DecimalByteParts), - ); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index e7dc95af84f..2eea785794b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs @@ -6,24 +6,18 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; use vortex_array::scalar_fn::fns::mask::MaskReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::decimal_dtype; +use crate::decimal_byte_parts::with_msp; impl MaskReduce for DecimalByteParts { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { + // Validity lives in the MSP, so only that part needs masking: the lower parts hold + // undefined bits in null slots, which is exactly what a masked-out row is. let masked_msp = MaskExpr::try_new(array.msp().clone(), mask.clone())?.into_array(); - Ok(Some( - DecimalByteParts::try_new( - masked_msp, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - )? - .into_array(), - )) + with_msp(array, masked_msp, decimal_dtype(array)).map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 6c2d0dabb31..844468545cf 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -19,10 +19,36 @@ mod tests { use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::i256; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use crate::DecimalByteParts; use crate::DecimalBytePartsArray; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_of; + use crate::decimal_byte_parts::testing::i256_parts; + + /// Values needing more than 64 bits, so the encoding carries lower parts. + fn wide_i128() -> Vec { + vec![ + 1 << 70, + -(1 << 70), + (1 << 64) - 1, + 0, + 99_999_999_999_999_999_999_999_999_999_999_999_999, + ] + } + + fn wide_i256() -> Vec { + vec![ + i256_of(1, 0), + i256_of(-1, 0), + i256_of(0, u128::MAX), + i256_of(1 << 64, 7), + i256_of(0, 0), + ] + } #[rstest] // Basic decimal byte parts arrays @@ -70,6 +96,11 @@ mod tests { PrimitiveArray::from_iter((0..2000i64).map(|i| i * 1000000)).into_array(), DecimalDType::new(19, 6) ).unwrap())] + // Wide decimals carrying lower parts + #[case::decimal_i128_one_lower_part(i128_parts(wide_i128(), Validity::NonNullable))] + #[case::decimal_i128_nullable(i128_parts(wide_i128(), Validity::from_iter([true, false, true, true, false])))] + #[case::decimal_i256_three_lower_parts(i256_parts(wide_i256(), Validity::NonNullable))] + #[case::decimal_i256_nullable(i256_parts(wide_i256(), Validity::from_iter([false, true, true, false, true])))] fn test_decimal_byte_parts_consistency(#[case] array: DecimalBytePartsArray) { let ctx = &mut array_session().create_execution_ctx(); @@ -89,6 +120,8 @@ mod tests { buffer![-100i32, -200, 300, -400, 500].into_array(), DecimalDType::new(10, 2) ).unwrap())] + #[case::decimal_i128_one_lower_part(i128_parts(wide_i128(), Validity::NonNullable))] + #[case::decimal_i256_three_lower_parts(i256_parts(wide_i256(), Validity::NonNullable))] fn test_decimal_byte_parts_binary_numeric(#[case] array: DecimalBytePartsArray) { let ctx = &mut array_session().create_execution_ctx(); test_binary_numeric_array(&array.into_array(), ctx); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index 7a18f7bf91b..578834635b8 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs @@ -3,28 +3,101 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; -use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::arrays::dict::TakeExecute; -use vortex_error::VortexExpect; +use vortex_array::arrays::dict::TakeReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::map_parts; -impl TakeExecute for DecimalByteParts { - fn take( - array: ArrayView<'_, Self>, - indices: &ArrayRef, - _ctx: &mut ExecutionCtx, - ) -> VortexResult> { - DecimalByteParts::try_new( - array.msp().take(indices.clone())?, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - ) - .map(|a| Some(a.into_array())) +impl TakeReduce for DecimalByteParts { + /// Taking wraps each part in a `Dict` without reading any buffer, so it reduces rather + /// than executes. + fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + // Taking with nullable indices makes every taken part nullable, but lower parts must + // stay non-nullable `u64` — validity belongs to the MSP alone. Fall back to the + // canonical path rather than rebuilding parts we would have to strip nullability from. + if indices.dtype().is_nullable() && !array.lower_parts().is_empty() { + return Ok(None); + } + + map_parts(array, |part| part.take(indices.clone())).map(|a| Some(a.into_array())) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DecimalDType; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::encode; + use crate::decimal_byte_parts::testing::i256_of; + + /// Taking pushes down into the parts during optimization, with no execution context in + /// play: `ArrayRef::take` wraps the array in a `Dict` and optimizes, and the reduce rule + /// must rewrite that into a `DecimalByteParts` of taken parts. + #[test] + fn take_pushes_down_without_executing() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + + let decimal = DecimalArray::new( + Buffer::from(vec![1i128 << 70, 2, 3]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let indices = buffer![0u64, 2].into_array(); + let taken = encode(&decimal)?.into_array().take(indices)?; + + assert!( + taken.is::(), + "expected the take to reduce into the encoding, got {}", + taken.encoding_id() + ); + Ok(()) + } + + /// Taking with nullable indices must still round-trip the wide values, including the + /// null row, on arrays that carry lower parts. + #[rstest] + #[case::one_lower_part(DecimalArray::new( + Buffer::from(vec![1i128 << 70, 2, 3]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ))] + #[case::three_lower_parts(DecimalArray::new( + Buffer::from(vec![i256_of(1, 1 << 70), i256_of(0, 2), i256_of(0, 3)]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ))] + fn take_with_nullable_indices(#[case] decimal: DecimalArray) -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let indices = PrimitiveArray::from_option_iter([Some(0u64), None, Some(2u64)]).into_array(); + let expected = decimal + .clone() + .into_array() + .take(indices.clone())? + .execute::(&mut ctx)?; + + let taken = encode(&decimal)?.into_array().take(indices)?; + let actual = taken.execute::(&mut ctx)?; + + assert_arrays_eq!(expected, actual, &mut ctx); + Ok(()) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs index a7f351e6e05..e3485463066 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -248,10 +248,6 @@ pub fn assemble_decimal( let lower: Vec<&[u64]> = lower_parts .iter() .map(|part| { - vortex_ensure!( - part.dtype() == &LOWER_PART_DTYPE, - "lower part must be non-nullable u64" - ); let part = part.as_slice::(); vortex_ensure!( part.len() == len, @@ -279,6 +275,30 @@ pub fn assemble_decimal( }) } +/// Combine a single row's parts into an `i128`. +#[inline] +pub(crate) fn combine_i128(msp: i64, lower: impl IntoIterator) -> i128 { + lower.into_iter().fold(i128::from(msp), |acc, part| { + (acc << LOWER_PART_BITS) | i128::from(part) + }) +} + +/// Combine a signed MSP and two or three lower parts into an `i256`. +#[inline] +pub(crate) fn combine_i256(msp: i64, lower: impl ExactSizeIterator) -> i256 { + let count = lower.len(); + let mut high = i128::from(msp); + let mut low = 0u128; + for (index, part) in lower.enumerate() { + if count == 3 && index == 0 { + high = (high << LOWER_PART_BITS) | i128::from(part); + } else { + low = (low << LOWER_PART_BITS) | u128::from(part); + } + } + i256::from_parts(low, high) +} + /// Reassemble a signed MSP and one `u64` lower part into `i128` values. /// /// For each row, the result is `msp * 2^64 + lower`. diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index a7f63bfa082..4a104e82c2b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -10,12 +10,14 @@ use vortex_array::ArrayParts; use vortex_array::ArrayView; pub(crate) mod compute; mod limbs; -pub use limbs::DecimalParts; -pub use limbs::MAX_LOWER_PARTS; -pub use limbs::split_decimal; mod rules; mod slice; +#[cfg(test)] +pub(crate) mod testing; +pub use limbs::DecimalParts; +pub use limbs::MAX_LOWER_PARTS; +pub use limbs::split_decimal; #[doc(hidden)] pub mod _benchmarking { pub use super::limbs::assemble_decimal; @@ -26,23 +28,22 @@ use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayRef; +use vortex_array::ArraySlots; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::array_slots; -use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; use vortex_array::dtype::PType; -use vortex_array::match_each_signed_integer_ptype; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityChild; @@ -51,10 +52,15 @@ 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::decimal_byte_parts::limbs::LOWER_PART_DTYPE; +use crate::decimal_byte_parts::limbs::assemble_decimal; +use crate::decimal_byte_parts::limbs::combine_i128; +use crate::decimal_byte_parts::limbs::combine_i256; use crate::decimal_byte_parts::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. @@ -78,6 +84,70 @@ pub struct DecimalBytesPartsMetadata { lower_part_count: u32, } +impl DecimalBytesPartsMetadata { + fn from_array(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { + Ok(Self { + zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, + lower_part_count: u32::try_from(array.lower_parts().len()) + .map_err(|_| vortex_err!("lower part count exceeds u32"))?, + }) + } + + fn into_array_parts( + self, + dtype: &DType, + len: usize, + children: &dyn ArrayChildren, + ) -> VortexResult> { + vortex_ensure!( + dtype.as_decimal_opt().is_some(), + "decoding decimal but given non decimal dtype {dtype}" + ); + + let encoded_dtype = DType::Primitive(self.zeroth_child_ptype(), dtype.nullability()); + + let lower_part_count = self.lower_part_count()?; + vortex_ensure!( + children.len() == DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, + "expected {} children, got {}", + DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, + children.len() + ); + + let msp = children.get(DecimalBytePartsSlots::MSP, &encoded_dtype, len)?; + + let mut slots = ArraySlots::with_capacity(children.len()); + slots.push(Some(msp)); + for idx in 0..lower_part_count { + slots.push(Some(children.get( + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx, + &LOWER_PART_DTYPE, + len, + )?)); + } + + Ok( + ArrayParts::new(DecimalByteParts, dtype.clone(), len, DecimalBytePartsData) + .with_slots(slots), + ) + } + + /// The number of lower parts encoded in this array. + /// + /// # Errors + /// + /// Returns an error if the count exceeds [`MAX_LOWER_PARTS`]. + fn lower_part_count(&self) -> VortexResult { + let count = usize::try_from(self.lower_part_count) + .map_err(|_| vortex_err!("lower part count {} out of range", self.lower_part_count))?; + vortex_ensure!( + count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {count}" + ); + Ok(count) + } +} + impl VTable for DecimalByteParts { type TypedArrayData = DecimalBytePartsData; @@ -99,8 +169,14 @@ impl VTable for DecimalByteParts { let Some(decimal_dtype) = dtype.as_decimal_opt() else { vortex_bail!("expected decimal dtype, got {}", dtype) }; - let msp = DecimalBytePartsSlotsView::from_slots(slots).msp; - DecimalBytePartsData::validate(msp, *decimal_dtype, dtype, len) + let slots = DecimalBytePartsSlotsView::from_slots(slots); + DecimalBytePartsData::validate( + slots.msp, + slots.lower_parts.iter(), + *decimal_dtype, + dtype, + len, + ) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -127,12 +203,12 @@ impl VTable for DecimalByteParts { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { + vortex_ensure!( + array.lower_parts().is_empty(), + "serializing DecimalByteParts with lower parts is not supported" + ); Ok(Some( - DecimalBytesPartsMetadata { - zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, - lower_part_count: 0, - } - .encode_to_vec(), + DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), )) } @@ -146,26 +222,15 @@ impl VTable for DecimalByteParts { _session: &VortexSession, ) -> VortexResult> { let metadata = DecimalBytesPartsMetadata::decode(metadata)?; - let Some(decimal_dtype) = dtype.as_decimal_opt() else { - vortex_bail!("decoding decimal but given non decimal dtype {}", dtype) - }; - - let encoded_dtype = DType::Primitive(metadata.zeroth_child_ptype(), dtype.nullability()); - - let msp = children.get(0, &encoded_dtype, len)?; - - assert_eq!( - metadata.lower_part_count, 0, - "lower_part_count > 0 not currently supported" + vortex_ensure!( + metadata.lower_part_count()? == 0, + "vortex.decimal_byte_parts must not carry lower parts" ); - - let slots = smallvec![Some(msp.clone())]; - let data = DecimalBytePartsData::try_new(msp.dtype(), msp.len(), *decimal_dtype)?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + metadata.into_array_parts(dtype, len, children) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - DecimalBytePartsSlots::NAMES[idx].to_string() + DecimalBytePartsSlots::slot_name(idx) } fn reduce_parent( @@ -186,20 +251,21 @@ pub struct DecimalBytePartsSlots { /// The most significant parts of the decimal values. #[slot(0)] pub msp: ArrayRef, + /// The remaining 64-bit windows of the decimal values, most significant first. + #[slot(1..)] + pub lower_parts: Vec, } /// This array encodes decimals as between 1-4 columns of primitive typed children. -/// The most significant part (msp) sorting the most significant decimal bits. +/// The most significant part (msp) storing the most significant decimal bits. /// This array must be signed and is nullable iff the decimal is nullable. +/// Every lower part is a non-nullable `u64` holding a raw 64-bit window of the value. +/// +/// e.g. for a decimal i128 \[ 127..64 | 63..0 \] msp = 127..64 and lower_part\[0\] = 63..0 /// -/// e.g. for a decimal i128 \[ 127..64 | 64..0 \] msp = 127..64 and lower_part\[0\] = 64..0 +/// All parts live in slots, so the array carries no additional data. #[derive(Clone, Debug)] -pub struct DecimalBytePartsData { - // NOTE: the lower_parts is currently unused, we reserve this field so that it is properly - // read/written during serde, but provide no constructor to initialize this to anything - // other than the empty Vec. - _lower_parts: Vec, -} +pub struct DecimalBytePartsData; impl Display for DecimalBytePartsData { fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { @@ -207,13 +273,17 @@ impl Display for DecimalBytePartsData { } } -pub struct DecimalBytePartsDataParts { - pub msp: ArrayRef, -} - impl DecimalBytePartsData { - pub fn validate( + /// Validate the parts of a [`DecimalBytePartsArray`]. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array of length `len`, if `dtype` + /// does not match the MSP's nullability, if there are more than [`MAX_LOWER_PARTS`] + /// lower parts, or if any lower part is not a non-nullable `u64` array of length `len`. + pub fn validate<'a>( msp: &ArrayRef, + lower_parts: impl ExactSizeIterator, decimal_dtype: DecimalDType, dtype: &DType, len: usize, @@ -228,24 +298,26 @@ impl DecimalBytePartsData { "expected dtype {expected_dtype}, got {dtype}" ); vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); - Ok(()) - } - pub(crate) fn try_new( - msp_dtype: &DType, - msp_len: usize, - decimal_dtype: DecimalDType, - ) -> VortexResult { - let expected_dtype = DType::Decimal(decimal_dtype, msp_dtype.nullability()); + let lower_part_count = lower_parts.len(); + // Physical storage may be wider than the declared precision, as for DecimalArray. vortex_ensure!( - msp_dtype.is_signed_int(), - "decimal bytes parts, first part must be a signed array" + lower_part_count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" ); - let _ = msp_len; - drop(expected_dtype); - Ok(Self { - _lower_parts: Vec::new(), - }) + for (idx, part) in lower_parts.enumerate() { + vortex_ensure!( + part.dtype() == &LOWER_PART_DTYPE, + "lower part {idx} must have dtype {LOWER_PART_DTYPE}, got {}", + part.dtype() + ); + vortex_ensure!( + part.len() == len, + "lower part {idx} has len {}, expected {len}", + part.len() + ); + } + Ok(()) } } @@ -254,47 +326,108 @@ pub struct DecimalByteParts; impl DecimalByteParts { /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array. pub fn try_new( msp: ArrayRef, decimal_dtype: DecimalDType, ) -> VortexResult { + Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + } + + /// Construct a new [`DecimalBytePartsArray`] from an MSP array, its lower parts, and a + /// decimal dtype. + /// + /// Lower parts are ordered most significant first and must each be a non-nullable `u64` + /// array of the same length as the MSP. See [`split_decimal`] for producing them from a + /// canonical decimal array. + /// + /// # Errors + /// + /// Returns an error if the parts do not describe a valid decimal, see + /// [`DecimalBytePartsData::validate`]. + pub fn try_new_with_lower_parts( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> VortexResult { + // Lower parts are supported in memory; the frozen serializer still rejects them. let len = msp.len(); let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); - let slots = smallvec![Some(msp.clone())]; - let data = DecimalBytePartsData::try_new(msp.dtype(), msp.len(), decimal_dtype)?; - Ok(unsafe { - Array::from_parts_unchecked( - ArrayParts::new(DecimalByteParts, dtype, len, data).with_slots(slots), - ) - }) + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + Array::try_from_parts( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + ) + } +} + +/// The decimal storage type this array canonicalizes to. +fn values_type(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { + match array.lower_parts().len() { + 0 => DecimalType::try_from(array.msp().dtype().as_ptype()), + 1 => Ok(DecimalType::I128), + 2 | 3 => Ok(DecimalType::I256), + count => vortex_bail!("at most {MAX_LOWER_PARTS} lower parts are supported, got {count}"), } } +/// The decimal dtype this array carries. +/// +/// Guaranteed to be a decimal by construction: [`DecimalBytePartsData::validate`] rejects +/// every other dtype. +pub(crate) fn decimal_dtype(array: ArrayView<'_, DecimalByteParts>) -> DecimalDType { + *array + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype") +} + +/// Rebuild the array by applying `f` to the MSP and to every lower part, in slot order. +/// +/// Part-wise operations must touch every part. Going through this rather than calling +/// [`DecimalByteParts::try_new_with_lower_parts`] directly makes dropping a lower part — +/// which silently corrupts wide values — unrepresentable. +pub(crate) fn map_parts( + array: ArrayView<'_, DecimalByteParts>, + mut f: impl FnMut(&ArrayRef) -> VortexResult, +) -> VortexResult { + let msp = f(array.msp())?; + let lower_parts = array + .lower_parts() + .iter() + .map(&mut f) + .collect::>>()?; + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal_dtype(array)) +} + +/// Rebuild the array with a replacement MSP, keeping its lower parts untouched. +/// +/// Only valid for operations that cannot change a row's magnitude bits — a nullability cast +/// or a mask — since the lower parts keep whatever bits they held. That is sound because +/// validity lives in the MSP alone, so lower-part bits in a null row are already undefined. +pub(crate) fn with_msp( + array: ArrayView<'_, DecimalByteParts>, + msp: ArrayRef, + decimal_dtype: DecimalDType, +) -> VortexResult { + DecimalByteParts::try_new_with_lower_parts(msp, array.lower_parts().to_vec(), decimal_dtype) +} + /// Converts a DecimalBytePartsArray to its canonical DecimalArray representation. fn to_canonical_decimal( array: &DecimalBytePartsArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): support parts len != 1 - let prim = array.msp().clone().execute::(ctx)?; - // Depending on the decimal type and the min/max of the primitive array we can choose - // the correct buffer size - - Ok(match_each_signed_integer_ptype!(prim.ptype(), |P| { - // SAFETY: The primitive array's buffer is already validated with correct type. - // The decimal dtype matches the array's dtype, and validity is preserved. - unsafe { - DecimalArray::new_unchecked( - prim.to_buffer::

(), - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - prim.validity()?, - ) - } - .into_array() - })) + let msp = array.msp().clone().execute::(ctx)?; + let lower_parts = array + .lower_parts() + .iter() + .map(|part| part.clone().execute::(ctx)) + .collect::>>()?; + + Ok(assemble_decimal(&msp, &lower_parts, decimal_dtype(array.as_view()))?.into_array()) } impl OperationsVTable for DecimalByteParts { @@ -303,17 +436,34 @@ impl OperationsVTable for DecimalByteParts { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): support parts len != 1 let scalar = array.msp().execute_scalar(index, ctx)?; // Note. values in msp, can only be signed integers upto size i64. let primitive_scalar = scalar.as_primitive(); - // TODO(joe): extend this to support multiple parts. - let value = primitive_scalar.as_::().vortex_expect("non-null"); - Scalar::try_new( - array.dtype().clone(), - Some(ScalarValue::Decimal(DecimalValue::I64(value))), - ) + let msp = primitive_scalar.as_::().vortex_expect("non-null"); + + let lower_parts = array + .lower_parts() + .iter() + .map(|part| { + Ok(part + .execute_scalar(index, ctx)? + .as_primitive() + .as_::() + .vortex_expect("lower parts are non-nullable")) + }) + .collect::>>()?; + + let value = if lower_parts.is_empty() { + DecimalValue::I64(msp) + } else { + match values_type(array)? { + DecimalType::I256 => DecimalValue::I256(combine_i256(msp, lower_parts.into_iter())), + _ => DecimalValue::I128(combine_i128(msp, lower_parts)), + } + }; + + Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value))) } } @@ -326,21 +476,32 @@ impl ValidityChild for DecimalByteParts { #[cfg(test)] mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; + use super::*; use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_of; + use crate::decimal_byte_parts::testing::i256_parts; #[test] fn test_scalar_at_decimal_parts() { @@ -380,4 +541,267 @@ mod tests { .unwrap() ); } + + /// The largest unscaled value a `Decimal(38, _)` can hold: `10^38 - 1`. + const MAX_PRECISION_38: i128 = 99_999_999_999_999_999_999_999_999_999_999_999_999; + + /// The largest unscaled value a `Decimal(76, _)` can hold: `10^76 - 1`. + fn max_precision_76() -> i256 { + i256::from_i128(10).wrapping_pow(76) - i256::ONE + } + + /// Values that exercise every 64-bit window of an `i128`, both signs, and the boundaries + /// where a lower part carries into the MSP. + fn wide_i128_values() -> Vec { + vec![ + 0, + 1, + -1, + (1 << 64) - 1, + 1 << 64, + -(1 << 64), + -((1 << 64) + 1), + MAX_PRECISION_38, + -MAX_PRECISION_38, + 1 << 100, + ] + } + + /// Values that exercise every 64-bit window of an `i256`. + fn wide_i256_values() -> Vec { + vec![ + i256::ZERO, + i256::ONE, + i256::ZERO - i256::ONE, + i256_of(0, u128::MAX), + i256_of(1, 0), + i256_of(-1, 0), + i256_of(-1, u128::MAX - 1), + i256_of(1 << 64, 12345), + max_precision_76(), + i256::ZERO - max_precision_76(), + ] + } + + #[rstest] + #[case::i128_non_nullable(i128_parts(wide_i128_values(), Validity::NonNullable))] + #[case::i256_non_nullable(i256_parts(wide_i256_values(), Validity::NonNullable))] + fn test_canonical_decimal_round_trips( + #[case] array: DecimalBytePartsArray, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(array, canonical, &mut ctx); + Ok(()) + } + + #[test] + fn test_lower_part_layout_i128() -> VortexResult<()> { + let array = i128_parts(vec![(3i128 << 64) | 7], Validity::NonNullable); + assert_eq!(array.lower_parts().len(), 1); + assert_eq!(array.msp().dtype().as_ptype(), PType::I64); + assert_eq!(array.lower_parts()[0].dtype(), &LOWER_PART_DTYPE); + + let mut ctx = array_session().create_execution_ctx(); + let msp = array.msp().clone().execute::(&mut ctx)?; + let lower = array.lower_parts()[0] + .clone() + .execute::(&mut ctx)?; + assert_eq!(msp.as_slice::(), &[3]); + assert_eq!(lower.as_slice::(), &[7]); + Ok(()) + } + + #[test] + fn test_lower_part_layout_i256() -> VortexResult<()> { + let array = i256_parts( + vec![i256_of((5i128 << 64) | 6, (7u128 << 64) | 8)], + Validity::NonNullable, + ); + assert_eq!(array.lower_parts().len(), MAX_LOWER_PARTS); + + let mut ctx = array_session().create_execution_ctx(); + let msp = array.msp().clone().execute::(&mut ctx)?; + assert_eq!(msp.as_slice::(), &[5]); + for (part, expected) in array.lower_parts().iter().zip([6u64, 7, 8]) { + let part = part.clone().execute::(&mut ctx)?; + assert_eq!(part.as_slice::(), &[expected]); + } + Ok(()) + } + + #[rstest] + #[case::i128(i128_parts(wide_i128_values(), Validity::AllValid))] + #[case::i256(i256_parts(wide_i256_values(), Validity::AllValid))] + fn test_scalar_at_matches_canonical(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)? + .into_array(); + let array = array.into_array(); + for idx in 0..array.len() { + assert_eq!( + array.execute_scalar(idx, &mut ctx)?, + canonical.execute_scalar(idx, &mut ctx)?, + "scalar mismatch at index {idx}" + ); + } + Ok(()) + } + + #[rstest] + fn test_scalar_at_matches_canonical_for_each_part_count( + #[values(false, true)] narrow_msp: bool, + #[values(0, 1, 2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let validity = Validity::from_iter([false, true, true]); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![0i8, 3, -3], validity) + } else { + PrimitiveArray::new(buffer![0i64, 3, -3], validity) + }; + let lower = [4u64, 1, 2] + .into_iter() + .take(lower_count) + .map(|word| PrimitiveArray::new(buffer![word; 3], Validity::NonNullable).into_array()) + .collect(); + let dtype = DecimalDType::new(if lower_count <= 1 { 38 } else { 76 }, 0); + let array = DecimalByteParts::try_new_with_lower_parts(msp.into_array(), lower, dtype)?; + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)?; + for row in 0..array.len() { + assert_eq!( + array.execute_scalar(row, &mut ctx)?, + canonical.execute_scalar(row, &mut ctx)? + ); + } + Ok(()) + } + + #[test] + fn test_scalar_at_null_with_lower_parts() -> VortexResult<()> { + let array = i128_parts( + vec![1i128 << 100, 2, 3], + Validity::Array(BoolArray::from_iter([false, true, true]).into_array()), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::decimal( + DecimalValue::I128(2), + DecimalDType::new(38, 2), + Nullability::Nullable + ) + ); + Ok(()) + } + + fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() + } + + fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() + } + + #[rstest] + #[case::signed_lower_part(vec![buffer![1i64, 2, 3].into_array()], DecimalDType::new(38, 2))] + #[case::nullable_lower_part( + vec![PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid).into_array()], + DecimalDType::new(38, 2) + )] + #[case::mismatched_length(vec![buffer![1u64, 2].into_array()], DecimalDType::new(38, 2))] + #[case::too_many_parts( + vec![lower_part(), lower_part(), lower_part(), lower_part()], + DecimalDType::new(76, 2) + )] + fn test_rejects_invalid_parts( + #[case] lower_parts: Vec, + #[case] decimal_dtype: DecimalDType, + ) { + assert!( + DecimalByteParts::try_new_with_lower_parts(msp(), lower_parts, decimal_dtype).is_err() + ); + } + + #[test] + fn test_wide_decimal_buffer_types() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let i128_array = i128_parts(vec![1i128 << 100], Validity::NonNullable); + let canonical = i128_array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + + let i256_array = i256_parts(vec![i256_of(1 << 100, 0)], Validity::NonNullable); + let canonical = i256_array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + + // A narrow MSP with a single lower part still fits 128 bits. + let array = DecimalByteParts::try_new_with_lower_parts( + buffer![1i8, -1, 0].into_array(), + vec![buffer![7u64, 7, 7].into_array()], + DecimalDType::new(38, 2), + )?; + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + assert_eq!( + canonical.buffer::().as_slice(), + &[(1i128 << 64) | 7, (-1i128 << 64) | 7, 7] + ); + + // Two lower parts under a narrow MSP overflow 128 bits, so the value widens. + let array = DecimalByteParts::try_new_with_lower_parts( + buffer![1i8].into_array(), + vec![buffer![0u64].into_array(), buffer![9u64].into_array()], + DecimalDType::new(76, 2), + )?; + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + assert_eq!(canonical.buffer::().as_slice(), &[i256_of(1, 9)]); + Ok(()) + } + + #[test] + fn test_unused_buffer_of_values_is_ignored_for_null_rows() -> VortexResult<()> { + // Null rows may hold arbitrary bits in the lower parts; they must stay null. + let array = DecimalByteParts::try_new_with_lower_parts( + PrimitiveArray::new( + buffer![0i64, 0, 0], + Validity::Array(BoolArray::from_iter([false, false, true]).into_array()), + ) + .into_array(), + vec![buffer![7u64, 9, 11].into_array()], + DecimalDType::new(38, 2), + )? + .into_array(); + + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + let canonical = array.clone().execute::(&mut ctx)?; + assert_arrays_eq!(array, canonical.into_array(), &mut ctx); + Ok(()) + } + #[test] + fn test_frozen_serializer_rejects_lower_parts() -> VortexResult<()> { + let session = array_session(); + let array = i128_parts(vec![1i128 << 70], Validity::NonNullable); + assert!(VTable::serialize(array.as_view(), &session).is_err()); + Ok(()) + } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs index d4052a4bed8..28503d5d8af 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs @@ -1,57 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::ArrayRef; -use vortex_array::ArrayView; -use vortex_array::IntoArray; -use vortex_array::arrays::Filter; +use vortex_array::arrays::dict::TakeReduceAdaptor; use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; -use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&DecimalBytePartsFilterPushDownRule), ParentRuleSet::lift(&CastReduceAdaptor(DecimalByteParts)), ParentRuleSet::lift(&FilterReduceAdaptor(DecimalByteParts)), ParentRuleSet::lift(&MaskReduceAdaptor(DecimalByteParts)), ParentRuleSet::lift(&SliceReduceAdaptor(DecimalByteParts)), + ParentRuleSet::lift(&TakeReduceAdaptor(DecimalByteParts)), ]); - -#[derive(Debug)] -struct DecimalBytePartsFilterPushDownRule; - -impl ArrayParentReduceRule for DecimalBytePartsFilterPushDownRule { - type Parent = Filter; - - fn reduce_parent( - &self, - child: ArrayView<'_, DecimalByteParts>, - parent: ArrayView<'_, Filter>, - _child_idx: usize, - ) -> VortexResult> { - // TODO(ngates): we should benchmark whether to push-down filters with "lower parts". - // For now, we only push down if there are no lower parts. - if !child._lower_parts.is_empty() { - return Ok(None); - } - - let new_msp = child.msp().filter(parent.filter_mask().clone())?; - let new_child = DecimalByteParts::try_new( - new_msp, - *child - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - )? - .into_array(); - Ok(Some(new_child)) - } -} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs index 14807421c73..e31f717d389 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs @@ -7,23 +7,13 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::slice::SliceReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::map_parts; impl SliceReduce for DecimalByteParts { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { - Ok(Some( - DecimalByteParts::try_new( - array.msp().slice(range)?, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - )? - .into_array(), - )) + map_parts(array, |part| part.slice(range.clone())).map(|d| Some(d.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs new file mode 100644 index 00000000000..d2ce68f3700 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Test-only helpers for building byte-parts arrays. + +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::DecimalByteParts; +use crate::DecimalBytePartsArray; +use crate::decimal_byte_parts::limbs::split_decimal; + +/// Encode a canonical decimal array as byte parts, splitting wide values into lower parts. +pub(crate) fn encode(decimal: &DecimalArray) -> VortexResult { + let parts = split_decimal(decimal, &mut array_session().create_execution_ctx())?; + DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ) +} + +/// An `i128`-backed decimal array, encoded as byte parts with one lower part. +pub(crate) fn i128_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { + encode(&DecimalArray::new( + Buffer::from(values), + DecimalDType::new(38, 2), + validity, + )) + .vortex_expect("valid decimal byte parts") +} + +/// An `i256`-backed decimal array, encoded as byte parts with three lower parts. +pub(crate) fn i256_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { + encode(&DecimalArray::new( + Buffer::from(values), + DecimalDType::new(76, 2), + validity, + )) + .vortex_expect("valid decimal byte parts") +} + +/// Build an `i256` from a signed high `i128` and unsigned low `u128`. +pub(crate) fn i256_of(high: i128, low: u128) -> i256 { + i256::from_parts(low, high) +} diff --git a/encodings/decimal-byte-parts/tests/props.rs b/encodings/decimal-byte-parts/tests/props.rs new file mode 100644 index 00000000000..e33606b2880 --- /dev/null +++ b/encodings/decimal-byte-parts/tests/props.rs @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Property tests for splitting decimals into byte parts and putting them back together. +//! +//! Every property here is the same shape: whatever the encoding does must be indistinguishable +//! from doing it to the canonical `DecimalArray`. Round tripping covers the split/assemble +//! pair directly; the compute properties cover it indirectly, since each one canonicalizes an +//! encoded array at the end. +//! +//! The generators deliberately reach the cases hand-written tests tend to miss: values that +//! straddle a 64-bit word boundary, negative values whose sign extension fills the words above +//! the most significant part, and null rows whose lower parts hold arbitrary bits. + +#![expect(clippy::tests_outside_test_module)] + +use hegel::TestCase; +use hegel::generators as gs; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArray; +use vortex_decimal_byte_parts::split_decimal; +use vortex_error::VortexExpect; + +/// Largest magnitude a `Decimal(38, _)` can hold: 38 nines. +const MAX_I128: i128 = 10i128.pow(38) - 1; + +/// Bound on the high `i128` half of an `i256` draw. `10^37 * 2^128` is about `3.4e75`, so any +/// value built from it stays inside the 76 digits a `Decimal(76, _)` can hold. +const MAX_I256_HIGH: i128 = 10i128.pow(37); + +/// Rows per generated array. Small enough to shrink usefully, large enough that a chunked or +/// vectorized path is not trivially degenerate. +const MAX_LEN: usize = 48; + +fn ctx() -> ExecutionCtx { + let session = array_session(); + vortex_decimal_byte_parts::initialize(&session); + session.create_execution_ctx() +} + +/// Encode a canonical decimal as byte parts, splitting wide values into lower parts. +fn encode(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> DecimalBytePartsArray { + let parts = split_decimal(decimal, ctx).vortex_expect("split"); + DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ) + .vortex_expect("valid byte parts") +} + +/// A validity mask of exactly `len` entries, so null rows exercise lower parts holding bits +/// that must never be read. +fn draw_validity(tc: &TestCase, len: usize) -> Validity { + let valid: Vec = tc.draw(gs::vecs(gs::booleans()).min_size(len).max_size(len)); + Validity::from_iter(valid) +} + +/// An `i128`-backed decimal. The bounds keep values inside `Decimal(38, 2)` while still +/// reaching both sides of the 64-bit word boundary the encoding splits on. +fn draw_i128_decimal(tc: &TestCase) -> DecimalArray { + let values: Vec = tc.draw( + gs::vecs( + gs::integers::() + .min_value(-MAX_I128) + .max_value(MAX_I128), + ) + .min_size(1) + .max_size(MAX_LEN), + ); + let validity = draw_validity(tc, values.len()); + DecimalArray::new(Buffer::from(values), DecimalDType::new(38, 2), validity) +} + +/// An `i256`-backed decimal, built from a signed high half and an unsigned low half so the +/// draw covers sign extension above the most significant part. +fn draw_i256_decimal(tc: &TestCase) -> DecimalArray { + let halves: Vec<(i128, u128)> = tc.draw( + gs::vecs(gs::tuples2( + gs::integers::() + .min_value(-MAX_I256_HIGH) + .max_value(MAX_I256_HIGH), + gs::integers::(), + )) + .min_size(1) + .max_size(MAX_LEN), + ); + let values: Vec = halves + .into_iter() + .map(|(high, low)| i256::from_parts(low, high)) + .collect(); + let validity = draw_validity(tc, values.len()); + DecimalArray::new(Buffer::from(values), DecimalDType::new(76, 2), validity) +} + +fn draw_decimal(tc: &TestCase) -> DecimalArray { + if tc.draw(gs::booleans()) { + draw_i128_decimal(tc) + } else { + draw_i256_decimal(tc) + } +} + +/// Canonicalize an encoded array back to a `DecimalArray`. +fn canonicalize(array: ArrayRef, ctx: &mut ExecutionCtx) -> DecimalArray { + array.execute::(ctx).vortex_expect("execute") +} + +/// A byte-parts array built directly from drawn parts, rather than by splitting a decimal. +/// +/// `split_decimal` only ever emits 0, 1 or 3 lower parts under an `i64` most significant +/// part, so drawing the part count here is the only way to reach the two-part shape and the +/// sign extension that sits above a most significant part below the top word. +fn draw_encoded(tc: &TestCase) -> (DecimalBytePartsArray, usize) { + let lower_part_count = tc.draw(gs::integers::().min_value(0).max_value(3)); + let msp: Vec = tc.draw( + gs::vecs(gs::integers::()) + .min_size(1) + .max_size(MAX_LEN), + ); + let len = msp.len(); + + let lower: Vec = (0..lower_part_count) + .map(|_| { + let part: Vec = + tc.draw(gs::vecs(gs::integers::()).min_size(len).max_size(len)); + PrimitiveArray::new(Buffer::from(part), Validity::NonNullable).into_array() + }) + .collect(); + + // The declared precision must be wide enough for what the parts assemble into. + let precision = match lower_part_count { + 0 => 18, + 1 => 38, + _ => 76, + }; + let msp = PrimitiveArray::new(Buffer::from(msp), draw_validity(tc, len)).into_array(); + let array = + DecimalByteParts::try_new_with_lower_parts(msp, lower, DecimalDType::new(precision, 2)) + .vortex_expect("valid byte parts"); + (array, len) +} + +/// Encoding a decimal and decoding it again must reproduce it exactly, including null rows +/// and the storage width. +#[hegel::test] +fn decoded_survives_encode_then_decode(tc: TestCase) { + let decimal = draw_decimal(&tc); + let mut ctx = ctx(); + + let round_tripped = canonicalize(encode(&decimal, &mut ctx).into_array(), &mut ctx); + + assert_eq!(round_tripped.values_type(), decimal.values_type()); + assert_arrays_eq!(decimal, round_tripped, &mut ctx); +} + +/// Decoding an encoded array and encoding it again must not change the values it decodes to. +/// +/// Starting from the encoded side reaches part counts `split_decimal` never produces, so this +/// covers layouts the property above cannot generate. It compares decoded values rather than +/// the arrays themselves because re-encoding normalizes the part count: splitting an `i256` +/// always yields three lower parts, whatever the original array carried. +#[hegel::test] +fn encoded_survives_decode_then_encode(tc: TestCase) { + let (array, _len) = draw_encoded(&tc); + let mut ctx = ctx(); + + let decoded = canonicalize(array.into_array(), &mut ctx); + let re_decoded = canonicalize(encode(&decoded, &mut ctx).into_array(), &mut ctx); + + assert_arrays_eq!(decoded, re_decoded, &mut ctx); +} + +// TODO(joe): restore the coverage removed alongside these two round trips. Each of the +// following was a property here and caught mutations that the round trips do not: +// +// - `scalar_at` against bulk canonicalization. `combine_i128`/`combine_i256` are a second +// implementation of the assembly loops and can drift from them silently. +// - filter, slice and take against the same operation on the canonical array. These caught +// part-order and word-placement mutations, though the round trips catch those too. +// - a serialize/decode round trip, which is the only property that exercised the metadata +// carrying the lower part count. +// - sign extension above a most significant part below the top word, checked against an +// expectation computed independently of the assembly loop. This is the one real gap: a +// round trip compares decode against decode, so a decode-side sign-extension bug is +// invisible to it. Dropping the sign extension is caught by neither property here. diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 07069f6309a..f173440f26d 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -418,7 +418,7 @@ fn trace_scan_filter_on_compressed_table() -> VortexResult<()> { optimize root=vortex.filter(i16, len=43) session=false reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(i16, len=43) child=vortex.dict(i16, len=4096) -> vortex.dict(i16, len=43) done output=vortex.dict(i16, len=43) - reduce_parent static:DecimalBytePartsFilterPushDownRule slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) + reduce_parent static:FilterReduceAdaptor(DecimalByteParts) slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) done output=vortex.decimal_byte_parts(decimal(15,2), len=43) optimize root=vortex.filter(vortex.date[days](i32), len=43) session=false optimize root=vortex.filter(i32, len=43) session=false @@ -454,6 +454,9 @@ fn trace_scan_take_on_compressed_table() -> VortexResult<()> { insta::assert_snapshot!(optimized.trace.to_string(), @" optimize root=vortex.dict({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) session=false + optimize root=vortex.dict(decimal(15,2), len=64) session=false + reduce_parent static:TakeReduceAdaptor(DecimalByteParts) slot=1 parent=vortex.dict(decimal(15,2), len=64) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=64) + done output=vortex.decimal_byte_parts(decimal(15,2), len=64) optimize root=vortex.dict(vortex.date[days](i32), len=64) session=false reduce_parent static:TakeReduceAdaptor(Extension) slot=1 parent=vortex.dict(vortex.date[days](i32), len=64) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.ext(vortex.date[days](i32), len=64) done output=vortex.ext(vortex.date[days](i32), len=64) diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 3475f26a175..a54df06fb4c 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -39,6 +39,13 @@ impl CudaExecute for DecimalBytePartsExecutor { .dtype() .as_decimal_opt() .vortex_expect("DecimalBytePartsArray dtype must be decimal"); + + // Reassembling lower parts into wide decimals is not implemented on the GPU; the MSP + // alone is not the value. + if !array.lower_parts().is_empty() { + vortex_bail!("DecimalBytePartsArray with lower parts is not supported on GPU") + } + let msp = array.msp().clone(); let PrimitiveDataParts { buffer, From 9a179ed402313d510eb2aacaebc4d619ac0bf4b1 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 23:38:41 -0400 Subject: [PATCH 08/16] Simplify decimal byte-parts array helpers Move array helpers onto a crate-private extension trait, preserve decimal precision and scale when replacing the MSP, and group slicing with the other compute operations. Inline canonical execution and select scalar storage directly from the lower-part count. Signed-off-by: "Matt Katz" --- .../src/decimal_byte_parts/compute/cast.rs | 6 +- .../src/decimal_byte_parts/compute/filter.rs | 6 +- .../src/decimal_byte_parts/compute/mask.rs | 5 +- .../src/decimal_byte_parts/compute/mod.rs | 1 + .../decimal_byte_parts/{ => compute}/slice.rs | 6 +- .../src/decimal_byte_parts/compute/take.rs | 6 +- .../src/decimal_byte_parts/mod.rs | 212 ++++++++---------- 7 files changed, 116 insertions(+), 126 deletions(-) rename encodings/decimal-byte-parts/src/decimal_byte_parts/{ => compute}/slice.rs (73%) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index 7b949fcd695..0594f15a7c6 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -10,8 +10,8 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; -use crate::decimal_byte_parts::with_msp; impl CastReduce for DecimalByteParts { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { @@ -20,7 +20,7 @@ impl CastReduce for DecimalByteParts { return Ok(None); } // DecimalBytePartsArray can only have Decimal dtype, so we only handle decimal-to-decimal casts - let DType::Decimal(target_decimal, target_nullability) = dtype else { + let DType::Decimal(_, target_nullability) = dtype else { // Cannot cast decimal to non-decimal types - delegate to canonical form return Ok(None); }; @@ -30,7 +30,7 @@ impl CastReduce for DecimalByteParts { .msp() .cast(array.msp().dtype().with_nullability(*target_nullability))?; - with_msp(array, new_msp, *target_decimal).map(|a| Some(a.into_array())) + array.with_msp(new_msp).map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index e4fb03a5ca0..49c4021dd18 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs @@ -9,11 +9,13 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use crate::DecimalByteParts; -use crate::decimal_byte_parts::map_parts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; impl FilterReduce for DecimalByteParts { fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { - map_parts(array, |part| part.filter(mask.clone())).map(|d| Some(d.into_array())) + array + .map_parts(|part| part.filter(mask.clone())) + .map(|d| Some(d.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index 2eea785794b..9a022ef34ce 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs @@ -9,15 +9,14 @@ use vortex_array::scalar_fn::fns::mask::MaskReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; -use crate::decimal_byte_parts::decimal_dtype; -use crate::decimal_byte_parts::with_msp; impl MaskReduce for DecimalByteParts { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { // Validity lives in the MSP, so only that part needs masking: the lower parts hold // undefined bits in null slots, which is exactly what a masked-out row is. let masked_msp = MaskExpr::try_new(array.msp().clone(), mask.clone())?.into_array(); - with_msp(array, masked_msp, decimal_dtype(array)).map(|a| Some(a.into_array())) + array.with_msp(masked_msp).map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 844468545cf..f9848e1b2e7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -7,6 +7,7 @@ mod filter; pub(crate) mod is_constant; pub(crate) mod kernel; mod mask; +mod slice; mod take; #[cfg(test)] diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs similarity index 73% rename from encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs rename to encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs index e31f717d389..1a2efc9034e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs @@ -10,10 +10,12 @@ use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::map_parts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; impl SliceReduce for DecimalByteParts { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { - map_parts(array, |part| part.slice(range.clone())).map(|d| Some(d.into_array())) + array + .map_parts(|part| part.slice(range.clone())) + .map(|d| Some(d.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index 578834635b8..bdf7dda4f74 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs @@ -8,8 +8,8 @@ use vortex_array::arrays::dict::TakeReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; -use crate::decimal_byte_parts::map_parts; impl TakeReduce for DecimalByteParts { /// Taking wraps each part in a `Dict` without reading any buffer, so it reduces rather @@ -22,7 +22,9 @@ impl TakeReduce for DecimalByteParts { return Ok(None); } - map_parts(array, |part| part.take(indices.clone())).map(|a| Some(a.into_array())) + array + .map_parts(|part| part.take(indices.clone())) + .map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 4a104e82c2b..ae02c23902a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -11,7 +11,6 @@ use vortex_array::ArrayView; pub(crate) mod compute; mod limbs; mod rules; -mod slice; #[cfg(test)] pub(crate) mod testing; @@ -32,13 +31,12 @@ use vortex_array::ArraySlots; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; -use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::DecimalType; use vortex_array::dtype::PType; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; @@ -148,6 +146,48 @@ impl DecimalBytesPartsMetadata { } } +#[derive(Clone, Debug)] +pub struct DecimalByteParts; + +impl DecimalByteParts { + /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array. + pub fn try_new( + msp: ArrayRef, + decimal_dtype: DecimalDType, + ) -> VortexResult { + Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + } + + /// Construct a new [`DecimalBytePartsArray`] from an MSP array, its lower parts, and a + /// decimal dtype. + /// + /// Lower parts are ordered most significant first and must each be a non-nullable `u64` + /// array of the same length as the MSP. See [`split_decimal`] for producing them from a + /// canonical decimal array. + /// + /// # Errors + /// + /// Returns an error if the parts do not describe a valid decimal, see + /// [`DecimalBytePartsData::validate`]. + pub fn try_new_with_lower_parts( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> VortexResult { + // Lower parts are supported in memory; the frozen serializer still rejects them. + let len = msp.len(); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + Array::try_from_parts( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + ) + } +} + impl VTable for DecimalByteParts { type TypedArrayData = DecimalBytePartsData; @@ -242,7 +282,17 @@ impl VTable for DecimalByteParts { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - to_canonical_decimal(&array, ctx).map(ExecutionResult::done) + // Reassemble DecimalArray from split parts + let msp = array.msp().clone().execute::(ctx)?; + let lower_parts = array + .lower_parts() + .iter() + .map(|part| part.clone().execute::(ctx)) + .collect::>>()?; + + let assembled = assemble_decimal(&msp, &lower_parts, array.decimal_dtype())?; + + Ok(ExecutionResult::done(assembled)) } } @@ -256,9 +306,11 @@ pub struct DecimalBytePartsSlots { pub lower_parts: Vec, } -/// This array encodes decimals as between 1-4 columns of primitive typed children. -/// The most significant part (msp) storing the most significant decimal bits. -/// This array must be signed and is nullable iff the decimal is nullable. +/// This array encodes decimals by splitting them between 1-4 columns of primitive typed children. +/// +/// The most significant part (MSP) stores the most significant decimal bits. It is signed and is +/// nullable iff the decimal is nullable. +/// /// Every lower part is a non-nullable `u64` holding a raw 64-bit window of the value. /// /// e.g. for a decimal i128 \[ 127..64 | 63..0 \] msp = 127..64 and lower_part\[0\] = 63..0 @@ -300,7 +352,7 @@ impl DecimalBytePartsData { vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); let lower_part_count = lower_parts.len(); - // Physical storage may be wider than the declared precision, as for DecimalArray. + vortex_ensure!( lower_part_count <= MAX_LOWER_PARTS, "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" @@ -321,114 +373,47 @@ impl DecimalBytePartsData { } } -#[derive(Clone, Debug)] -pub struct DecimalByteParts; +pub(crate) trait DecimalBytePartsArrayExt: DecimalBytePartsArraySlotsExt { + /// The decimal precision and scale, validated when the array was constructed. + fn decimal_dtype(&self) -> DecimalDType { + *self + .as_ref() + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype") + } -impl DecimalByteParts { - /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. - /// - /// # Errors + /// Rebuild the array by applying `f` to the MSP and every lower part, in slot order. /// - /// Returns an error if the MSP is not a signed integer array. - pub fn try_new( - msp: ArrayRef, - decimal_dtype: DecimalDType, + /// This applies row operations such as slicing and filtering to all parts together, + /// preserving the decimal precision and scale. + fn map_parts( + &self, + mut f: impl FnMut(&ArrayRef) -> VortexResult, ) -> VortexResult { - Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + let msp = f(self.msp())?; + let lower_parts = self + .lower_parts() + .iter() + .map(&mut f) + .collect::>>()?; + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, self.decimal_dtype()) } - /// Construct a new [`DecimalBytePartsArray`] from an MSP array, its lower parts, and a - /// decimal dtype. + /// Rebuild the array with a replacement MSP, preserving its lower parts, precision and scale. /// - /// Lower parts are ordered most significant first and must each be a non-nullable `u64` - /// array of the same length as the MSP. See [`split_decimal`] for producing them from a - /// canonical decimal array. - /// - /// # Errors - /// - /// Returns an error if the parts do not describe a valid decimal, see - /// [`DecimalBytePartsData::validate`]. - pub fn try_new_with_lower_parts( - msp: ArrayRef, - lower_parts: Vec, - decimal_dtype: DecimalDType, - ) -> VortexResult { - // Lower parts are supported in memory; the frozen serializer still rejects them. - let len = msp.len(); - let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); - let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); - Array::try_from_parts( - ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + /// Use this for operations such as masking and nullability casts that only affect the MSP. + /// The replacement MSP determines the result's nullability. + fn with_msp(&self, msp: ArrayRef) -> VortexResult { + DecimalByteParts::try_new_with_lower_parts( + msp, + self.lower_parts().to_vec(), + self.decimal_dtype(), ) } } -/// The decimal storage type this array canonicalizes to. -fn values_type(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { - match array.lower_parts().len() { - 0 => DecimalType::try_from(array.msp().dtype().as_ptype()), - 1 => Ok(DecimalType::I128), - 2 | 3 => Ok(DecimalType::I256), - count => vortex_bail!("at most {MAX_LOWER_PARTS} lower parts are supported, got {count}"), - } -} - -/// The decimal dtype this array carries. -/// -/// Guaranteed to be a decimal by construction: [`DecimalBytePartsData::validate`] rejects -/// every other dtype. -pub(crate) fn decimal_dtype(array: ArrayView<'_, DecimalByteParts>) -> DecimalDType { - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype") -} - -/// Rebuild the array by applying `f` to the MSP and to every lower part, in slot order. -/// -/// Part-wise operations must touch every part. Going through this rather than calling -/// [`DecimalByteParts::try_new_with_lower_parts`] directly makes dropping a lower part — -/// which silently corrupts wide values — unrepresentable. -pub(crate) fn map_parts( - array: ArrayView<'_, DecimalByteParts>, - mut f: impl FnMut(&ArrayRef) -> VortexResult, -) -> VortexResult { - let msp = f(array.msp())?; - let lower_parts = array - .lower_parts() - .iter() - .map(&mut f) - .collect::>>()?; - DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal_dtype(array)) -} - -/// Rebuild the array with a replacement MSP, keeping its lower parts untouched. -/// -/// Only valid for operations that cannot change a row's magnitude bits — a nullability cast -/// or a mask — since the lower parts keep whatever bits they held. That is sound because -/// validity lives in the MSP alone, so lower-part bits in a null row are already undefined. -pub(crate) fn with_msp( - array: ArrayView<'_, DecimalByteParts>, - msp: ArrayRef, - decimal_dtype: DecimalDType, -) -> VortexResult { - DecimalByteParts::try_new_with_lower_parts(msp, array.lower_parts().to_vec(), decimal_dtype) -} - -/// Converts a DecimalBytePartsArray to its canonical DecimalArray representation. -fn to_canonical_decimal( - array: &DecimalBytePartsArray, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let msp = array.msp().clone().execute::(ctx)?; - let lower_parts = array - .lower_parts() - .iter() - .map(|part| part.clone().execute::(ctx)) - .collect::>>()?; - - Ok(assemble_decimal(&msp, &lower_parts, decimal_dtype(array.as_view()))?.into_array()) -} +impl> DecimalBytePartsArrayExt for T {} impl OperationsVTable for DecimalByteParts { fn scalar_at( @@ -438,7 +423,8 @@ impl OperationsVTable for DecimalByteParts { ) -> VortexResult { let scalar = array.msp().execute_scalar(index, ctx)?; - // Note. values in msp, can only be signed integers upto size i64. + // Widen the MSP's signed value (i8/i16/i32/i64) to i64 for scalar reconstruction. + // The array retains its original MSP storage type. let primitive_scalar = scalar.as_primitive(); let msp = primitive_scalar.as_::().vortex_expect("non-null"); @@ -454,13 +440,10 @@ impl OperationsVTable for DecimalByteParts { }) .collect::>>()?; - let value = if lower_parts.is_empty() { - DecimalValue::I64(msp) - } else { - match values_type(array)? { - DecimalType::I256 => DecimalValue::I256(combine_i256(msp, lower_parts.into_iter())), - _ => DecimalValue::I128(combine_i128(msp, lower_parts)), - } + let value = match lower_parts.len() { + 0 => DecimalValue::I64(msp), + 1 => DecimalValue::I128(combine_i128(msp, lower_parts)), + _ => DecimalValue::I256(combine_i256(msp, lower_parts.into_iter())), }; Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value))) @@ -487,6 +470,7 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::DecimalType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; From 5d2d3ebbf0c04b20cd9d28f5acc9f5cfe5c00354 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 16:29:46 -0400 Subject: [PATCH 09/16] Register versioned decimal byte-part serialization Use one ArrayPlugin for the frozen single-part format and the new wide format. Preserve frozen files with wider physical storage and add wire contract tests plus an opt-in compatibility fixture. Signed-off-by: "Matt Katz" --- .../src/decimal_byte_parts/mod.rs | 390 +++++++++++++++++- encodings/decimal-byte-parts/src/lib.rs | 4 +- .../decimal-byte-parts/tests/format_v2.rs | 257 ++++++++++++ vortex-test/compat-gen/Cargo.toml | 5 + .../encodings/decimal_byte_parts_v2.rs | 116 ++++++ .../arrays/synthetic/encodings/mod.rs | 10 +- 6 files changed, 770 insertions(+), 12 deletions(-) create mode 100644 encodings/decimal-byte-parts/tests/format_v2.rs create mode 100644 vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index ae02c23902a..7d19dbf3f02 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -6,7 +6,9 @@ use std::fmt::Formatter; use std::hash::Hasher; use vortex_array::Array; +use vortex_array::ArrayDeserialization; use vortex_array::ArrayParts; +use vortex_array::ArraySerialization; use vortex_array::ArrayView; pub(crate) mod compute; mod limbs; @@ -26,11 +28,13 @@ use prost::Message as _; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; +use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; use vortex_array::ArraySlots; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; +use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::arrays::PrimitiveArray; @@ -178,7 +182,9 @@ impl DecimalByteParts { lower_parts: Vec, decimal_dtype: DecimalDType, ) -> VortexResult { - // Lower parts are supported in memory; the frozen serializer still rejects them. + // Building lower parts in memory is never gated — reading a file requires it. What is + // gated is the serialized form: an array carrying lower parts serializes under the + // `vortex.decimal_byte_parts_v2` format ID, which only editions that contain it may write. let len = msp.len(); let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); @@ -245,7 +251,7 @@ impl VTable for DecimalByteParts { ) -> VortexResult>> { vortex_ensure!( array.lower_parts().is_empty(), - "serializing DecimalByteParts with lower parts is not supported" + "serializing DecimalByteParts with lower parts requires DecimalBytePartsPlugin" ); Ok(Some( DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), @@ -415,6 +421,97 @@ pub(crate) trait DecimalBytePartsArrayExt: DecimalBytePartsArraySlotsExt { impl> DecimalBytePartsArrayExt for T {} +/// The `vortex.decimal_byte_parts_v2` serialized format ID: byte parts carrying lower parts. +/// +/// This is a serialized format, not a second in-memory encoding. `vortex.decimal_byte_parts` +/// froze promising a single child, so an array with lower parts serializes under this ID +/// instead, and both IDs deserialize back into the same [`DecimalBytePartsArray`]. A reader +/// that predates lower parts fails on this ID with an unknown-encoding error rather than +/// misreading the children. +pub fn decimal_byte_parts_v2_id() -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts_v2"); + *ID +} + +/// The [`ArrayPlugin`] for [`DecimalByteParts`], owning both of its serialized formats. +/// +/// An array without lower parts serializes under the frozen `vortex.decimal_byte_parts` ID, +/// byte-identical to files written before lower parts existed. An array carrying lower parts +/// serializes under [`decimal_byte_parts_v2_id`]. Reading holds each ID to its own contract: +/// the frozen ID carries no lower parts and the v2 ID carries at least one, so recognizing the +/// newer format never widens what the frozen one may mean. +/// +/// Register this plugin, or call [`crate::initialize`], to enable both formats. Registering +/// [`DecimalByteParts`] directly only supports the frozen format. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsPlugin; + +impl ArrayPlugin for DecimalBytePartsPlugin { + fn id(&self) -> ArrayId { + VTable::id(&DecimalByteParts) + } + + fn serialized_ids(&self) -> Vec { + vec![VTable::id(&DecimalByteParts), decimal_byte_parts_v2_id()] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!( + "DecimalByteParts plugin cannot serialize {}", + array.encoding_id() + ) + })?; + let serialized_id = if view.lower_parts().is_empty() { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + Ok(Some(ArraySerialization::from_array( + serialized_id, + array, + DecimalBytesPartsMetadata::from_array(view)?.encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let metadata = DecimalBytesPartsMetadata::decode(parts.metadata)?; + let lower_part_count = metadata.lower_part_count()?; + if parts.serialized_id == decimal_byte_parts_v2_id() { + vortex_ensure!( + lower_part_count > 0, + "{} must carry at least one lower part", + parts.serialized_id + ); + } else { + vortex_ensure!( + parts.serialized_id == VTable::id(&DecimalByteParts), + "DecimalByteParts plugin does not recognize serialized ID {}", + parts.serialized_id + ); + vortex_ensure!( + lower_part_count == 0, + "{} must not carry lower parts, got {lower_part_count}", + parts.serialized_id + ); + } + Ok(Array::try_from_parts(metadata.into_array_parts( + parts.dtype, + parts.len, + parts.children, + )?)? + .into_array()) + } +} + impl OperationsVTable for DecimalByteParts { fn scalar_at( array: ArrayView<'_, DecimalByteParts>, @@ -460,12 +557,15 @@ impl ValidityChild for DecimalByteParts { #[cfg(test)] mod tests { use rstest::rstest; + use vortex_array::ArrayContext; use vortex_array::ArrayRef; + use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; @@ -477,12 +577,18 @@ mod tests { use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; + use vortex_array::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; + use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; use super::*; use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::encode; use crate::decimal_byte_parts::testing::i128_parts; use crate::decimal_byte_parts::testing::i256_of; use crate::decimal_byte_parts::testing::i256_parts; @@ -693,6 +799,127 @@ mod tests { Ok(()) } + #[rstest] + #[case::one_lower_part(i128_parts(wide_i128_values(), Validity::NonNullable))] + #[case::three_lower_parts(i256_parts(wide_i256_values(), Validity::NonNullable))] + #[case::nullable_three_lower_parts(i256_parts(wide_i256_values(), Validity::AllValid))] + fn test_serde_round_trip_with_lower_parts( + #[case] array: DecimalBytePartsArray, + ) -> VortexResult<()> { + test_serde_round_trip(array) + } + + #[rstest] + #[case::no_lower_parts( + encode(&DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(9, 2), Validity::NonNullable)) + .vortex_expect("valid decimal byte parts") + )] + fn test_serde_round_trip_flat(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { + test_serde_round_trip(array) + } + + #[rstest] + fn test_deserialize_frozen_with_wider_storage( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] + validity: Validity, + ) -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(2, 0); + let expected = DecimalArray::new(buffer![1i8, 2, 3], decimal_dtype, validity.clone()); + let children = vec![PrimitiveArray::new(buffer![1i64, 2, 3], validity).into_array()]; + + // Metadata emitted by the frozen serializer for a single i64 child. + let decoded = DecimalBytePartsPlugin.deserialize( + ArrayDeserialization::new( + VTable::id(&DecimalByteParts), + expected.dtype(), + expected.len(), + &[8, 7], + &[], + &children, + ), + &session, + )?; + assert_arrays_eq!(expected, decoded, &mut ctx); + test_serde_round_trip(decoded.as_::().into_owned()) + } + + #[rstest] + #[case::i64(DecimalArray::new( + buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, + ))] + #[case::i128(DecimalArray::new( + buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, + ))] + #[case::i256(DecimalArray::new( + buffer![i256::from_i128(-99), i256::ZERO, i256::from_i128(99)], + DecimalDType::new(2, 0), Validity::NonNullable, + ))] + fn test_serde_round_trip_wider_storage(#[case] decimal: DecimalArray) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let encoded = encode(&decimal)?; + assert_arrays_eq!(decimal, encoded, &mut ctx); + assert_eq!( + encoded.execute_scalar(0, &mut ctx)?, + decimal.execute_scalar(0, &mut ctx)?, + ); + test_serde_round_trip(encoded) + } + + fn test_serde_round_trip(array: DecimalBytePartsArray) -> VortexResult<()> { + let session = array_session(); + // Both serialized formats must be registered: an array with lower parts comes back + // under the v2 format id. + crate::initialize(&session); + + let array = array.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + let lower_part_count = array + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len(); + + let expected_id = if lower_part_count == 0 { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + assert_eq!( + session + .array_serialize(&array)? + .vortex_expect("byte parts arrays are serializable") + .serialized_id, + expected_id + ); + + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; + let mut concat = ByteBufferMut::empty(); + for buf in serialized { + concat.extend_from_slice(buf.as_ref()); + } + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?; + + assert_eq!( + decoded + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len(), + lower_part_count, + "lower parts must survive serde" + ); + + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) + } + fn msp() -> ArrayRef { buffer![1i64, 2, 3].into_array() } @@ -721,6 +948,158 @@ mod tests { ); } + fn deserialize_with( + lower_part_count: u32, + children: Vec, + ) -> VortexResult { + let serialized_id = if lower_part_count == 0 { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + plugin_deserialize_with(serialized_id, lower_part_count, children) + .map(|array| array.as_::().into_owned()) + } + + #[test] + fn test_deserialize_reads_lower_parts() -> VortexResult<()> { + let array = deserialize_with(1, vec![msp(), lower_part()])?; + assert_eq!(array.lower_parts().len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!( + canonical.buffer::().as_slice(), + &[(1i128 << 64) | 1, (2i128 << 64) | 2, (3i128 << 64) | 3] + ); + Ok(()) + } + + /// An array read from a file can be handed straight back to a writer, bypassing both the + /// constructor and the compressor. Its serialized id must still be the v2 format, so a + /// writer whose permitted encodings predate the v2 format refuses it. + #[test] + fn read_lower_parts_serialize_under_the_wide_format() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + + let array = deserialize_with(1, vec![msp(), lower_part()])?.into_array(); + + let serialization = session + .array_serialize(&array)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialization.serialized_id, decimal_byte_parts_v2_id()); + + let restricted = ArrayContext::empty() + .with_allowed_ids([VTable::id(&DecimalByteParts)].into_iter().collect()); + let err = array + .serialize(&restricted, &session, &SerializeOptions::default()) + .expect_err("expected the permitted-encoding check to refuse the v2 format"); + assert!( + err.to_string().contains("not permitted"), + "error should name the permitted-encoding check, got: {err}" + ); + Ok(()) + } + + /// Reading back an array that already carries lower parts, and computing over it, must + /// always work: the v2 format only restricts which writers may emit it. If reading or + /// the rebuild that every compute kernel does were blocked, a session whose editions + /// predate the v2 format could not read a file written by one that includes it. + #[test] + fn compute_over_existing_lower_parts_is_not_gated() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + // Stands in for an array materialized from a file: the parts already exist. + let array = deserialize_with(1, vec![msp(), lower_part()])?.into_array(); + + let sliced = array.slice(0..2)?; + assert_eq!(sliced.execute::(&mut ctx)?.len(), 2); + Ok(()) + } + + #[rstest] + fn test_deserialize_redundant_lower_parts( + #[values(2, 3)] lower_part_count: u32, + ) -> VortexResult<()> { + let mut children = vec![buffer![0i64; 3].into_array()]; + children.extend((1..lower_part_count).map(|_| buffer![0u64; 3].into_array())); + children.push(lower_part()); + let array = deserialize_with(lower_part_count, children)?; + let expected = DecimalArray::new( + buffer![1i128, 2, 3], + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(expected, array, &mut ctx); + test_serde_round_trip(array) + } + + #[test] + fn test_deserialize_rejects_child_count_mismatch() { + // Metadata claiming a lower part that was not serialized. + assert!(deserialize_with(1, vec![msp()]).is_err()); + // Metadata claiming fewer lower parts than there are children. + assert!(deserialize_with(0, vec![msp(), lower_part()]).is_err()); + // Metadata claiming more lower parts than the encoding supports. + assert!( + deserialize_with( + 4, + vec![ + msp(), + lower_part(), + lower_part(), + lower_part(), + lower_part() + ] + ) + .is_err() + ); + } + + fn plugin_deserialize_with( + serialized_id: ArrayId, + lower_part_count: u32, + children: Vec, + ) -> VortexResult { + let metadata = DecimalBytesPartsMetadata { + zeroth_child_ptype: PType::I64 as i32, + lower_part_count, + } + .encode_to_vec(); + let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); + DecimalBytePartsPlugin.deserialize( + ArrayDeserialization::new(serialized_id, &dtype, 3, &metadata, &[], &children), + &array_session(), + ) + } + + /// Each serialized ID keeps its own contract: the frozen ID never carries lower parts, and + /// the v2 ID is never written without them. + #[rstest] + #[case::frozen_without_lower_parts(VTable::id(&DecimalByteParts), 0, vec![msp()], true)] + #[case::frozen_with_lower_parts( + VTable::id(&DecimalByteParts), + 1, + vec![msp(), lower_part()], + false + )] + #[case::v2_with_lower_parts(decimal_byte_parts_v2_id(), 1, vec![msp(), lower_part()], true)] + #[case::v2_without_lower_parts(decimal_byte_parts_v2_id(), 0, vec![msp()], false)] + #[case::unknown_id(ArrayVTable::id(&Primitive), 0, vec![msp()], false)] + fn plugin_holds_each_id_to_its_contract( + #[case] serialized_id: ArrayId, + #[case] lower_part_count: u32, + #[case] children: Vec, + #[case] accepted: bool, + ) { + let result = plugin_deserialize_with(serialized_id, lower_part_count, children); + assert_eq!(result.is_ok(), accepted, "{serialized_id}: {result:?}"); + } + #[test] fn test_wide_decimal_buffer_types() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -781,11 +1160,4 @@ mod tests { assert_arrays_eq!(array, canonical.into_array(), &mut ctx); Ok(()) } - #[test] - fn test_frozen_serializer_rejects_lower_parts() -> VortexResult<()> { - let session = array_session(); - let array = i128_parts(vec![1i128 << 70], Validity::NonNullable); - assert!(VTable::serialize(array.as_view(), &session).is_err()); - Ok(()) - } } diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 36a53c3a614..2557555eac8 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -22,7 +22,9 @@ use vortex_session::VortexSession; /// Initialize decimal-byte-parts encoding in the given session. pub fn initialize(session: &VortexSession) { - session.arrays().register(DecimalByteParts); + // One plugin owns both serialized formats: registering it reads either ID and writes the + // one that fits the array. Which of them a writer may emit is decided by its editions. + session.arrays().register(DecimalBytePartsPlugin); compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/decimal-byte-parts/tests/format_v2.rs b/encodings/decimal-byte-parts/tests/format_v2.rs new file mode 100644 index 00000000000..338598df032 --- /dev/null +++ b/encodings/decimal-byte-parts/tests/format_v2.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The v2 serialized format. +//! +//! Lower parts can be built and computed over freely. What changes with them is the bytes: +//! an array carrying lower parts serializes under `vortex.decimal_byte_parts_v2` rather +//! than the frozen `vortex.decimal_byte_parts` format, so a writer restricted to editions +//! without the v2 format refuses it, and a reader that predates lower parts fails with an +//! unknown-encoding error instead of misreading the children. These tests pin all of that: +//! construction always works, the serialized id tracks the parts, and the permitted-encoding +//! check applies to the serialized id. + +#![expect(clippy::tests_outside_test_module)] + +use vortex_array::ArrayContext; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::serde::SerializeOptions; +use vortex_array::session::ArraySessionExt; +use vortex_buffer::buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() +} + +fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() +} + +fn session() -> VortexSession { + let session = vortex_array::array_session(); + vortex_decimal_byte_parts::initialize(&session); + session +} + +/// The wire ID the session's plugin picks for `array`. +fn serialized_id(session: &VortexSession, array: &ArrayRef) -> VortexResult { + Ok(session + .array_serialize(array)? + .ok_or_else(|| vortex_err!("byte parts arrays are serializable"))? + .serialized_id) +} + +/// A single-child array is the stable shape and is always constructible. +#[test] +fn single_child_is_always_allowed() { + assert!(DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2)).is_ok()); + assert!( + DecimalByteParts::try_new_with_lower_parts(msp(), vec![], DecimalDType::new(19, 2)).is_ok() + ); +} + +/// Building lower parts in memory is always allowed — reading a file requires it. What +/// changes is the serialized format, not what can be constructed. +#[test] +fn lower_parts_can_always_be_constructed() { + assert!( + DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + ) + .is_ok() + ); +} + +/// A single-child array keeps the frozen format id, byte-compatible with every reader since +/// the format froze; lower parts move the array onto the v2 format id. +#[test] +fn serialized_id_tracks_lower_parts() -> VortexResult<()> { + let session = session(); + + let flat = DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2))?.into_array(); + assert_eq!( + serialized_id(&session, &flat)?, + ArrayVTable::id(&DecimalByteParts) + ); + + let wide = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + assert_eq!(serialized_id(&session, &wide)?, decimal_byte_parts_v2_id()); + + Ok(()) +} + +/// The permitted-encoding check applies to the serialized id. A context restricted to the +/// frozen format — a writer whose enabled editions predate the v2 format — must refuse an +/// array carrying lower parts, however it was obtained. +/// +/// `ArrayParts` is public and `DecimalBytePartsData` is a public unit struct, so a caller can +/// assemble slots by hand and go straight to `Array::try_from_parts`, bypassing +/// `try_new_with_lower_parts` entirely. That back door is left open on purpose — it is the +/// same path `deserialize` uses. What must hold is that the resulting array cannot become +/// bytes under the frozen id. +#[test] +fn wide_format_is_refused_where_not_permitted() -> VortexResult<()> { + use vortex_array::Array; + use vortex_array::ArrayParts; + use vortex_array::ArraySlots; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_decimal_byte_parts::DecimalBytePartsData; + + let session = session(); + + let mut slots = ArraySlots::with_capacity(2); + slots.push(Some(msp())); + slots.push(Some(lower_part())); + + // Assembling the array by hand succeeds: this is the shape a file read produces. + let array = Array::try_from_parts( + ArrayParts::new( + DecimalByteParts, + DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable), + 3, + DecimalBytePartsData, + ) + .with_slots(slots), + )? + .into_array(); + assert_eq!(array.nchildren(), 2, "expected two limbs"); + + // A context permitting only the frozen format refuses to write it. + let restricted = ArrayContext::empty() + .with_allowed_ids([ArrayVTable::id(&DecimalByteParts)].into_iter().collect()); + let err = array + .serialize(&restricted, &session, &SerializeOptions::default()) + .expect_err("expected the permitted-encoding check to refuse the v2 format"); + assert!( + err.to_string().contains("not permitted"), + "error should name the permitted-encoding check, got: {err}" + ); + + // Permitting the v2 format id is exactly what allows the same array through. + let permissive = ArrayContext::empty().with_allowed_ids( + [ + ArrayVTable::id(&DecimalByteParts), + decimal_byte_parts_v2_id(), + ArrayVTable::id(&vortex_array::arrays::Primitive), + ] + .into_iter() + .collect(), + ); + let serialized = array.serialize(&permissive, &session, &SerializeOptions::default())?; + assert!(!serialized.is_empty()); + assert!( + permissive.to_ids().contains(&decimal_byte_parts_v2_id()), + "the file's encoding table must carry the v2 format id" + ); + + Ok(()) +} + +#[test] +fn bare_vtable_refuses_wide_serialization() -> VortexResult<()> { + let session = vortex_array::array_session(); + session.arrays().register(DecimalByteParts); + let array = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + let restricted = ArrayContext::empty().with_allowed_ids( + [ + ArrayVTable::id(&DecimalByteParts), + ArrayVTable::id(&vortex_array::arrays::Primitive), + ] + .into_iter() + .collect(), + ); + + assert!( + array + .serialize(&restricted, &session, &SerializeOptions::default()) + .is_err(), + "bare VTable registration must not write lower parts under the frozen ID" + ); + Ok(()) +} + +#[test] +fn bare_vtable_refuses_lower_parts_on_frozen_id() -> VortexResult<()> { + let session = vortex_array::array_session(); + session.arrays().register(DecimalByteParts); + let array = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + let id = ArrayVTable::id(&DecimalByteParts); + let plugin = session + .arrays() + .registry() + .get(&id) + .ok_or_else(|| vortex_err!("missing decimal plugin"))?; + let children = array.children(); + + // i64 MSP and one lower part, mislabeled as the frozen format. + let parts = ArrayDeserialization::new( + id, + array.dtype(), + array.len(), + &[8, 7, 16, 1], + &[], + &children, + ); + assert!(plugin.deserialize(parts, &session).is_err()); + Ok(()) +} + +#[test] +fn bare_vtable_keeps_frozen_serde() -> VortexResult<()> { + let session = vortex_array::array_session(); + session.arrays().register(DecimalByteParts); + let array = DecimalByteParts::try_new(msp(), DecimalDType::new(2, 0))?.into_array(); + let serialized = session + .array_serialize(&array)? + .ok_or_else(|| vortex_err!("missing decimal serialization"))?; + assert_eq!(serialized.serialized_id, ArrayVTable::id(&DecimalByteParts)); + assert_eq!(serialized.metadata, [8, 7]); + let plugin = session + .arrays() + .registry() + .get(&serialized.serialized_id) + .ok_or_else(|| vortex_err!("missing decimal plugin"))?; + let decoded = plugin.deserialize( + ArrayDeserialization::new( + serialized.serialized_id, + array.dtype(), + array.len(), + &serialized.metadata, + &[], + &serialized.children, + ), + &session, + )?; + assert_arrays_eq!(array, decoded, &mut session.create_execution_ctx()); + Ok(()) +} diff --git a/vortex-test/compat-gen/Cargo.toml b/vortex-test/compat-gen/Cargo.toml index 4a62aca3671..2a5fe657d9b 100644 --- a/vortex-test/compat-gen/Cargo.toml +++ b/vortex-test/compat-gen/Cargo.toml @@ -20,6 +20,11 @@ name = "vortex-compat" path = "src/main.rs" test = false +[features] +# Fixtures for encodings whose on-disk shape is not yet stable. Kept out of the default +# fixture set so a default build never publishes a file older readers cannot open. +unstable_encodings = ["vortex/unstable_encodings"] + [dependencies] # Vortex crates vortex = { workspace = true, features = ["files", "tokio", "zstd"] } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs new file mode 100644 index 00000000000..dfdcd893860 --- /dev/null +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Wide `DecimalByteParts` fixtures: values that need lower parts. +//! +//! These live in their own fixture file rather than as extra columns on +//! `decimal_byte_parts.vortex` because a fixture's `build()` is immutable once published. +//! `check` compares files written by older releases against what `build()` produces today, +//! so changing an existing fixture's schema fails the check against every previously +//! published version — see "Fixture evolution" in `DESIGN.md`, which requires a new fixture +//! file with a new name for a new type, encoding, or structural pattern. +//! +//! So `decimal_byte_parts.vortex` keeps testing exactly what it always did, decimals whose +//! values fit a single signed part, and the MSP-plus-lower-parts layout added alongside it +//! is covered here instead. + +use vortex::array::ArrayId; +use vortex::array::ArrayRef; +use vortex::array::ArrayVTable; +use vortex::array::IntoArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::StructArray; +use vortex::array::dtype::DecimalDType; +use vortex::array::dtype::FieldNames; +use vortex::array::dtype::i256; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::encodings::decimal_byte_parts::DecimalByteParts; +use vortex::encodings::decimal_byte_parts::DecimalBytePartsArray; +use vortex::encodings::decimal_byte_parts::split_decimal; +use vortex::error::VortexResult; +use vortex_array::ExecutionCtx; + +use super::N; +use crate::fixtures::FlatLayoutFixture; + +/// Encode a canonical decimal as byte parts, splitting wide values into lower parts. +fn encode_byte_parts( + decimal: &DecimalArray, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let parts = split_decimal(decimal, ctx)?; + DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ) +} + +pub struct DecimalBytePartsV2Fixture; + +impl FlatLayoutFixture for DecimalBytePartsV2Fixture { + fn name(&self) -> &str { + "decimal_byte_parts_v2.vortex" + } + + fn description(&self) -> &str { + "Wide decimal arrays split into a most significant part plus 64-bit lower parts" + } + + fn expected_encodings(&self) -> Vec { + vec![DecimalByteParts.id()] + } + + fn build(&self, ctx: &mut ExecutionCtx) -> VortexResult { + // An `i128` magnitude above 2^64, so the encoding must carry one lower part. + let wide_128_dtype = DecimalDType::new(38, 2); + let wide_128 = DecimalArray::new( + (0..N as i128) + .map(|i| 10i128.pow(25) + i * 7) + .collect::>(), + wide_128_dtype, + Validity::NonNullable, + ); + let wide_128_arr = encode_byte_parts(&wide_128, ctx)?; + + // Negative values, so the sign extension above the MSP is exercised on read back. + let wide_128_negative = DecimalArray::new( + (0..N as i128) + .map(|i| -(10i128.pow(25)) - i * 7) + .collect::>(), + wide_128_dtype, + Validity::NonNullable, + ); + let wide_128_negative_arr = encode_byte_parts(&wide_128_negative, ctx)?; + + // An `i256` magnitude beyond 128 bits, so all three lower parts are populated, with + // nulls to pin that validity is carried by the MSP alone. + let wide_256_dtype = DecimalDType::new(76, 2); + let base = i256::from_i128(10).wrapping_pow(40); + let wide_256 = DecimalArray::new( + (0..N as i128) + .map(|i| base + i256::from_i128(i * 7)) + .collect::>(), + wide_256_dtype, + Validity::from_iter((0..N).map(|i| i % 7 != 0)), + ); + let wide_256_arr = encode_byte_parts(&wide_256, ctx)?; + + let arr = StructArray::try_new( + FieldNames::from([ + "dec_wide_128", + "dec_wide_128_negative", + "dec_wide_256_nullable", + ]), + vec![ + wide_128_arr.into_array(), + wide_128_negative_arr.into_array(), + wide_256_arr.into_array(), + ], + N, + Validity::NonNullable, + )?; + Ok(arr.into_array()) + } +} diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs index 830b50450da..5af7596ca7b 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs @@ -12,6 +12,8 @@ mod bytebool; mod constant; mod datetimeparts; mod decimal_byte_parts; +#[cfg(feature = "unstable_encodings")] +mod decimal_byte_parts_v2; mod delta; mod dict; mod for_; @@ -31,7 +33,8 @@ pub(crate) const N: usize = 1024; /// All per-encoding fixtures. pub fn fixtures() -> Vec> { - vec![ + #[allow(unused_mut)] + let mut fixtures: Vec> = vec![ Box::new(alp::AlpFixture), Box::new(alprd::AlprdFixture), Box::new(bitpacked::BitPackedFixture), @@ -53,5 +56,8 @@ pub fn fixtures() -> Vec> { Box::new(zstd::ZstdFixture), Box::new(zigzag::ZigZagFixture), Box::new(constant::ConstantFixture), - ] + ]; + #[cfg(feature = "unstable_encodings")] + fixtures.push(Box::new(decimal_byte_parts_v2::DecimalBytePartsV2Fixture)); + fixtures } From 42e43028695bf8d912012ca49b686f1520f9ca99 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 00:43:35 -0400 Subject: [PATCH 10/16] Extract decimal byte-parts serde plugin and tests Move the plugin and serde coverage into plugin.rs while preserving metadata and frozen-format VTable serde. Share wide decimal test fixtures and exercise frozen compatibility through both registration paths. Include the v2 compatibility fixture in the default suite without enabling unstable encodings. Signed-off-by: "Matt Katz" --- .../src/decimal_byte_parts/mod.rs | 425 +----------- .../src/decimal_byte_parts/plugin.rs | 649 ++++++++++++++++++ .../src/decimal_byte_parts/testing.rs | 41 ++ .../decimal-byte-parts/tests/format_v2.rs | 257 ------- vortex-test/compat-gen/Cargo.toml | 5 - .../arrays/synthetic/encodings/mod.rs | 10 +- 6 files changed, 699 insertions(+), 688 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin.rs delete mode 100644 encodings/decimal-byte-parts/tests/format_v2.rs diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 7d19dbf3f02..455a32f98eb 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -5,13 +5,13 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use prost::Message as _; use vortex_array::Array; -use vortex_array::ArrayDeserialization; use vortex_array::ArrayParts; -use vortex_array::ArraySerialization; use vortex_array::ArrayView; pub(crate) mod compute; mod limbs; +mod plugin; mod rules; #[cfg(test)] pub(crate) mod testing; @@ -24,17 +24,16 @@ pub mod _benchmarking { pub use super::limbs::assemble_decimal; } -use prost::Message as _; +pub use plugin::DecimalBytePartsPlugin; +pub use plugin::decimal_byte_parts_v2_id; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; -use vortex_array::ArrayPlugin; use vortex_array::ArrayRef; use vortex_array::ArraySlots; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; -use vortex_array::IntoArray; use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::arrays::PrimitiveArray; @@ -421,97 +420,6 @@ pub(crate) trait DecimalBytePartsArrayExt: DecimalBytePartsArraySlotsExt { impl> DecimalBytePartsArrayExt for T {} -/// The `vortex.decimal_byte_parts_v2` serialized format ID: byte parts carrying lower parts. -/// -/// This is a serialized format, not a second in-memory encoding. `vortex.decimal_byte_parts` -/// froze promising a single child, so an array with lower parts serializes under this ID -/// instead, and both IDs deserialize back into the same [`DecimalBytePartsArray`]. A reader -/// that predates lower parts fails on this ID with an unknown-encoding error rather than -/// misreading the children. -pub fn decimal_byte_parts_v2_id() -> ArrayId { - static ID: CachedId = CachedId::new("vortex.decimal_byte_parts_v2"); - *ID -} - -/// The [`ArrayPlugin`] for [`DecimalByteParts`], owning both of its serialized formats. -/// -/// An array without lower parts serializes under the frozen `vortex.decimal_byte_parts` ID, -/// byte-identical to files written before lower parts existed. An array carrying lower parts -/// serializes under [`decimal_byte_parts_v2_id`]. Reading holds each ID to its own contract: -/// the frozen ID carries no lower parts and the v2 ID carries at least one, so recognizing the -/// newer format never widens what the frozen one may mean. -/// -/// Register this plugin, or call [`crate::initialize`], to enable both formats. Registering -/// [`DecimalByteParts`] directly only supports the frozen format. -#[derive(Clone, Debug)] -pub struct DecimalBytePartsPlugin; - -impl ArrayPlugin for DecimalBytePartsPlugin { - fn id(&self) -> ArrayId { - VTable::id(&DecimalByteParts) - } - - fn serialized_ids(&self) -> Vec { - vec![VTable::id(&DecimalByteParts), decimal_byte_parts_v2_id()] - } - - fn serialize( - &self, - array: &ArrayRef, - _session: &VortexSession, - ) -> VortexResult> { - let view = array.as_opt::().ok_or_else(|| { - vortex_err!( - "DecimalByteParts plugin cannot serialize {}", - array.encoding_id() - ) - })?; - let serialized_id = if view.lower_parts().is_empty() { - VTable::id(&DecimalByteParts) - } else { - decimal_byte_parts_v2_id() - }; - Ok(Some(ArraySerialization::from_array( - serialized_id, - array, - DecimalBytesPartsMetadata::from_array(view)?.encode_to_vec(), - ))) - } - - fn deserialize( - &self, - parts: ArrayDeserialization<'_>, - _session: &VortexSession, - ) -> VortexResult { - let metadata = DecimalBytesPartsMetadata::decode(parts.metadata)?; - let lower_part_count = metadata.lower_part_count()?; - if parts.serialized_id == decimal_byte_parts_v2_id() { - vortex_ensure!( - lower_part_count > 0, - "{} must carry at least one lower part", - parts.serialized_id - ); - } else { - vortex_ensure!( - parts.serialized_id == VTable::id(&DecimalByteParts), - "DecimalByteParts plugin does not recognize serialized ID {}", - parts.serialized_id - ); - vortex_ensure!( - lower_part_count == 0, - "{} must not carry lower parts, got {lower_part_count}", - parts.serialized_id - ); - } - Ok(Array::try_from_parts(metadata.into_array_parts( - parts.dtype, - parts.len, - parts.children, - )?)? - .into_array()) - } -} - impl OperationsVTable for DecimalByteParts { fn scalar_at( array: ArrayView<'_, DecimalByteParts>, @@ -557,15 +465,12 @@ impl ValidityChild for DecimalByteParts { #[cfg(test)] mod tests { use rstest::rstest; - use vortex_array::ArrayContext; use vortex_array::ArrayRef; - use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::DecimalArray; - use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; @@ -577,21 +482,17 @@ mod tests { use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; - use vortex_array::serde::SerializeOptions; - use vortex_array::serde::SerializedArray; - use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; - use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_session::registry::ReadContext; use super::*; use crate::DecimalByteParts; - use crate::decimal_byte_parts::testing::encode; use crate::decimal_byte_parts::testing::i128_parts; use crate::decimal_byte_parts::testing::i256_of; use crate::decimal_byte_parts::testing::i256_parts; + use crate::decimal_byte_parts::testing::wide_i128_values; + use crate::decimal_byte_parts::testing::wide_i256_values; #[test] fn test_scalar_at_decimal_parts() { @@ -632,47 +533,6 @@ mod tests { ); } - /// The largest unscaled value a `Decimal(38, _)` can hold: `10^38 - 1`. - const MAX_PRECISION_38: i128 = 99_999_999_999_999_999_999_999_999_999_999_999_999; - - /// The largest unscaled value a `Decimal(76, _)` can hold: `10^76 - 1`. - fn max_precision_76() -> i256 { - i256::from_i128(10).wrapping_pow(76) - i256::ONE - } - - /// Values that exercise every 64-bit window of an `i128`, both signs, and the boundaries - /// where a lower part carries into the MSP. - fn wide_i128_values() -> Vec { - vec![ - 0, - 1, - -1, - (1 << 64) - 1, - 1 << 64, - -(1 << 64), - -((1 << 64) + 1), - MAX_PRECISION_38, - -MAX_PRECISION_38, - 1 << 100, - ] - } - - /// Values that exercise every 64-bit window of an `i256`. - fn wide_i256_values() -> Vec { - vec![ - i256::ZERO, - i256::ONE, - i256::ZERO - i256::ONE, - i256_of(0, u128::MAX), - i256_of(1, 0), - i256_of(-1, 0), - i256_of(-1, u128::MAX - 1), - i256_of(1 << 64, 12345), - max_precision_76(), - i256::ZERO - max_precision_76(), - ] - } - #[rstest] #[case::i128_non_nullable(i128_parts(wide_i128_values(), Validity::NonNullable))] #[case::i256_non_nullable(i256_parts(wide_i256_values(), Validity::NonNullable))] @@ -799,127 +659,6 @@ mod tests { Ok(()) } - #[rstest] - #[case::one_lower_part(i128_parts(wide_i128_values(), Validity::NonNullable))] - #[case::three_lower_parts(i256_parts(wide_i256_values(), Validity::NonNullable))] - #[case::nullable_three_lower_parts(i256_parts(wide_i256_values(), Validity::AllValid))] - fn test_serde_round_trip_with_lower_parts( - #[case] array: DecimalBytePartsArray, - ) -> VortexResult<()> { - test_serde_round_trip(array) - } - - #[rstest] - #[case::no_lower_parts( - encode(&DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(9, 2), Validity::NonNullable)) - .vortex_expect("valid decimal byte parts") - )] - fn test_serde_round_trip_flat(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { - test_serde_round_trip(array) - } - - #[rstest] - fn test_deserialize_frozen_with_wider_storage( - #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] - validity: Validity, - ) -> VortexResult<()> { - let session = array_session(); - crate::initialize(&session); - let mut ctx = session.create_execution_ctx(); - let decimal_dtype = DecimalDType::new(2, 0); - let expected = DecimalArray::new(buffer![1i8, 2, 3], decimal_dtype, validity.clone()); - let children = vec![PrimitiveArray::new(buffer![1i64, 2, 3], validity).into_array()]; - - // Metadata emitted by the frozen serializer for a single i64 child. - let decoded = DecimalBytePartsPlugin.deserialize( - ArrayDeserialization::new( - VTable::id(&DecimalByteParts), - expected.dtype(), - expected.len(), - &[8, 7], - &[], - &children, - ), - &session, - )?; - assert_arrays_eq!(expected, decoded, &mut ctx); - test_serde_round_trip(decoded.as_::().into_owned()) - } - - #[rstest] - #[case::i64(DecimalArray::new( - buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, - ))] - #[case::i128(DecimalArray::new( - buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, - ))] - #[case::i256(DecimalArray::new( - buffer![i256::from_i128(-99), i256::ZERO, i256::from_i128(99)], - DecimalDType::new(2, 0), Validity::NonNullable, - ))] - fn test_serde_round_trip_wider_storage(#[case] decimal: DecimalArray) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let encoded = encode(&decimal)?; - assert_arrays_eq!(decimal, encoded, &mut ctx); - assert_eq!( - encoded.execute_scalar(0, &mut ctx)?, - decimal.execute_scalar(0, &mut ctx)?, - ); - test_serde_round_trip(encoded) - } - - fn test_serde_round_trip(array: DecimalBytePartsArray) -> VortexResult<()> { - let session = array_session(); - // Both serialized formats must be registered: an array with lower parts comes back - // under the v2 format id. - crate::initialize(&session); - - let array = array.into_array(); - let dtype = array.dtype().clone(); - let len = array.len(); - let lower_part_count = array - .as_opt::() - .vortex_expect("byte parts array") - .lower_parts() - .len(); - - let expected_id = if lower_part_count == 0 { - VTable::id(&DecimalByteParts) - } else { - decimal_byte_parts_v2_id() - }; - assert_eq!( - session - .array_serialize(&array)? - .vortex_expect("byte parts arrays are serializable") - .serialized_id, - expected_id - ); - - let array_ctx = ArrayContext::empty(); - let serialized = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; - let mut concat = ByteBufferMut::empty(); - for buf in serialized { - concat.extend_from_slice(buf.as_ref()); - } - let parts = SerializedArray::try_from(concat.freeze())?; - let decoded = parts.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?; - - assert_eq!( - decoded - .as_opt::() - .vortex_expect("byte parts array") - .lower_parts() - .len(), - lower_part_count, - "lower parts must survive serde" - ); - - let mut ctx = session.create_execution_ctx(); - assert_arrays_eq!(array, decoded, &mut ctx); - Ok(()) - } - fn msp() -> ArrayRef { buffer![1i64, 2, 3].into_array() } @@ -948,158 +687,6 @@ mod tests { ); } - fn deserialize_with( - lower_part_count: u32, - children: Vec, - ) -> VortexResult { - let serialized_id = if lower_part_count == 0 { - VTable::id(&DecimalByteParts) - } else { - decimal_byte_parts_v2_id() - }; - plugin_deserialize_with(serialized_id, lower_part_count, children) - .map(|array| array.as_::().into_owned()) - } - - #[test] - fn test_deserialize_reads_lower_parts() -> VortexResult<()> { - let array = deserialize_with(1, vec![msp(), lower_part()])?; - assert_eq!(array.lower_parts().len(), 1); - - let mut ctx = array_session().create_execution_ctx(); - let canonical = array.into_array().execute::(&mut ctx)?; - assert_eq!( - canonical.buffer::().as_slice(), - &[(1i128 << 64) | 1, (2i128 << 64) | 2, (3i128 << 64) | 3] - ); - Ok(()) - } - - /// An array read from a file can be handed straight back to a writer, bypassing both the - /// constructor and the compressor. Its serialized id must still be the v2 format, so a - /// writer whose permitted encodings predate the v2 format refuses it. - #[test] - fn read_lower_parts_serialize_under_the_wide_format() -> VortexResult<()> { - let session = array_session(); - crate::initialize(&session); - - let array = deserialize_with(1, vec![msp(), lower_part()])?.into_array(); - - let serialization = session - .array_serialize(&array)? - .vortex_expect("byte parts arrays are serializable"); - assert_eq!(serialization.serialized_id, decimal_byte_parts_v2_id()); - - let restricted = ArrayContext::empty() - .with_allowed_ids([VTable::id(&DecimalByteParts)].into_iter().collect()); - let err = array - .serialize(&restricted, &session, &SerializeOptions::default()) - .expect_err("expected the permitted-encoding check to refuse the v2 format"); - assert!( - err.to_string().contains("not permitted"), - "error should name the permitted-encoding check, got: {err}" - ); - Ok(()) - } - - /// Reading back an array that already carries lower parts, and computing over it, must - /// always work: the v2 format only restricts which writers may emit it. If reading or - /// the rebuild that every compute kernel does were blocked, a session whose editions - /// predate the v2 format could not read a file written by one that includes it. - #[test] - fn compute_over_existing_lower_parts_is_not_gated() -> VortexResult<()> { - let session = array_session(); - crate::initialize(&session); - let mut ctx = session.create_execution_ctx(); - - // Stands in for an array materialized from a file: the parts already exist. - let array = deserialize_with(1, vec![msp(), lower_part()])?.into_array(); - - let sliced = array.slice(0..2)?; - assert_eq!(sliced.execute::(&mut ctx)?.len(), 2); - Ok(()) - } - - #[rstest] - fn test_deserialize_redundant_lower_parts( - #[values(2, 3)] lower_part_count: u32, - ) -> VortexResult<()> { - let mut children = vec![buffer![0i64; 3].into_array()]; - children.extend((1..lower_part_count).map(|_| buffer![0u64; 3].into_array())); - children.push(lower_part()); - let array = deserialize_with(lower_part_count, children)?; - let expected = DecimalArray::new( - buffer![1i128, 2, 3], - DecimalDType::new(38, 2), - Validity::NonNullable, - ); - let mut ctx = array_session().create_execution_ctx(); - assert_arrays_eq!(expected, array, &mut ctx); - test_serde_round_trip(array) - } - - #[test] - fn test_deserialize_rejects_child_count_mismatch() { - // Metadata claiming a lower part that was not serialized. - assert!(deserialize_with(1, vec![msp()]).is_err()); - // Metadata claiming fewer lower parts than there are children. - assert!(deserialize_with(0, vec![msp(), lower_part()]).is_err()); - // Metadata claiming more lower parts than the encoding supports. - assert!( - deserialize_with( - 4, - vec![ - msp(), - lower_part(), - lower_part(), - lower_part(), - lower_part() - ] - ) - .is_err() - ); - } - - fn plugin_deserialize_with( - serialized_id: ArrayId, - lower_part_count: u32, - children: Vec, - ) -> VortexResult { - let metadata = DecimalBytesPartsMetadata { - zeroth_child_ptype: PType::I64 as i32, - lower_part_count, - } - .encode_to_vec(); - let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); - DecimalBytePartsPlugin.deserialize( - ArrayDeserialization::new(serialized_id, &dtype, 3, &metadata, &[], &children), - &array_session(), - ) - } - - /// Each serialized ID keeps its own contract: the frozen ID never carries lower parts, and - /// the v2 ID is never written without them. - #[rstest] - #[case::frozen_without_lower_parts(VTable::id(&DecimalByteParts), 0, vec![msp()], true)] - #[case::frozen_with_lower_parts( - VTable::id(&DecimalByteParts), - 1, - vec![msp(), lower_part()], - false - )] - #[case::v2_with_lower_parts(decimal_byte_parts_v2_id(), 1, vec![msp(), lower_part()], true)] - #[case::v2_without_lower_parts(decimal_byte_parts_v2_id(), 0, vec![msp()], false)] - #[case::unknown_id(ArrayVTable::id(&Primitive), 0, vec![msp()], false)] - fn plugin_holds_each_id_to_its_contract( - #[case] serialized_id: ArrayId, - #[case] lower_part_count: u32, - #[case] children: Vec, - #[case] accepted: bool, - ) { - let result = plugin_deserialize_with(serialized_id, lower_part_count, children); - assert_eq!(result.is_ok(), accepted, "{serialized_id}: {result:?}"); - } - #[test] fn test_wide_decimal_buffer_types() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin.rs new file mode 100644 index 00000000000..4b31fb8b53e --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin.rs @@ -0,0 +1,649 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serialization of decimal byte parts under the frozen and v2 format IDs. + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::vtable::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::DecimalByteParts; +use super::DecimalBytePartsArraySlotsExt; +use super::DecimalBytesPartsMetadata; + +/// The `vortex.decimal_byte_parts_v2` serialized format ID: byte parts carrying lower parts. +/// +/// This is a serialized format, not a second in-memory encoding. `vortex.decimal_byte_parts` +/// froze promising a single child, so an array with lower parts serializes under this ID +/// instead, and both IDs deserialize back into the same [`crate::DecimalBytePartsArray`]. A reader +/// that predates lower parts fails on this ID with an unknown-encoding error rather than +/// misreading the children. +pub fn decimal_byte_parts_v2_id() -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts_v2"); + *ID +} + +/// The [`ArrayPlugin`] for [`DecimalByteParts`], owning both of its serialized formats. +/// +/// An array without lower parts serializes under the frozen `vortex.decimal_byte_parts` ID, +/// byte-identical to files written before lower parts existed. An array carrying lower parts +/// serializes under [`decimal_byte_parts_v2_id`]. Reading holds each ID to its own contract: +/// the frozen ID carries no lower parts and the v2 ID carries at least one, so recognizing the +/// newer format never widens what the frozen one may mean. +/// +/// Register this plugin, or call [`crate::initialize`], to enable both formats. Registering +/// [`DecimalByteParts`] directly only supports the frozen format. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsPlugin; + +impl ArrayPlugin for DecimalBytePartsPlugin { + fn id(&self) -> ArrayId { + VTable::id(&DecimalByteParts) + } + + fn serialized_ids(&self) -> Vec { + vec![VTable::id(&DecimalByteParts), decimal_byte_parts_v2_id()] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!( + "DecimalByteParts plugin cannot serialize {}", + array.encoding_id() + ) + })?; + let serialized_id = if view.lower_parts().is_empty() { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + Ok(Some(ArraySerialization::from_array( + serialized_id, + array, + DecimalBytesPartsMetadata::from_array(view)?.encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let metadata = DecimalBytesPartsMetadata::decode(parts.metadata)?; + let lower_part_count = metadata.lower_part_count()?; + if parts.serialized_id == decimal_byte_parts_v2_id() { + vortex_ensure!( + lower_part_count > 0, + "{} must carry at least one lower part", + parts.serialized_id + ); + } else { + vortex_ensure!( + parts.serialized_id == VTable::id(&DecimalByteParts), + "DecimalByteParts plugin does not recognize serialized ID {}", + parts.serialized_id + ); + vortex_ensure!( + lower_part_count == 0, + "{} must not carry lower parts, got {lower_part_count}", + parts.serialized_id + ); + } + Ok(Array::try_from_parts(metadata.into_array_parts( + parts.dtype, + parts.len, + parts.children, + )?)? + .into_array()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::ArrayParts; + use vortex_array::ArraySlots; + use vortex_array::ArrayVTable; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::Primitive; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_array::session::ArraySessionExt; + use vortex_array::validity::Validity; + use vortex_buffer::ByteBufferMut; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_session::registry::ReadContext; + + use super::*; + use crate::DecimalBytePartsArray; + use crate::DecimalBytePartsData; + use crate::decimal_byte_parts::testing::encode; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_parts; + use crate::decimal_byte_parts::testing::wide_i128_values; + use crate::decimal_byte_parts::testing::wide_i256_values; + + #[rstest] + #[case::one_lower_part(i128_parts(wide_i128_values(), Validity::NonNullable))] + #[case::three_lower_parts(i256_parts(wide_i256_values(), Validity::NonNullable))] + #[case::nullable_three_lower_parts(i256_parts(wide_i256_values(), Validity::AllValid))] + fn test_serde_round_trip_with_lower_parts( + #[case] array: DecimalBytePartsArray, + ) -> VortexResult<()> { + test_serde_round_trip(array) + } + + #[rstest] + #[case::no_lower_parts( + encode(&DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(9, 2), Validity::NonNullable)) + .vortex_expect("valid decimal byte parts") + )] + fn test_serde_round_trip_flat(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { + test_serde_round_trip(array) + } + + #[rstest] + fn test_deserialize_frozen_with_wider_storage( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] + validity: Validity, + ) -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(2, 0); + let expected = DecimalArray::new(buffer![1i8, 2, 3], decimal_dtype, validity.clone()); + let children = vec![PrimitiveArray::new(buffer![1i64, 2, 3], validity).into_array()]; + + // Metadata emitted by the frozen serializer for a single i64 child. + let decoded = DecimalBytePartsPlugin.deserialize( + ArrayDeserialization::new( + VTable::id(&DecimalByteParts), + expected.dtype(), + expected.len(), + &[8, 7], + &[], + &children, + ), + &session, + )?; + assert_arrays_eq!(expected, decoded, &mut ctx); + test_serde_round_trip(decoded.as_::().into_owned()) + } + + #[rstest] + #[case::i64(DecimalArray::new( + buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, + ))] + #[case::i128(DecimalArray::new( + buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, + ))] + #[case::i256(DecimalArray::new( + buffer![i256::from_i128(-99), i256::ZERO, i256::from_i128(99)], + DecimalDType::new(2, 0), Validity::NonNullable, + ))] + fn test_serde_round_trip_wider_storage(#[case] decimal: DecimalArray) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let encoded = encode(&decimal)?; + assert_arrays_eq!(decimal, encoded, &mut ctx); + assert_eq!( + encoded.execute_scalar(0, &mut ctx)?, + decimal.execute_scalar(0, &mut ctx)?, + ); + test_serde_round_trip(encoded) + } + + fn test_serde_round_trip(array: DecimalBytePartsArray) -> VortexResult<()> { + let session = array_session(); + // Both serialized formats must be registered: an array with lower parts comes back + // under the v2 format id. + crate::initialize(&session); + + let array = array.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + let lower_part_count = array + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len(); + + let expected_id = if lower_part_count == 0 { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + assert_eq!( + session + .array_serialize(&array)? + .vortex_expect("byte parts arrays are serializable") + .serialized_id, + expected_id + ); + + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; + let mut concat = ByteBufferMut::empty(); + for buf in serialized { + concat.extend_from_slice(buf.as_ref()); + } + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?; + + assert_eq!( + decoded + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len(), + lower_part_count, + "lower parts must survive serde" + ); + + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) + } + + fn deserialize_with( + lower_part_count: u32, + children: Vec, + ) -> VortexResult { + let serialized_id = if lower_part_count == 0 { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + plugin_deserialize_with(serialized_id, lower_part_count, children) + .map(|array| array.as_::().into_owned()) + } + + #[test] + fn test_deserialize_reads_lower_parts() -> VortexResult<()> { + let array = deserialize_with(1, vec![msp(), lower_part()])?; + assert_eq!(array.lower_parts().len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!( + canonical.buffer::().as_slice(), + &[(1i128 << 64) | 1, (2i128 << 64) | 2, (3i128 << 64) | 3] + ); + Ok(()) + } + + /// An array read from a file can be handed straight back to a writer, bypassing both the + /// constructor and the compressor. Its serialized id must still be the v2 format, so a + /// writer whose permitted encodings predate the v2 format refuses it. + #[test] + fn read_lower_parts_serialize_under_the_wide_format() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + + let array = deserialize_with(1, vec![msp(), lower_part()])?.into_array(); + + let serialization = session + .array_serialize(&array)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialization.serialized_id, decimal_byte_parts_v2_id()); + + let restricted = ArrayContext::empty() + .with_allowed_ids([VTable::id(&DecimalByteParts)].into_iter().collect()); + let err = array + .serialize(&restricted, &session, &SerializeOptions::default()) + .expect_err("expected the permitted-encoding check to refuse the v2 format"); + assert!( + err.to_string().contains("not permitted"), + "error should name the permitted-encoding check, got: {err}" + ); + Ok(()) + } + + /// Reading back an array that already carries lower parts, and computing over it, must + /// always work: the v2 format only restricts which writers may emit it. If reading or + /// the rebuild that every compute kernel does were blocked, a session whose editions + /// predate the v2 format could not read a file written by one that includes it. + #[test] + fn compute_over_existing_lower_parts_is_not_gated() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + // Stands in for an array materialized from a file: the parts already exist. + let array = deserialize_with(1, vec![msp(), lower_part()])?.into_array(); + + let sliced = array.slice(0..2)?; + assert_eq!(sliced.execute::(&mut ctx)?.len(), 2); + Ok(()) + } + + #[rstest] + fn test_deserialize_redundant_lower_parts( + #[values(2, 3)] lower_part_count: u32, + ) -> VortexResult<()> { + let mut children = vec![buffer![0i64; 3].into_array()]; + children.extend((1..lower_part_count).map(|_| buffer![0u64; 3].into_array())); + children.push(lower_part()); + let array = deserialize_with(lower_part_count, children)?; + let expected = DecimalArray::new( + buffer![1i128, 2, 3], + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(expected, array, &mut ctx); + test_serde_round_trip(array) + } + + #[test] + fn test_deserialize_rejects_child_count_mismatch() { + // Metadata claiming a lower part that was not serialized. + assert!(deserialize_with(1, vec![msp()]).is_err()); + // Metadata claiming fewer lower parts than there are children. + assert!(deserialize_with(0, vec![msp(), lower_part()]).is_err()); + // Metadata claiming more lower parts than the encoding supports. + assert!( + deserialize_with( + 4, + vec![ + msp(), + lower_part(), + lower_part(), + lower_part(), + lower_part() + ] + ) + .is_err() + ); + } + + fn plugin_deserialize_with( + serialized_id: ArrayId, + lower_part_count: u32, + children: Vec, + ) -> VortexResult { + let metadata = DecimalBytesPartsMetadata { + zeroth_child_ptype: PType::I64 as i32, + lower_part_count, + } + .encode_to_vec(); + let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); + DecimalBytePartsPlugin.deserialize( + ArrayDeserialization::new(serialized_id, &dtype, 3, &metadata, &[], &children), + &array_session(), + ) + } + + /// Each serialized ID keeps its own contract: the frozen ID never carries lower parts, and + /// the v2 ID is never written without them. + #[rstest] + #[case::frozen_without_lower_parts(VTable::id(&DecimalByteParts), 0, vec![msp()], true)] + #[case::frozen_with_lower_parts( + VTable::id(&DecimalByteParts), + 1, + vec![msp(), lower_part()], + false + )] + #[case::v2_with_lower_parts(decimal_byte_parts_v2_id(), 1, vec![msp(), lower_part()], true)] + #[case::v2_without_lower_parts(decimal_byte_parts_v2_id(), 0, vec![msp()], false)] + #[case::unknown_id(ArrayVTable::id(&Primitive), 0, vec![msp()], false)] + fn plugin_holds_each_id_to_its_contract( + #[case] serialized_id: ArrayId, + #[case] lower_part_count: u32, + #[case] children: Vec, + #[case] accepted: bool, + ) { + let result = plugin_deserialize_with(serialized_id, lower_part_count, children); + assert_eq!(result.is_ok(), accepted, "{serialized_id}: {result:?}"); + } + + fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() + } + + fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() + } + + fn session() -> VortexSession { + let session = array_session(); + crate::initialize(&session); + session + } + + /// The wire ID the session's plugin picks for `array`. + fn serialized_id(session: &VortexSession, array: &ArrayRef) -> VortexResult { + Ok(session + .array_serialize(array)? + .ok_or_else(|| vortex_err!("byte parts arrays are serializable"))? + .serialized_id) + } + + /// A single-child array is the stable shape and is always constructible. + #[test] + fn single_child_is_always_allowed() { + assert!(DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2)).is_ok()); + assert!( + DecimalByteParts::try_new_with_lower_parts(msp(), vec![], DecimalDType::new(19, 2)) + .is_ok() + ); + } + + /// Building lower parts in memory is always allowed — reading a file requires it. What + /// changes is the serialized format, not what can be constructed. + #[test] + fn lower_parts_can_always_be_constructed() { + assert!( + DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + ) + .is_ok() + ); + } + + /// A single-child array keeps the frozen format id, byte-compatible with every reader since + /// the format froze; lower parts move the array onto the v2 format id. + #[test] + fn serialized_id_tracks_lower_parts() -> VortexResult<()> { + let session = session(); + + let flat = DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2))?.into_array(); + assert_eq!( + serialized_id(&session, &flat)?, + ArrayVTable::id(&DecimalByteParts) + ); + + let wide = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + assert_eq!(serialized_id(&session, &wide)?, decimal_byte_parts_v2_id()); + + Ok(()) + } + + /// The permitted-encoding check applies to the serialized id. A context restricted to the + /// frozen format — a writer whose enabled editions predate the v2 format — must refuse an + /// array carrying lower parts, however it was obtained. + /// + /// `ArrayParts` is public and `DecimalBytePartsData` is a public unit struct, so a caller can + /// assemble slots by hand and go straight to `Array::try_from_parts`, bypassing + /// `try_new_with_lower_parts` entirely. That back door is left open on purpose — it is the + /// same path `deserialize` uses. What must hold is that the resulting array cannot become + /// bytes under the frozen id. + #[test] + fn wide_format_is_refused_where_not_permitted() -> VortexResult<()> { + let session = session(); + + let mut slots = ArraySlots::with_capacity(2); + slots.push(Some(msp())); + slots.push(Some(lower_part())); + + // Assembling the array by hand succeeds: this is the shape a file read produces. + let array = Array::try_from_parts( + ArrayParts::new( + DecimalByteParts, + DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable), + 3, + DecimalBytePartsData, + ) + .with_slots(slots), + )? + .into_array(); + assert_eq!(array.nchildren(), 2, "expected two limbs"); + + // A context permitting only the frozen format refuses to write it. + let restricted = ArrayContext::empty() + .with_allowed_ids([ArrayVTable::id(&DecimalByteParts)].into_iter().collect()); + let err = array + .serialize(&restricted, &session, &SerializeOptions::default()) + .expect_err("expected the permitted-encoding check to refuse the v2 format"); + assert!( + err.to_string().contains("not permitted"), + "error should name the permitted-encoding check, got: {err}" + ); + + // Permitting the v2 format id is exactly what allows the same array through. + let permissive = ArrayContext::empty().with_allowed_ids( + [ + ArrayVTable::id(&DecimalByteParts), + decimal_byte_parts_v2_id(), + ArrayVTable::id(&Primitive), + ] + .into_iter() + .collect(), + ); + let serialized = array.serialize(&permissive, &session, &SerializeOptions::default())?; + assert!(!serialized.is_empty()); + assert!( + permissive.to_ids().contains(&decimal_byte_parts_v2_id()), + "the file's encoding table must carry the v2 format id" + ); + + Ok(()) + } + + #[test] + fn bare_vtable_refuses_wide_serialization() -> VortexResult<()> { + let session = array_session(); + session.arrays().register(DecimalByteParts); + let array = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + let restricted = ArrayContext::empty().with_allowed_ids( + [ + ArrayVTable::id(&DecimalByteParts), + ArrayVTable::id(&Primitive), + ] + .into_iter() + .collect(), + ); + + assert!( + array + .serialize(&restricted, &session, &SerializeOptions::default()) + .is_err(), + "bare VTable registration must not write lower parts under the frozen ID" + ); + Ok(()) + } + + #[test] + fn bare_vtable_refuses_lower_parts_on_frozen_id() -> VortexResult<()> { + let session = array_session(); + session.arrays().register(DecimalByteParts); + let array = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + let id = ArrayVTable::id(&DecimalByteParts); + let plugin = session + .arrays() + .registry() + .get(&id) + .ok_or_else(|| vortex_err!("missing decimal plugin"))?; + let children = array.children(); + + // i64 MSP and one lower part, mislabeled as the frozen format. + let parts = ArrayDeserialization::new( + id, + array.dtype(), + array.len(), + &[8, 7, 16, 1], + &[], + &children, + ); + assert!(plugin.deserialize(parts, &session).is_err()); + Ok(()) + } + + #[rstest] + #[case::vtable(false)] + #[case::plugin(true)] + fn frozen_serde_is_compatible(#[case] use_plugin: bool) -> VortexResult<()> { + let session = array_session(); + if use_plugin { + session.arrays().register(DecimalBytePartsPlugin); + } else { + session.arrays().register(DecimalByteParts); + } + let array = DecimalByteParts::try_new(msp(), DecimalDType::new(2, 0))?.into_array(); + let serialized = session + .array_serialize(&array)? + .ok_or_else(|| vortex_err!("missing decimal serialization"))?; + assert_eq!(serialized.serialized_id, ArrayVTable::id(&DecimalByteParts)); + assert_eq!(serialized.metadata, [8, 7]); + let plugin = session + .arrays() + .registry() + .get(&serialized.serialized_id) + .ok_or_else(|| vortex_err!("missing decimal plugin"))?; + let decoded = plugin.deserialize( + ArrayDeserialization::new( + serialized.serialized_id, + array.dtype(), + array.len(), + &serialized.metadata, + &[], + &serialized.children, + ), + &session, + )?; + assert_arrays_eq!(array, decoded, &mut session.create_execution_ctx()); + Ok(()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs index d2ce68f3700..950c963f7a2 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -51,3 +51,44 @@ pub(crate) fn i256_parts(values: Vec, validity: Validity) -> DecimalBytePa pub(crate) fn i256_of(high: i128, low: u128) -> i256 { i256::from_parts(low, high) } + +/// The largest unscaled value a `Decimal(38, _)` can hold: `10^38 - 1`. +const MAX_PRECISION_38: i128 = 99_999_999_999_999_999_999_999_999_999_999_999_999; + +/// The largest unscaled value a `Decimal(76, _)` can hold: `10^76 - 1`. +fn max_precision_76() -> i256 { + i256::from_i128(10).wrapping_pow(76) - i256::ONE +} + +/// Values that exercise every 64-bit window of an `i128`, both signs, and the boundaries +/// where a lower part carries into the MSP. +pub(crate) fn wide_i128_values() -> Vec { + vec![ + 0, + 1, + -1, + (1 << 64) - 1, + 1 << 64, + -(1 << 64), + -((1 << 64) + 1), + MAX_PRECISION_38, + -MAX_PRECISION_38, + 1 << 100, + ] +} + +/// Values that exercise every 64-bit window of an `i256`. +pub(crate) fn wide_i256_values() -> Vec { + vec![ + i256::ZERO, + i256::ONE, + i256::ZERO - i256::ONE, + i256_of(0, u128::MAX), + i256_of(1, 0), + i256_of(-1, 0), + i256_of(-1, u128::MAX - 1), + i256_of(1 << 64, 12345), + max_precision_76(), + i256::ZERO - max_precision_76(), + ] +} diff --git a/encodings/decimal-byte-parts/tests/format_v2.rs b/encodings/decimal-byte-parts/tests/format_v2.rs deleted file mode 100644 index 338598df032..00000000000 --- a/encodings/decimal-byte-parts/tests/format_v2.rs +++ /dev/null @@ -1,257 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! The v2 serialized format. -//! -//! Lower parts can be built and computed over freely. What changes with them is the bytes: -//! an array carrying lower parts serializes under `vortex.decimal_byte_parts_v2` rather -//! than the frozen `vortex.decimal_byte_parts` format, so a writer restricted to editions -//! without the v2 format refuses it, and a reader that predates lower parts fails with an -//! unknown-encoding error instead of misreading the children. These tests pin all of that: -//! construction always works, the serialized id tracks the parts, and the permitted-encoding -//! check applies to the serialized id. - -#![expect(clippy::tests_outside_test_module)] - -use vortex_array::ArrayContext; -use vortex_array::ArrayDeserialization; -use vortex_array::ArrayId; -use vortex_array::ArrayRef; -use vortex_array::ArrayVTable; -use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; -use vortex_array::assert_arrays_eq; -use vortex_array::dtype::DecimalDType; -use vortex_array::serde::SerializeOptions; -use vortex_array::session::ArraySessionExt; -use vortex_buffer::buffer; -use vortex_decimal_byte_parts::DecimalByteParts; -use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_session::VortexSession; - -fn msp() -> ArrayRef { - buffer![1i64, 2, 3].into_array() -} - -fn lower_part() -> ArrayRef { - buffer![1u64, 2, 3].into_array() -} - -fn session() -> VortexSession { - let session = vortex_array::array_session(); - vortex_decimal_byte_parts::initialize(&session); - session -} - -/// The wire ID the session's plugin picks for `array`. -fn serialized_id(session: &VortexSession, array: &ArrayRef) -> VortexResult { - Ok(session - .array_serialize(array)? - .ok_or_else(|| vortex_err!("byte parts arrays are serializable"))? - .serialized_id) -} - -/// A single-child array is the stable shape and is always constructible. -#[test] -fn single_child_is_always_allowed() { - assert!(DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2)).is_ok()); - assert!( - DecimalByteParts::try_new_with_lower_parts(msp(), vec![], DecimalDType::new(19, 2)).is_ok() - ); -} - -/// Building lower parts in memory is always allowed — reading a file requires it. What -/// changes is the serialized format, not what can be constructed. -#[test] -fn lower_parts_can_always_be_constructed() { - assert!( - DecimalByteParts::try_new_with_lower_parts( - msp(), - vec![lower_part()], - DecimalDType::new(38, 2), - ) - .is_ok() - ); -} - -/// A single-child array keeps the frozen format id, byte-compatible with every reader since -/// the format froze; lower parts move the array onto the v2 format id. -#[test] -fn serialized_id_tracks_lower_parts() -> VortexResult<()> { - let session = session(); - - let flat = DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2))?.into_array(); - assert_eq!( - serialized_id(&session, &flat)?, - ArrayVTable::id(&DecimalByteParts) - ); - - let wide = DecimalByteParts::try_new_with_lower_parts( - msp(), - vec![lower_part()], - DecimalDType::new(38, 2), - )? - .into_array(); - assert_eq!(serialized_id(&session, &wide)?, decimal_byte_parts_v2_id()); - - Ok(()) -} - -/// The permitted-encoding check applies to the serialized id. A context restricted to the -/// frozen format — a writer whose enabled editions predate the v2 format — must refuse an -/// array carrying lower parts, however it was obtained. -/// -/// `ArrayParts` is public and `DecimalBytePartsData` is a public unit struct, so a caller can -/// assemble slots by hand and go straight to `Array::try_from_parts`, bypassing -/// `try_new_with_lower_parts` entirely. That back door is left open on purpose — it is the -/// same path `deserialize` uses. What must hold is that the resulting array cannot become -/// bytes under the frozen id. -#[test] -fn wide_format_is_refused_where_not_permitted() -> VortexResult<()> { - use vortex_array::Array; - use vortex_array::ArrayParts; - use vortex_array::ArraySlots; - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_decimal_byte_parts::DecimalBytePartsData; - - let session = session(); - - let mut slots = ArraySlots::with_capacity(2); - slots.push(Some(msp())); - slots.push(Some(lower_part())); - - // Assembling the array by hand succeeds: this is the shape a file read produces. - let array = Array::try_from_parts( - ArrayParts::new( - DecimalByteParts, - DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable), - 3, - DecimalBytePartsData, - ) - .with_slots(slots), - )? - .into_array(); - assert_eq!(array.nchildren(), 2, "expected two limbs"); - - // A context permitting only the frozen format refuses to write it. - let restricted = ArrayContext::empty() - .with_allowed_ids([ArrayVTable::id(&DecimalByteParts)].into_iter().collect()); - let err = array - .serialize(&restricted, &session, &SerializeOptions::default()) - .expect_err("expected the permitted-encoding check to refuse the v2 format"); - assert!( - err.to_string().contains("not permitted"), - "error should name the permitted-encoding check, got: {err}" - ); - - // Permitting the v2 format id is exactly what allows the same array through. - let permissive = ArrayContext::empty().with_allowed_ids( - [ - ArrayVTable::id(&DecimalByteParts), - decimal_byte_parts_v2_id(), - ArrayVTable::id(&vortex_array::arrays::Primitive), - ] - .into_iter() - .collect(), - ); - let serialized = array.serialize(&permissive, &session, &SerializeOptions::default())?; - assert!(!serialized.is_empty()); - assert!( - permissive.to_ids().contains(&decimal_byte_parts_v2_id()), - "the file's encoding table must carry the v2 format id" - ); - - Ok(()) -} - -#[test] -fn bare_vtable_refuses_wide_serialization() -> VortexResult<()> { - let session = vortex_array::array_session(); - session.arrays().register(DecimalByteParts); - let array = DecimalByteParts::try_new_with_lower_parts( - msp(), - vec![lower_part()], - DecimalDType::new(38, 2), - )? - .into_array(); - let restricted = ArrayContext::empty().with_allowed_ids( - [ - ArrayVTable::id(&DecimalByteParts), - ArrayVTable::id(&vortex_array::arrays::Primitive), - ] - .into_iter() - .collect(), - ); - - assert!( - array - .serialize(&restricted, &session, &SerializeOptions::default()) - .is_err(), - "bare VTable registration must not write lower parts under the frozen ID" - ); - Ok(()) -} - -#[test] -fn bare_vtable_refuses_lower_parts_on_frozen_id() -> VortexResult<()> { - let session = vortex_array::array_session(); - session.arrays().register(DecimalByteParts); - let array = DecimalByteParts::try_new_with_lower_parts( - msp(), - vec![lower_part()], - DecimalDType::new(38, 2), - )? - .into_array(); - let id = ArrayVTable::id(&DecimalByteParts); - let plugin = session - .arrays() - .registry() - .get(&id) - .ok_or_else(|| vortex_err!("missing decimal plugin"))?; - let children = array.children(); - - // i64 MSP and one lower part, mislabeled as the frozen format. - let parts = ArrayDeserialization::new( - id, - array.dtype(), - array.len(), - &[8, 7, 16, 1], - &[], - &children, - ); - assert!(plugin.deserialize(parts, &session).is_err()); - Ok(()) -} - -#[test] -fn bare_vtable_keeps_frozen_serde() -> VortexResult<()> { - let session = vortex_array::array_session(); - session.arrays().register(DecimalByteParts); - let array = DecimalByteParts::try_new(msp(), DecimalDType::new(2, 0))?.into_array(); - let serialized = session - .array_serialize(&array)? - .ok_or_else(|| vortex_err!("missing decimal serialization"))?; - assert_eq!(serialized.serialized_id, ArrayVTable::id(&DecimalByteParts)); - assert_eq!(serialized.metadata, [8, 7]); - let plugin = session - .arrays() - .registry() - .get(&serialized.serialized_id) - .ok_or_else(|| vortex_err!("missing decimal plugin"))?; - let decoded = plugin.deserialize( - ArrayDeserialization::new( - serialized.serialized_id, - array.dtype(), - array.len(), - &serialized.metadata, - &[], - &serialized.children, - ), - &session, - )?; - assert_arrays_eq!(array, decoded, &mut session.create_execution_ctx()); - Ok(()) -} diff --git a/vortex-test/compat-gen/Cargo.toml b/vortex-test/compat-gen/Cargo.toml index 2a5fe657d9b..4a62aca3671 100644 --- a/vortex-test/compat-gen/Cargo.toml +++ b/vortex-test/compat-gen/Cargo.toml @@ -20,11 +20,6 @@ name = "vortex-compat" path = "src/main.rs" test = false -[features] -# Fixtures for encodings whose on-disk shape is not yet stable. Kept out of the default -# fixture set so a default build never publishes a file older readers cannot open. -unstable_encodings = ["vortex/unstable_encodings"] - [dependencies] # Vortex crates vortex = { workspace = true, features = ["files", "tokio", "zstd"] } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs index 5af7596ca7b..4d799e33e74 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs @@ -12,7 +12,6 @@ mod bytebool; mod constant; mod datetimeparts; mod decimal_byte_parts; -#[cfg(feature = "unstable_encodings")] mod decimal_byte_parts_v2; mod delta; mod dict; @@ -33,14 +32,14 @@ pub(crate) const N: usize = 1024; /// All per-encoding fixtures. pub fn fixtures() -> Vec> { - #[allow(unused_mut)] - let mut fixtures: Vec> = vec![ + vec![ Box::new(alp::AlpFixture), Box::new(alprd::AlprdFixture), Box::new(bitpacked::BitPackedFixture), Box::new(bytebool::ByteBoolFixture), Box::new(datetimeparts::DateTimePartsFixture), Box::new(decimal_byte_parts::DecimalBytePartsFixture), + Box::new(decimal_byte_parts_v2::DecimalBytePartsV2Fixture), // Re-enable this once delta is stable // Box::new(delta::DeltaFixture), Box::new(dict::DictFixture), @@ -56,8 +55,5 @@ pub fn fixtures() -> Vec> { Box::new(zstd::ZstdFixture), Box::new(zigzag::ZigZagFixture), Box::new(constant::ConstantFixture), - ]; - #[cfg(feature = "unstable_encodings")] - fixtures.push(Box::new(decimal_byte_parts_v2::DecimalBytePartsV2Fixture)); - fixtures + ] } From a1bbf1f9adf40a8edb1295b924b7bcdb2802578e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 4 Sep 2026 11:15:20 -0400 Subject: [PATCH 11/16] Tell the compressor which serialized IDs the writer may emit CascadingCompressor carries the snapshot of serialized IDs the writer may emit, filled by the file writer from the enabled editions through BtrBlocksCompressorBuilder::allow_serialized_ids. A scheme whose encoding has more than one wire format picks its compression mode from it with allows_serialized_id, the newest permitted one; without a restriction every ID is allowed. No scheme consults the set yet. This is the mechanism docs/specs/editions.md describes under compression with replacement encodings (#9779). Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 36 +++++++++- vortex-compressor/src/compressor/mod.rs | 34 ++++++++++ vortex-compressor/src/compressor/tests.rs | 81 +++++++++++++++++++++++ vortex-file/src/writer.rs | 25 +++++-- 4 files changed, 169 insertions(+), 7 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index fe8072d5e66..9701579c5c5 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -91,12 +91,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + allowed_serialized_ids: Option>, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + allowed_serialized_ids: None, } } } @@ -108,6 +110,7 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + allowed_serialized_ids: None, } } @@ -214,15 +217,33 @@ impl BtrBlocksCompressorBuilder { self } + /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier + /// call. A scheme whose encoding has several wire formats picks its compression mode from this + /// set: the newest permitted one. + /// + /// The file writer passes the serialized IDs its enabled editions permit. + pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { + self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(allowed).copied().collect(), + None => allowed.clone(), + }); + self + } + /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + let compressor = CascadingCompressor::new(self.schemes); + BtrBlocksCompressor(match self.allowed_serialized_ids { + Some(allowed) => compressor.with_allowed_serialized_ids(allowed), + None => compressor, + }) } } #[cfg(test)] mod tests { use vortex_array::VTable; + use vortex_array::arrays::Bool; use vortex_fastlanes::FoR; use super::*; @@ -287,6 +308,19 @@ mod tests { } } + /// Every serialized ID is allowed until the writer narrows the set to its editions. + #[test] + fn allowed_serialized_ids_reach_the_compressor() { + let default = BtrBlocksCompressorBuilder::default().build(); + assert!(default.0.allows_serialized_id(Bool.id())); + + let narrowed = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .build(); + assert!(narrowed.0.allows_serialized_id(FoR.id())); + assert!(!narrowed.0.allows_serialized_id(Bool.id())); + } + #[test] fn cuda_compatible_uses_fsst_for_strings() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..774513c66a7 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,9 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -46,6 +49,10 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// The serialized IDs the output may use, or `None` for no restriction. See + /// [`allows_serialized_id`](Self::allows_serialized_id). + allowed_serialized_ids: Option>, } impl CascadingCompressor { @@ -63,9 +70,36 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + allowed_serialized_ids: None, } } + /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier + /// call. + /// + /// The file writer passes the serialized IDs its enabled editions permit. A scheme whose + /// encoding has several wire formats picks its compression mode from this set, the newest + /// permitted one, before estimating or compressing. + pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { + self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(&allowed).copied().collect(), + None => allowed, + }); + self + } + + /// Returns whether the writer may emit the serialized ID `id`. + /// + /// Schemes whose encoding has several wire formats consult this to pick their compression + /// mode. Without a restriction every ID is allowed, so the newest mode is chosen. The + /// serializer still emits the oldest wire form the resulting array fits, and the + /// serialization context validates that ID. + pub fn allows_serialized_id(&self, id: ArrayId) -> bool { + self.allowed_serialized_ids + .as_ref() + .is_none_or(|allowed| allowed.contains(&id)) + } + /// Returns whether the compressor was configured with `scheme`. pub fn has_scheme(&self, scheme: SchemeId) -> bool { self.schemes diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index ec14383ce36..69afcd3e9ea 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -9,11 +9,14 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Bool; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -26,6 +29,7 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -96,6 +100,48 @@ impl Scheme for DirectRatioScheme { } } +/// What the last `FormatRecordingScheme::compress` call saw for `allows_serialized_id`. +static SEEN_FORMAT: Mutex> = Mutex::new(None); + +/// Stands in for a scheme whose encoding has several wire formats: it asks the compressor whether +/// the newer one is allowed and records the answer. +#[derive(Debug)] +struct FormatRecordingScheme; + +impl Scheme for FormatRecordingScheme { + fn scheme_name(&self) -> &'static str { + "test.format_recording" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + *SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id())); + Ok(data.array().clone()) + } +} + #[derive(Debug)] struct ImmediateAlwaysUseScheme; @@ -841,3 +887,38 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } + +#[test] +fn allowed_serialized_ids_default_to_everything_and_intersect() { + let compressor = compressor(); + assert!(compressor.allows_serialized_id(Constant.id())); + assert!(compressor.allows_serialized_id(Bool.id())); + + let restricted = + compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); + assert!(restricted.allows_serialized_id(Constant.id())); + assert!(!restricted.allows_serialized_id(Bool.id())); + + let narrowed = + restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); + assert!(narrowed.allows_serialized_id(Primitive.id())); + assert!(!narrowed.allows_serialized_id(Constant.id())); + assert!(!narrowed.allows_serialized_id(Bool.id())); +} + +/// A scheme sees the restriction through the compressor it is handed: everything is allowed until +/// the writer narrows the set to its editions. +#[test] +fn schemes_see_the_allowed_serialized_ids() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0..4096i32).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let unrestricted = CascadingCompressor::new(vec![&FormatRecordingScheme]); + unrestricted.compress(&array, &mut exec_ctx)?; + assert_eq!(*SEEN_FORMAT.lock(), Some(true)); + + let restricted = unrestricted.with_allowed_serialized_ids(HashSet::from([Primitive.id()])); + restricted.compress(&array, &mut exec_ctx)?; + assert_eq!(*SEEN_FORMAT.lock(), Some(false)); + Ok(()) +} diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index ec45653f5c1..220599733fb 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,7 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, allowed_array_encodings) = + let (array_ctx, allowed_array_encodings, allowed_serialized_ids) = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); @@ -253,7 +253,8 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings), + .retain_allowed_encodings(&allowed_array_encodings) + .allow_serialized_ids(&allowed_serialized_ids), ) .build(), }; @@ -384,10 +385,12 @@ impl VortexWriteOptions { } } +/// Returns the array context, the in-memory encodings the compressor may produce, and the +/// serialized IDs its output may use. fn new_array_context( session: &VortexSession, enforce_editions: bool, -) -> (ArrayContext, HashSet) { +) -> (ArrayContext, HashSet, HashSet) { // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated. // This is preferred for now over having an empty context here, because only the // serialised array order is deterministic. The serialisation of arrays are done @@ -406,6 +409,9 @@ fn new_array_context( .filter_map(|serialized_id| arrays.registry().get(serialized_id)) .map(|plugin| plugin.id()) .collect(); + // The compressor sees the same set, so an encoding with several wire formats produces the + // newest one the editions permit. + let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { // Only permit serialized IDs in the enabled editions. @@ -413,7 +419,7 @@ fn new_array_context( } else { array_ctx }; - (array_ctx, allowed_array_encodings) + (array_ctx, allowed_array_encodings, allowed_serialized_ids) } /// The ids of `kind` the enabled editions permit. @@ -787,10 +793,12 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings) = new_array_context(&session, true); + let (ctx, allowed_array_encodings, allowed_serialized_ids) = + new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()])); + assert_eq!(allowed_serialized_ids, HashSet::from([Primitive.id()])); Ok(()) } @@ -807,9 +815,14 @@ mod tests { ) }); - let (ctx, allowed_array_encodings) = new_array_context(&session, false); + let (ctx, allowed_array_encodings, allowed_serialized_ids) = + new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); assert_eq!(allowed_array_encodings, registered_encodings); + assert_eq!( + allowed_serialized_ids, + registered_ids.iter().copied().collect::>() + ); assert!(ctx.intern(&Bool.id()).is_some()); } From 565ca086dec15586a72bdd05571410d26eb4f0f8 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 10:47:58 -0400 Subject: [PATCH 12/16] Enable a scheme when any of its serialized IDs is permitted Scheme::produced_encodings now names the serialized IDs a scheme may write its output under, oldest first. BtrBlocksCompressorBuilder::allow_serialized_ids replaces retain_allowed_encodings: it keeps a scheme when at least one of those IDs is permitted and hands the set to the compressor, so the writer makes one call from the serialized IDs its editions permit instead of mapping them back to in-memory encodings, which could not tell two wire formats of one encoding apart. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 97 +++++++++++++++++++++++------ vortex-compressor/src/scheme/mod.rs | 12 ++-- vortex-file/src/writer.rs | 40 ++++-------- 3 files changed, 98 insertions(+), 51 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 9701579c5c5..2c9a54456b5 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -207,26 +207,22 @@ impl BtrBlocksCompressorBuilder { self } - /// Retains only schemes whose produced encodings all belong to `allowed`. + /// Restricts compression to the serialized IDs in `allowed`, intersecting with any earlier + /// call. /// - /// The file writer uses this to restrict compression to the encodings of its configured - /// editions. - pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); - self - } - - /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier - /// call. A scheme whose encoding has several wire formats picks its compression mode from this - /// set: the newest permitted one. + /// A scheme stays when at least one of its [produced IDs](Scheme::produced_encodings) is + /// permitted, and the compressor is handed the set so a scheme whose encoding has several + /// wire formats picks its compression mode from it: the newest permitted one. /// /// The file writer passes the serialized IDs its enabled editions permit. pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { - self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + let allowed: HashSet = match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(allowed).copied().collect(), None => allowed.clone(), - }); + }; + self.schemes + .retain(|s| s.produced_encodings().iter().any(|id| allowed.contains(id))); + self.allowed_serialized_ids = Some(allowed); self } @@ -242,11 +238,20 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; use vortex_array::VTable; use vortex_array::arrays::Bool; + use vortex_array::arrays::Primitive; + use vortex_compressor::scheme::CompressionEstimate; + use vortex_compressor::scheme::EstimateVerdict; + use vortex_error::VortexResult; use vortex_fastlanes::FoR; use super::*; + use crate::ArrayAndStats; + use crate::CompressorContext; #[test] fn empty_starts_with_no_schemes() { @@ -261,26 +266,80 @@ mod tests { } #[test] - fn retain_allowed_encodings_filters_schemes() { + fn allow_serialized_ids_filters_schemes() { let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); assert_eq!(builder.schemes.len(), 1); assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); + let none = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&HashSet::new()); assert!(none.schemes.is_empty()); } #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { + fn allowing_all_declared_outputs_keeps_every_scheme() { let allowed: HashSet = ALL_SCHEMES .iter() .flat_map(|scheme| scheme.produced_encodings()) .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } + /// Stands in for a scheme whose encoding has two wire formats. + #[derive(Debug)] + struct TwoFormatScheme; + + impl Scheme for TwoFormatScheme { + fn scheme_name(&self) -> &'static str { + "test.two_formats" + } + + fn matches(&self, _canonical: &Canonical) -> bool { + false + } + + fn produced_encodings(&self) -> Vec { + vec![FoR.id(), Bool.id()] + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Skip) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper never matches") + } + } + + /// A scheme with several wire formats stays while any of them is permitted; which one it + /// produces is decided when compressing. + #[test] + fn any_permitted_format_keeps_the_scheme() { + static TWO_FORMATS: TwoFormatScheme = TwoFormatScheme; + + let newer_only = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&TWO_FORMATS) + .allow_serialized_ids(&HashSet::from([Bool.id()])); + assert_eq!(newer_only.schemes.len(), 1); + + let neither = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&TWO_FORMATS) + .allow_serialized_ids(&HashSet::from([Primitive.id()])); + assert!(neither.schemes.is_empty()); + } + #[test] fn cuda_compatible_excludes_alprd() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index de9e67690d4..fa42231d422 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,11 +124,15 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// The array encodings this scheme itself may introduce into its compressed output. + /// The serialized IDs this scheme may write its output under. /// - /// Cascaded children are compressed by other schemes, which declare their own encodings, - /// so only encodings constructed directly by [`compress`](Scheme::compress) belong here. - /// Canonical arrays the scheme merely rearranges do not need to be declared. + /// Cascaded children are compressed by other schemes, which declare their own IDs, so only + /// arrays constructed directly by [`compress`](Scheme::compress) belong here. Canonical + /// arrays the scheme merely rearranges do not need to be declared. + /// + /// An encoding with several wire formats lists every one of them, oldest first. The writer + /// keeps the scheme while any of them is permitted, and the scheme picks the newest + /// permitted one as its compression mode. fn produced_encodings(&self) -> Vec; /// Returns the stats generation options this scheme requires. The compressor merges all diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 220599733fb..79d485f9b1f 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,7 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, allowed_array_encodings, allowed_serialized_ids) = + let (array_ctx, allowed_serialized_ids) = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); @@ -253,7 +253,6 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings) .allow_serialized_ids(&allowed_serialized_ids), ) .build(), @@ -385,12 +384,11 @@ impl VortexWriteOptions { } } -/// Returns the array context, the in-memory encodings the compressor may produce, and the -/// serialized IDs its output may use. +/// Returns the array context and the serialized IDs the compressor may write its output under. fn new_array_context( session: &VortexSession, enforce_editions: bool, -) -> (ArrayContext, HashSet, HashSet) { +) -> (ArrayContext, HashSet) { // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated. // This is preferred for now over having an empty context here, because only the // serialised array order is deterministic. The serialisation of arrays are done @@ -404,13 +402,8 @@ fn new_array_context( .registry() .read(|registry| registry.keys().copied().collect()) }; - let allowed_array_encodings = serialized_ids - .iter() - .filter_map(|serialized_id| arrays.registry().get(serialized_id)) - .map(|plugin| plugin.id()) - .collect(); - // The compressor sees the same set, so an encoding with several wire formats produces the - // newest one the editions permit. + // The compressor sees the same set: it keeps the schemes that can write one of these IDs, and + // an encoding with several wire formats produces the newest one permitted. let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { @@ -419,7 +412,7 @@ fn new_array_context( } else { array_ctx }; - (array_ctx, allowed_array_encodings, allowed_serialized_ids) + (array_ctx, allowed_serialized_ids) } /// The ids of `kind` the enabled editions permit. @@ -793,11 +786,9 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings, allowed_serialized_ids) = - new_array_context(&session, true); + let (ctx, allowed_serialized_ids) = new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); - assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()])); assert_eq!(allowed_serialized_ids, HashSet::from([Primitive.id()])); Ok(()) } @@ -805,20 +796,13 @@ mod tests { #[test] fn disabling_editions_allows_all_registered_array_ids() { let session = array_session(); - let (registered_ids, registered_encodings) = session.arrays().registry().read(|registry| { - ( - registry.keys().copied().sorted().collect::>(), - registry - .values() - .map(|plugin| plugin.id()) - .collect::>(), - ) - }); + let registered_ids = session + .arrays() + .registry() + .read(|registry| registry.keys().copied().sorted().collect::>()); - let (ctx, allowed_array_encodings, allowed_serialized_ids) = - new_array_context(&session, false); + let (ctx, allowed_serialized_ids) = new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); - assert_eq!(allowed_array_encodings, registered_encodings); assert_eq!( allowed_serialized_ids, registered_ids.iter().copied().collect::>() From 6721ceaef80ad393e3ced7d63c20c8472411eb46 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 10:49:53 -0400 Subject: [PATCH 13/16] Read the permitted serialized IDs from the CompressorContext The compressor seeds each root CompressorContext with its permitted serialized IDs and every descent inherits them, so a scheme asks compress_ctx.allows_serialized_id both while estimating and while compressing and picks the same mode in both. The per-compressor accessor goes; allowed_serialized_ids remains for inspection. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 8 +- vortex-compressor/src/compressor/cascade.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 38 +++---- vortex-compressor/src/compressor/tests.rs | 110 ++++++++++++++------ vortex-compressor/src/scheme/ctx.rs | 27 ++++- 5 files changed, 127 insertions(+), 58 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 2c9a54456b5..ad25969a028 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -371,13 +371,15 @@ mod tests { #[test] fn allowed_serialized_ids_reach_the_compressor() { let default = BtrBlocksCompressorBuilder::default().build(); - assert!(default.0.allows_serialized_id(Bool.id())); + assert!(default.0.allowed_serialized_ids().is_none()); let narrowed = BtrBlocksCompressorBuilder::default() .allow_serialized_ids(&HashSet::from([FoR.id()])) .build(); - assert!(narrowed.0.allows_serialized_id(FoR.id())); - assert!(!narrowed.0.allows_serialized_id(Bool.id())); + assert_eq!( + narrowed.0.allowed_serialized_ids(), + Some(&HashSet::from([FoR.id()])) + ); } #[test] diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..dd98f4ea3c6 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -59,7 +59,7 @@ impl CascadingCompressor { let canonical = array.clone().execute::(exec_ctx)?.0; let compact = canonical.compact(exec_ctx)?; - let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + let compressed = self.compress_canonical(compact, self.root_context(), exec_ctx)?; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 774513c66a7..159842f2595 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,11 +9,14 @@ mod sample; mod select; mod structural; +use std::sync::Arc; + use vortex_array::ArrayId; use vortex_utils::aliases::hash_set::HashSet; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; +use crate::scheme::CompressorContext; use crate::scheme::DescendantExclusion; use crate::scheme::Scheme; use crate::scheme::SchemeExt; @@ -50,9 +53,9 @@ pub struct CascadingCompressor { /// list offsets). root_exclusions: Vec, - /// The serialized IDs the output may use, or `None` for no restriction. See - /// [`allows_serialized_id`](Self::allows_serialized_id). - allowed_serialized_ids: Option>, + /// The serialized IDs the writer may emit, or `None` for no restriction. Seeds every root + /// [`CompressorContext`], where schemes read it. + allowed_serialized_ids: Option>>, } impl CascadingCompressor { @@ -77,27 +80,26 @@ impl CascadingCompressor { /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier /// call. /// - /// The file writer passes the serialized IDs its enabled editions permit. A scheme whose - /// encoding has several wire formats picks its compression mode from this set, the newest - /// permitted one, before estimating or compressing. + /// The file writer passes the serialized IDs its enabled editions permit. Schemes read the + /// set through [`CompressorContext::allows_serialized_id`], so a scheme whose encoding has + /// several wire formats picks the newest permitted one as its mode, while estimating and + /// while compressing alike. pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { - self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + self.allowed_serialized_ids = Some(Arc::new(match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(&allowed).copied().collect(), None => allowed, - }); + })); self } - /// Returns whether the writer may emit the serialized ID `id`. - /// - /// Schemes whose encoding has several wire formats consult this to pick their compression - /// mode. Without a restriction every ID is allowed, so the newest mode is chosen. The - /// serializer still emits the oldest wire form the resulting array fits, and the - /// serialization context validates that ID. - pub fn allows_serialized_id(&self, id: ArrayId) -> bool { - self.allowed_serialized_ids - .as_ref() - .is_none_or(|allowed| allowed.contains(&id)) + /// The serialized IDs the writer may emit, or `None` when unrestricted. + pub fn allowed_serialized_ids(&self) -> Option<&HashSet> { + self.allowed_serialized_ids.as_deref() + } + + /// The context a compress call starts from. + pub(crate) fn root_context(&self) -> CompressorContext { + CompressorContext::new(self.allowed_serialized_ids.clone()) } /// Returns whether the compressor was configured with `scheme`. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 69afcd3e9ea..42a98245fef 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -132,12 +132,12 @@ impl Scheme for FormatRecordingScheme { fn compress( &self, - compressor: &CascadingCompressor, + _compressor: &CascadingCompressor, data: &ArrayAndStats, - _compress_ctx: CompressorContext, + compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - *SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id())); + *SEEN_FORMAT.lock() = Some(compress_ctx.allows_serialized_id(Constant.id())); Ok(data.array().clone()) } } @@ -420,8 +420,12 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -438,8 +442,12 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -456,8 +464,12 @@ fn callback_skip_is_ignored() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -474,8 +486,12 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -492,8 +508,12 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -510,8 +530,12 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -528,8 +552,12 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(winner.is_none()); Ok(()) @@ -630,8 +658,12 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -651,7 +683,7 @@ fn threshold_reflects_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -672,7 +704,7 @@ fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner // `None`) because the zero-byte sample is never stored as the best. @@ -691,7 +723,7 @@ fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert_eq!(observed, Some(None)); @@ -711,7 +743,7 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -732,8 +764,12 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -780,7 +816,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // A context with default stats_options (count_distinct_values = false) and // marked as a sample so the function skips the sampling step and compresses // the array directly. - let ctx = CompressorContext::new().with_sampling(); + let ctx = CompressorContext::new(None).with_sampling(); // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" @@ -891,19 +927,27 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { #[test] fn allowed_serialized_ids_default_to_everything_and_intersect() { let compressor = compressor(); - assert!(compressor.allows_serialized_id(Constant.id())); - assert!(compressor.allows_serialized_id(Bool.id())); + let root = compressor.root_context(); + assert!(root.allows_serialized_id(Constant.id())); + assert!(root.allows_serialized_id(Bool.id())); let restricted = compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); - assert!(restricted.allows_serialized_id(Constant.id())); - assert!(!restricted.allows_serialized_id(Bool.id())); + let root = restricted.root_context(); + assert!(root.allows_serialized_id(Constant.id())); + assert!(!root.allows_serialized_id(Bool.id())); let narrowed = restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); - assert!(narrowed.allows_serialized_id(Primitive.id())); - assert!(!narrowed.allows_serialized_id(Constant.id())); - assert!(!narrowed.allows_serialized_id(Bool.id())); + let root = narrowed.root_context(); + assert!(root.allows_serialized_id(Primitive.id())); + assert!(!root.allows_serialized_id(Constant.id())); + assert!(!root.allows_serialized_id(Bool.id())); + + // Descending keeps the set. + let child = root.descend_with_scheme(IntDictScheme.id(), 0); + assert!(child.allows_serialized_id(Primitive.id())); + assert!(!child.allows_serialized_id(Constant.id())); } /// A scheme sees the restriction through the compressor it is handed: everything is allowed until diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..83685031c33 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,8 +4,11 @@ //! Compression context for recursive compression. use std::fmt; +use std::sync::Arc; +use vortex_array::ArrayId; use vortex_error::VortexExpect; +use vortex_utils::aliases::hash_set::HashSet; use crate::compressor::ROOT_SCHEME_ID; use crate::scheme::SchemeId; @@ -38,18 +41,24 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, + + /// The serialized IDs the writer may emit, or `None` for no restriction. Shared by every + /// context of one compress call, so cloning at each descent is a pointer bump. + allowed_serialized_ids: Option>>, } impl CompressorContext { - /// Creates a new `CompressorContext`. + /// Creates a new root `CompressorContext` for a compressor that may emit the given serialized + /// IDs, or any ID when `None`. /// /// This should **only** be created by the compressor. - pub(crate) fn new() -> Self { + pub(crate) fn new(allowed_serialized_ids: Option>>) -> Self { Self { is_sample: false, allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), + allowed_serialized_ids, } } } @@ -57,7 +66,7 @@ impl CompressorContext { #[cfg(test)] impl Default for CompressorContext { fn default() -> Self { - Self::new() + Self::new(None) } } @@ -67,6 +76,18 @@ impl CompressorContext { self.is_sample } + /// Returns whether the writer may emit the serialized ID `id`. + /// + /// A scheme whose encoding has several wire formats picks its compression mode from this, + /// the newest permitted one, and the same answer is available while estimating and while + /// compressing. Without a restriction every ID is allowed. The serializer still emits the + /// oldest wire form the resulting array fits, and the serialization context validates it. + pub fn allows_serialized_id(&self, id: ArrayId) -> bool { + self.allowed_serialized_ids + .as_ref() + .is_none_or(|allowed| allowed.contains(&id)) + } + /// Returns the merged stats generation options for this compression site. pub fn merged_stats_options(&self) -> GenerateStatsOptions { self.merged_stats_options From a618e33c959e490b7adb25f5bf95908293bec071 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 14:25:55 -0400 Subject: [PATCH 14/16] add predecessor schemes and make produced_encodings refer to serialized ids Signed-off-by: Matt Katz --- .../src/{builder.rs => builder/mod.rs} | 217 +++----------- vortex-btrblocks/src/builder/tests.rs | 216 ++++++++++++++ .../src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- .../src/schemes/string/zstd_buffers.rs | 2 +- vortex-compressor/src/compressor/cascade.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 98 +++++-- vortex-compressor/src/compressor/select.rs | 21 +- vortex-compressor/src/compressor/tests.rs | 171 ++--------- .../src/compressor/version_tests.rs | 266 ++++++++++++++++++ vortex-compressor/src/scheme/ctx.rs | 27 +- vortex-compressor/src/scheme/exclusion.rs | 6 +- vortex-compressor/src/scheme/mod.rs | 22 +- vortex-file/src/tests.rs | 35 +++ vortex-file/src/writer.rs | 2 - 16 files changed, 686 insertions(+), 405 deletions(-) rename vortex-btrblocks/src/{builder.rs => builder/mod.rs} (58%) create mode 100644 vortex-btrblocks/src/builder/tests.rs create mode 100644 vortex-compressor/src/compressor/version_tests.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder/mod.rs similarity index 58% rename from vortex-btrblocks/src/builder.rs rename to vortex-btrblocks/src/builder/mod.rs index ad25969a028..ea0ef849dd2 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -18,7 +18,7 @@ use crate::schemes::integer; use crate::schemes::string; use crate::schemes::temporal; -/// All available compression schemes. +/// The newest versions of all available compression schemes. /// /// This list is order-sensitive: the builder preserves this order when constructing /// the final scheme list, so that tie-breaking is deterministic. @@ -117,7 +117,7 @@ impl BtrBlocksCompressorBuilder { /// Adds an external compression scheme not in [`ALL_SCHEMES`]. /// /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes - /// with the compressor. + /// with the compressor. Register only the newest version of a scheme. /// /// # Panics /// @@ -201,32 +201,53 @@ impl BtrBlocksCompressorBuilder { } /// Removes the specified compression schemes by their [`SchemeId`]. + /// + /// An ID anywhere in a registered predecessor chain removes the entire chain. + /// + /// # Panics + /// + /// Panics if a traversed predecessor chain contains a cycle. pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { let ids: HashSet<_> = ids.into_iter().collect(); - self.schemes.retain(|s| !ids.contains(&s.id())); + self.schemes.retain(|scheme| { + let mut seen = HashSet::new(); + let mut candidate = Some(*scheme); + while let Some(version) = candidate { + assert!( + seen.insert(version.id()), + "cycle in scheme predecessor chain" + ); + if ids.contains(&version.id()) { + return false; + } + candidate = version.predecessor(); + } + true + }); self } /// Restricts compression to the serialized IDs in `allowed`, intersecting with any earlier /// call. /// - /// A scheme stays when at least one of its [produced IDs](Scheme::produced_encodings) is - /// permitted, and the compressor is handed the set so a scheme whose encoding has several - /// wire formats picks its compression mode from it: the newest permitted one. - /// - /// The file writer passes the serialized IDs its enabled editions permit. + /// At build time, each scheme is replaced by the newest version in its predecessor chain + /// whose [`required_serialized_ids`](Scheme::required_serialized_ids) are all permitted. + /// Schemes with no eligible version are removed. This also applies to schemes added after + /// this call. The file writer passes the serialized IDs its enabled editions permit. pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { let allowed: HashSet = match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(allowed).copied().collect(), None => allowed.clone(), }; - self.schemes - .retain(|s| s.produced_encodings().iter().any(|id| allowed.contains(id))); self.allowed_serialized_ids = Some(allowed); self } /// Builds the configured [`BtrBlocksCompressor`]. + /// + /// # Panics + /// + /// Panics if predecessor chains contain a cycle or share a scheme ID. pub fn build(self) -> BtrBlocksCompressor { let compressor = CascadingCompressor::new(self.schemes); BtrBlocksCompressor(match self.allowed_serialized_ids { @@ -237,178 +258,4 @@ impl BtrBlocksCompressorBuilder { } #[cfg(test)] -mod tests { - use vortex_array::ArrayRef; - use vortex_array::Canonical; - use vortex_array::ExecutionCtx; - use vortex_array::VTable; - use vortex_array::arrays::Bool; - use vortex_array::arrays::Primitive; - use vortex_compressor::scheme::CompressionEstimate; - use vortex_compressor::scheme::EstimateVerdict; - use vortex_error::VortexResult; - use vortex_fastlanes::FoR; - - use super::*; - use crate::ArrayAndStats; - use crate::CompressorContext; - - #[test] - fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(); - assert!(builder.schemes.is_empty()); - } - - #[test] - fn default_includes_all_schemes() { - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn allow_serialized_ids_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - - let none = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn allowing_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - /// Stands in for a scheme whose encoding has two wire formats. - #[derive(Debug)] - struct TwoFormatScheme; - - impl Scheme for TwoFormatScheme { - fn scheme_name(&self) -> &'static str { - "test.two_formats" - } - - fn matches(&self, _canonical: &Canonical) -> bool { - false - } - - fn produced_encodings(&self) -> Vec { - vec![FoR.id(), Bool.id()] - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::Skip) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper never matches") - } - } - - /// A scheme with several wire formats stays while any of them is permitted; which one it - /// produces is decided when compressing. - #[test] - fn any_permitted_format_keeps_the_scheme() { - static TWO_FORMATS: TwoFormatScheme = TwoFormatScheme; - - let newer_only = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&TWO_FORMATS) - .allow_serialized_ids(&HashSet::from([Bool.id()])); - assert_eq!(newer_only.schemes.len(), 1); - - let neither = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&TWO_FORMATS) - .allow_serialized_ids(&HashSet::from([Primitive.id()])); - assert!(neither.schemes.is_empty()); - } - - #[test] - fn cuda_compatible_excludes_alprd() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - !builder - .schemes - .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) - ); - } - - /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. - #[test] - fn cuda_compatible_excludes_every_sparse_scheme() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - for excluded in [ - integer::SparseScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - ] { - assert!( - !builder.schemes.iter().any(|s| s.id() == excluded), - "{excluded} should be excluded" - ); - } - } - - /// Every serialized ID is allowed until the writer narrows the set to its editions. - #[test] - fn allowed_serialized_ids_reach_the_compressor() { - let default = BtrBlocksCompressorBuilder::default().build(); - assert!(default.0.allowed_serialized_ids().is_none()); - - let narrowed = BtrBlocksCompressorBuilder::default() - .allow_serialized_ids(&HashSet::from([FoR.id()])) - .build(); - assert_eq!( - narrowed.0.allowed_serialized_ids(), - Some(&HashSet::from([FoR.id()])) - ); - } - - #[test] - fn cuda_compatible_uses_fsst_for_strings() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::FSSTScheme.id()) - ); - #[cfg(feature = "zstd")] - assert!( - !builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::ZstdScheme.id()) - ); - } - - #[test] - #[cfg(feature = "pco")] - fn cuda_compatible_excludes_pco() { - let builder = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme) - .only_cuda_compatible(); - for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { - assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); - } - } -} +mod tests; diff --git a/vortex-btrblocks/src/builder/tests.rs b/vortex-btrblocks/src/builder/tests.rs new file mode 100644 index 00000000000..1b736ba895b --- /dev/null +++ b/vortex-btrblocks/src/builder/tests.rs @@ -0,0 +1,216 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::VarBin; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexResult; +use vortex_fastlanes::FoR; +use vortex_fsst::FSST; +use vortex_session::registry::CachedId; + +use super::*; +use crate::ArrayAndStats; +use crate::CompressorContext; + +#[test] +fn empty_starts_with_no_schemes() { + assert!(BtrBlocksCompressorBuilder::empty().schemes.is_empty()); +} + +#[test] +fn default_includes_all_schemes() { + assert_eq!( + BtrBlocksCompressorBuilder::default().schemes.len(), + ALL_SCHEMES.len() + ); +} + +#[test] +fn allowed_serialized_ids_filter_schemes_at_build() { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .build(); + for scheme in ALL_SCHEMES { + assert_eq!( + compressor.has_scheme(scheme.id()), + scheme.id() == integer::FoRScheme.id() + ); + } +} + +#[test] +fn allowing_all_declared_outputs_keeps_every_scheme() { + let allowed = ALL_SCHEMES + .iter() + .flat_map(|s| s.produced_encodings()) + .collect(); + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&allowed) + .build(); + for scheme in ALL_SCHEMES { + assert!(compressor.has_scheme(scheme.id())); + } +} + +#[rstest] +#[case::neither(vec![], false)] +#[case::fsst_only(vec![FSST.id()], false)] +#[case::varbin_only(vec![VarBin.id()], false)] +#[case::both(vec![FSST.id(), VarBin.id()], true)] +fn all_required_outputs_must_be_allowed(#[case] allowed: Vec, #[case] expected: bool) { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&allowed.into_iter().collect()) + .build(); + assert_eq!(compressor.has_scheme(string::FSSTScheme.id()), expected); +} + +#[rstest] +#[case::forbidden(HashSet::new(), false)] +#[case::permitted(HashSet::from([FoR.id()]), true)] +fn restriction_applies_to_schemes_added_later( + #[case] allowed: HashSet, + #[case] expected: bool, +) { + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_serialized_ids(&allowed) + .with_new_scheme(&integer::FoRScheme) + .build(); + assert_eq!(compressor.has_scheme(integer::FoRScheme.id()), expected); +} + +#[test] +fn repeated_restrictions_intersect() { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id(), FSST.id()])) + .allow_serialized_ids(&HashSet::from([FSST.id(), VarBin.id()])) + .build(); + assert!(!compressor.has_scheme(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(string::FSSTScheme.id())); +} + +#[test] +fn cuda_compatible_excludes_alprd() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + !builder + .schemes + .iter() + .any(|s| s.id() == float::ALPRDScheme.id()) + ); +} + +/// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. +#[test] +fn cuda_compatible_excludes_every_sparse_scheme() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + for excluded in [ + integer::SparseScheme.id(), + float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), + ] { + assert!( + !builder.schemes.iter().any(|s| s.id() == excluded), + "{excluded} should be excluded" + ); + } +} + +#[test] +fn cuda_compatible_uses_fsst_for_strings() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::FSSTScheme.id()) + ); + #[cfg(feature = "zstd")] + assert!( + !builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::ZstdScheme.id()) + ); +} + +#[test] +#[cfg(feature = "pco")] +fn cuda_compatible_excludes_pco() { + let builder = BtrBlocksCompressorBuilder::default() + .with_new_scheme(&integer::PcoScheme) + .with_new_scheme(&float::PcoScheme) + .only_cuda_compatible(); + for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { + assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); + } +} + +static FOR_V2_ID: CachedId = CachedId::new("test.for_v2"); + +#[derive(Debug)] +struct NewFoRScheme; + +impl Scheme for NewFoRScheme { + fn scheme_name(&self) -> &'static str { + "test.for_v2" + } + + fn matches(&self, canonical: &Canonical) -> bool { + integer::FoRScheme.matches(canonical) + } + + fn produced_encodings(&self) -> Vec { + vec![*FOR_V2_ID] + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + Some(&integer::FoRScheme) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Skip) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +#[test] +fn restrictions_select_predecessors_of_schemes_added_later() { + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .with_new_scheme(&NewFoRScheme) + .build(); + assert!(compressor.has_scheme(integer::FoRScheme.id())); + assert!(compressor.has_scheme(NewFoRScheme.id())); +} + +#[rstest] +#[case::old(integer::FoRScheme.id())] +#[case::new(NewFoRScheme.id())] +fn excluding_any_version_removes_the_chain(#[case] excluded: SchemeId) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&NewFoRScheme) + .exclude_schemes([excluded]) + .build(); + assert!(!compressor.has_scheme(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(NewFoRScheme.id())); +} diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 3f06d65b061..5204e4e2478 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 46b2f1e302e..86b69bba47e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index a1bc8643775..dddaa349de1 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index cf691c70fcb..98d5feee2d9 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index dd98f4ea3c6..ecfd3c2c542 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -93,7 +93,7 @@ impl CascadingCompressor { let child_ctx = parent_ctx .clone() - .descend_with_scheme(parent_id, child_index); + .descend_with_scheme(self.resolve_scheme_id(parent_id), child_index); self.compress_canonical(compact, child_ctx, exec_ctx) } diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 159842f2595..03733aca5e1 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,9 +9,8 @@ mod sample; mod select; mod structural; -use std::sync::Arc; - use vortex_array::ArrayId; +use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; use crate::builtins::IntDictScheme; @@ -53,16 +52,38 @@ pub struct CascadingCompressor { /// list offsets). root_exclusions: Vec, - /// The serialized IDs the writer may emit, or `None` for no restriction. Seeds every root - /// [`CompressorContext`], where schemes read it. - allowed_serialized_ids: Option>>, + /// Maps every registered version to the version selected for compression. + scheme_aliases: HashMap, + + /// Configuration only: retained so repeated restrictions intersect exactly. + allowed_serialized_ids: Option>, } impl CascadingCompressor { /// Creates a new compressor with the given schemes. /// + /// Register only the newest version of each scheme. Predecessor IDs are aliases for the + /// selected version in exclusions and [`has_scheme`](Self::has_scheme) checks. /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. + /// + /// # Panics + /// + /// Panics if predecessor chains contain a cycle or share a scheme ID, including when multiple + /// versions of the same scheme are registered separately. pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + let mut scheme_aliases = HashMap::new(); + for &scheme in &schemes { + let mut candidate = Some(scheme); + while let Some(version) = candidate { + assert!( + scheme_aliases.insert(version.id(), scheme.id()).is_none(), + "scheme {} appears more than once in the registered predecessor chains", + version.id(), + ); + candidate = version.predecessor(); + } + } + // Root exclusion: exclude IntDict from list/listview offsets (monotonically // increasing data where dictionary encoding is wasteful). let root_exclusions = vec![DescendantExclusion { @@ -73,40 +94,68 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + scheme_aliases, allowed_serialized_ids: None, } } - /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier - /// call. + /// Selects the newest eligible version of each scheme, intersecting with any earlier call. /// - /// The file writer passes the serialized IDs its enabled editions permit. Schemes read the - /// set through [`CompressorContext::allows_serialized_id`], so a scheme whose encoding has - /// several wire formats picks the newest permitted one as its mode, while estimating and - /// while compressing alike. + /// A version is eligible only when all of its [`Scheme::required_serialized_ids`] are allowed. + /// Otherwise its predecessors are tried in order; the scheme is removed if none is eligible. + /// Selection preserves registration order and happens before any compression or estimation. pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { - self.allowed_serialized_ids = Some(Arc::new(match self.allowed_serialized_ids.take() { + let allowed = match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(&allowed).copied().collect(), None => allowed, - })); + }; + let mut replacements = HashMap::new(); + self.schemes = self + .schemes + .into_iter() + .filter_map(|scheme| { + let mut candidate = Some(scheme); + while let Some(version) = candidate { + if version + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + { + replacements.insert(scheme.id(), version.id()); + return Some(version); + } + candidate = version.predecessor(); + } + None + }) + .collect(); + self.scheme_aliases.retain(|_, selected| { + if let Some(replacement) = replacements.get(selected) { + *selected = *replacement; + true + } else { + false + } + }); + self.allowed_serialized_ids = Some(allowed); self } - /// The serialized IDs the writer may emit, or `None` when unrestricted. - pub fn allowed_serialized_ids(&self) -> Option<&HashSet> { - self.allowed_serialized_ids.as_deref() - } - /// The context a compress call starts from. pub(crate) fn root_context(&self) -> CompressorContext { - CompressorContext::new(self.allowed_serialized_ids.clone()) + CompressorContext::new() } - /// Returns whether the compressor was configured with `scheme`. + /// Returns whether a version of `scheme` is enabled. + /// + /// Any ID in a registered predecessor chain refers to the selected version, including when + /// the selected version is older or newer than the specified ID. pub fn has_scheme(&self, scheme: SchemeId) -> bool { - self.schemes - .iter() - .any(|candidate| candidate.id() == scheme) + self.scheme_aliases.contains_key(&scheme) + } + + fn resolve_scheme_id(&self, scheme: SchemeId) -> SchemeId { + self.scheme_aliases.get(&scheme).copied().unwrap_or(scheme) } } @@ -114,3 +163,6 @@ impl CascadingCompressor { #[cfg(test)] mod tests; + +#[cfg(test)] +mod version_tests; diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..c492729f77c 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -152,10 +152,9 @@ impl CascadingCompressor { // The root entry is always first in the history (if present). Check if the root has // excluded us. if let Some((_, child_idx)) = iter.next_if(|&(sid, _)| sid == ROOT_SCHEME_ID) - && self - .root_exclusions - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + && self.root_exclusions.iter().any(|rule| { + self.resolve_scheme_id(rule.excluded) == id && rule.children.contains(child_idx) + }) { return true; } @@ -163,10 +162,9 @@ impl CascadingCompressor { // Push rules: Check if any of our ancestors have excluded us. for (ancestor_id, child_idx) in iter { if let Some(ancestor) = self.schemes.iter().find(|s| s.id() == ancestor_id) - && ancestor - .descendant_exclusions() - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + && ancestor.descendant_exclusions().iter().any(|rule| { + self.resolve_scheme_id(rule.excluded) == id && rule.children.contains(child_idx) + }) { return true; } @@ -174,10 +172,9 @@ impl CascadingCompressor { // Pull rules: Check if we have excluded ourselves because of our ancestors. for rule in candidate.ancestor_exclusions() { - if history - .iter() - .any(|(sid, cidx)| *sid == rule.ancestor && rule.children.contains(*cidx)) - { + if history.iter().any(|(sid, cidx)| { + *sid == self.resolve_scheme_id(rule.ancestor) && rule.children.contains(*cidx) + }) { return true; } } diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 42a98245fef..ec14383ce36 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -9,14 +9,11 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::VTable; use vortex_array::VortexSessionExecute; -use vortex_array::arrays::Bool; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; -use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -29,7 +26,6 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -100,48 +96,6 @@ impl Scheme for DirectRatioScheme { } } -/// What the last `FormatRecordingScheme::compress` call saw for `allows_serialized_id`. -static SEEN_FORMAT: Mutex> = Mutex::new(None); - -/// Stands in for a scheme whose encoding has several wire formats: it asks the compressor whether -/// the newer one is allowed and records the answer. -#[derive(Debug)] -struct FormatRecordingScheme; - -impl Scheme for FormatRecordingScheme { - fn scheme_name(&self) -> &'static str { - "test.format_recording" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn produced_encodings(&self) -> Vec { - Vec::new() - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - *SEEN_FORMAT.lock() = Some(compress_ctx.allows_serialized_id(Constant.id())); - Ok(data.array().clone()) - } -} - #[derive(Debug)] struct ImmediateAlwaysUseScheme; @@ -420,12 +374,8 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -442,12 +392,8 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -464,12 +410,8 @@ fn callback_skip_is_ignored() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -486,12 +428,8 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -508,12 +446,8 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -530,12 +464,8 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -552,12 +482,8 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(winner.is_none()); Ok(()) @@ -658,12 +584,8 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -683,7 +605,7 @@ fn threshold_reflects_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -704,7 +626,7 @@ fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner // `None`) because the zero-byte sample is never stored as the best. @@ -723,7 +645,7 @@ fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert_eq!(observed, Some(None)); @@ -743,7 +665,7 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -764,12 +686,8 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -816,7 +734,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // A context with default stats_options (count_distinct_values = false) and // marked as a sample so the function skips the sampling step and compresses // the array directly. - let ctx = CompressorContext::new(None).with_sampling(); + let ctx = CompressorContext::new().with_sampling(); // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" @@ -923,46 +841,3 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } - -#[test] -fn allowed_serialized_ids_default_to_everything_and_intersect() { - let compressor = compressor(); - let root = compressor.root_context(); - assert!(root.allows_serialized_id(Constant.id())); - assert!(root.allows_serialized_id(Bool.id())); - - let restricted = - compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); - let root = restricted.root_context(); - assert!(root.allows_serialized_id(Constant.id())); - assert!(!root.allows_serialized_id(Bool.id())); - - let narrowed = - restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); - let root = narrowed.root_context(); - assert!(root.allows_serialized_id(Primitive.id())); - assert!(!root.allows_serialized_id(Constant.id())); - assert!(!root.allows_serialized_id(Bool.id())); - - // Descending keeps the set. - let child = root.descend_with_scheme(IntDictScheme.id(), 0); - assert!(child.allows_serialized_id(Primitive.id())); - assert!(!child.allows_serialized_id(Constant.id())); -} - -/// A scheme sees the restriction through the compressor it is handed: everything is allowed until -/// the writer narrows the set to its editions. -#[test] -fn schemes_see_the_allowed_serialized_ids() -> VortexResult<()> { - let array = PrimitiveArray::from_iter(0..4096i32).into_array(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let unrestricted = CascadingCompressor::new(vec![&FormatRecordingScheme]); - unrestricted.compress(&array, &mut exec_ctx)?; - assert_eq!(*SEEN_FORMAT.lock(), Some(true)); - - let restricted = unrestricted.with_allowed_serialized_ids(HashSet::from([Primitive.id()])); - restricted.compress(&array, &mut exec_ctx)?; - assert_eq!(*SEEN_FORMAT.lock(), Some(false)); - Ok(()) -} diff --git a/vortex-compressor/src/compressor/version_tests.rs b/vortex-compressor/src/compressor/version_tests.rs new file mode 100644 index 00000000000..197583331cf --- /dev/null +++ b/vortex-compressor/src/compressor/version_tests.rs @@ -0,0 +1,266 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use super::*; +use crate::scheme::AncestorExclusion; +use crate::scheme::CompressionEstimate; +use crate::scheme::EstimateVerdict; +use crate::stats::ArrayAndStats; +use crate::stats::GenerateStatsOptions; + +static V1_ID: CachedId = CachedId::new("test.version_1"); +static V2_ID: CachedId = CachedId::new("test.version_2"); +static V3_ID: CachedId = CachedId::new("test.version_3"); +static AUX_ID: CachedId = CachedId::new("test.auxiliary"); + +#[derive(Debug)] +struct TestScheme { + name: &'static str, + version: u8, + predecessor: Option<&'static dyn Scheme>, + push: Option<&'static dyn Scheme>, + pull: Option<&'static dyn Scheme>, +} + +impl TestScheme { + const fn new( + name: &'static str, + version: u8, + predecessor: Option<&'static dyn Scheme>, + ) -> Self { + Self { + name, + version, + predecessor, + push: None, + pull: None, + } + } +} + +impl Scheme for TestScheme { + fn scheme_name(&self) -> &'static str { + self.name + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_int() + } + + fn produced_encodings(&self) -> Vec { + match self.version { + 1 => vec![*V1_ID], + 2 => vec![*V2_ID, *AUX_ID], + 3 => vec![*V3_ID], + _ => vec![], + } + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + self.predecessor + } + + fn num_children(&self) -> usize { + 2 + } + + fn descendant_exclusions(&self) -> Vec { + self.push + .map(|scheme| DescendantExclusion { + excluded: scheme.id(), + children: ChildSelection::One(1), + }) + .into_iter() + .collect() + } + + fn ancestor_exclusions(&self) -> Vec { + self.pull + .map(|scheme| AncestorExclusion { + ancestor: scheme.id(), + children: ChildSelection::One(1), + }) + .into_iter() + .collect() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // Older versions would beat newer versions if they reached estimation together. + CompressionEstimate::Verdict(EstimateVerdict::Ratio(5.0 - f64::from(self.version))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +static V1: TestScheme = TestScheme::new("test.scheme_v1", 1, None); +static V2: TestScheme = TestScheme::new("test.scheme_v2", 2, Some(&V1)); +static V3: TestScheme = TestScheme::new("test.scheme_v3", 3, Some(&V2)); +static OTHER: TestScheme = TestScheme::new("test.other", 0, None); + +#[test] +fn newest_eligible_version_is_selected_before_estimation() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut exec_ctx = session.create_execution_ctx(); + let data = ArrayAndStats::new( + PrimitiveArray::from_iter(0..128i32).into_array(), + GenerateStatsOptions::default(), + ); + for (allowed, expected) in [ + (None, V3.id()), + ( + Some(HashSet::from([*V1_ID, *V2_ID, *AUX_ID, *V3_ID])), + V3.id(), + ), + (Some(HashSet::from([*V1_ID, *V2_ID, *AUX_ID])), V2.id()), + (Some(HashSet::from([*V2_ID, *AUX_ID])), V2.id()), + (Some(HashSet::from([*V1_ID, *V2_ID])), V1.id()), + (Some(HashSet::from([*V1_ID])), V1.id()), + ] { + let mut compressor = CascadingCompressor::new(vec![&V3]); + if let Some(allowed) = allowed { + compressor = compressor.with_allowed_serialized_ids(allowed); + } + assert_eq!(compressor.schemes.len(), 1); + let winner = compressor.choose_best_scheme( + &compressor.schemes, + &data, + compressor.root_context(), + &mut exec_ctx, + )?; + assert_eq!(winner.map(|(scheme, _)| scheme.id()), Some(expected)); + for version in [&V1, &V2, &V3] { + assert!(compressor.has_scheme(version.id())); + } + } + Ok(()) +} + +#[test] +fn no_eligible_version_removes_the_entire_chain() { + for allowed in [HashSet::new(), HashSet::from([*V2_ID])] { + let compressor = CascadingCompressor::new(vec![&V3]).with_allowed_serialized_ids(allowed); + assert!(compressor.schemes.is_empty()); + for version in [&V1, &V2, &V3] { + assert!(!compressor.has_scheme(version.id())); + } + } +} + +#[test] +fn fallback_preserves_registration_order() { + let compressor = CascadingCompressor::new(vec![&V3, &OTHER]) + .with_allowed_serialized_ids(HashSet::from([*V1_ID])); + assert_eq!( + compressor + .schemes + .iter() + .map(|s| s.id()) + .collect::>(), + vec![V1.id(), OTHER.id()] + ); +} + +#[test] +fn successive_restrictions_keep_aliases_and_intersect_wire_ids() { + let compressor = CascadingCompressor::new(vec![&V3]) + .with_allowed_serialized_ids(HashSet::from([*V1_ID, *V2_ID, *AUX_ID])) + .with_allowed_serialized_ids(HashSet::from([*V1_ID])); + assert_eq!(compressor.schemes[0].id(), V1.id()); + assert_eq!(compressor.resolve_scheme_id(V3.id()), V1.id()); + + let compressor = compressor.with_allowed_serialized_ids(HashSet::from([*V2_ID, *AUX_ID])); + assert!(compressor.schemes.is_empty()); + assert!(!compressor.has_scheme(V3.id())); +} + +static PUSH_OLD: TestScheme = TestScheme { + push: Some(&V1), + ..TestScheme::new("test.push_old", 0, None) +}; +static PUSH_NEW: TestScheme = TestScheme { + push: Some(&V3), + ..TestScheme::new("test.push_new", 0, None) +}; +static PULL_OLD: TestScheme = TestScheme { + pull: Some(&V1), + ..TestScheme::new("test.pull_old", 0, None) +}; +static PULL_NEW: TestScheme = TestScheme { + pull: Some(&V3), + ..TestScheme::new("test.pull_new", 0, None) +}; + +#[test] +fn exclusions_follow_upgrades_and_fallbacks() { + for allowed in [HashSet::from([*V1_ID]), HashSet::from([*V3_ID])] { + let compressor = + CascadingCompressor::new(vec![&V3, &PUSH_OLD, &PUSH_NEW, &PULL_OLD, &PULL_NEW]) + .with_allowed_serialized_ids(allowed); + let selected = compressor.schemes[0]; + for child in [0, 1] { + for pusher in [&PUSH_OLD, &PUSH_NEW] { + let ctx = compressor + .root_context() + .descend_with_scheme(pusher.id(), child); + assert_eq!(compressor.is_excluded(selected, &ctx), child == 1); + } + let ctx = compressor + .root_context() + .descend_with_scheme(selected.id(), child); + for puller in [&PULL_OLD, &PULL_NEW] { + assert_eq!(compressor.is_excluded(puller, &ctx), child == 1); + } + assert!(compressor.is_excluded(selected, &ctx)); + } + } +} + +#[test] +fn root_exclusions_follow_new_versions() { + static DICT_V2: TestScheme = TestScheme::new("test.dict_v2", 3, Some(&IntDictScheme)); + let compressor = CascadingCompressor::new(vec![&DICT_V2]); + let ctx = compressor + .root_context() + .descend_with_scheme(ROOT_SCHEME_ID, structural::root_list_children::OFFSETS); + assert!(compressor.is_excluded(&DICT_V2, &ctx)); + let ctx = compressor + .root_context() + .descend_with_scheme(ROOT_SCHEME_ID, structural::root_list_children::SIZES); + assert!(!compressor.is_excluded(&DICT_V2, &ctx)); +} + +#[test] +#[should_panic(expected = "appears more than once")] +fn predecessor_cycles_are_rejected() { + static CYCLE: TestScheme = TestScheme::new("test.cycle", 1, Some(&CYCLE)); + CascadingCompressor::new(vec![&CYCLE]); +} + +#[test] +#[should_panic(expected = "appears more than once")] +fn registering_multiple_versions_is_rejected() { + CascadingCompressor::new(vec![&V3, &V1]); +} diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 83685031c33..0b9d8e3d4b8 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,11 +4,8 @@ //! Compression context for recursive compression. use std::fmt; -use std::sync::Arc; -use vortex_array::ArrayId; use vortex_error::VortexExpect; -use vortex_utils::aliases::hash_set::HashSet; use crate::compressor::ROOT_SCHEME_ID; use crate::scheme::SchemeId; @@ -41,24 +38,18 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, - - /// The serialized IDs the writer may emit, or `None` for no restriction. Shared by every - /// context of one compress call, so cloning at each descent is a pointer bump. - allowed_serialized_ids: Option>>, } impl CompressorContext { - /// Creates a new root `CompressorContext` for a compressor that may emit the given serialized - /// IDs, or any ID when `None`. + /// Creates a new root `CompressorContext`. /// /// This should **only** be created by the compressor. - pub(crate) fn new(allowed_serialized_ids: Option>>) -> Self { + pub(crate) fn new() -> Self { Self { is_sample: false, allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), - allowed_serialized_ids, } } } @@ -66,7 +57,7 @@ impl CompressorContext { #[cfg(test)] impl Default for CompressorContext { fn default() -> Self { - Self::new(None) + Self::new() } } @@ -76,18 +67,6 @@ impl CompressorContext { self.is_sample } - /// Returns whether the writer may emit the serialized ID `id`. - /// - /// A scheme whose encoding has several wire formats picks its compression mode from this, - /// the newest permitted one, and the same answer is available while estimating and while - /// compressing. Without a restriction every ID is allowed. The serializer still emits the - /// oldest wire form the resulting array fits, and the serialization context validates it. - pub fn allows_serialized_id(&self, id: ArrayId) -> bool { - self.allowed_serialized_ids - .as_ref() - .is_none_or(|allowed| allowed.contains(&id)) - } - /// Returns the merged stats generation options for this compression site. pub fn merged_stats_options(&self) -> GenerateStatsOptions { self.merged_stats_options diff --git a/vortex-compressor/src/scheme/exclusion.rs b/vortex-compressor/src/scheme/exclusion.rs index 2dba6b85046..46ca12d7735 100644 --- a/vortex-compressor/src/scheme/exclusion.rs +++ b/vortex-compressor/src/scheme/exclusion.rs @@ -34,7 +34,8 @@ impl ChildSelection { /// `ZigZag` excludes `Dict` from all its children. #[derive(Debug, Clone, Copy)] pub struct DescendantExclusion { - /// The scheme to exclude from descendants. + /// The scheme to exclude from descendants. Any version in its registered predecessor chain + /// refers to the selected version. pub excluded: SchemeId, /// Which children of the declaring scheme this rule applies to. pub children: ChildSelection, @@ -47,7 +48,8 @@ pub struct DescendantExclusion { /// `Sequence` excludes itself when `IntDict` is an ancestor on its codes child. #[derive(Debug, Clone, Copy)] pub struct AncestorExclusion { - /// The ancestor scheme that makes the declaring scheme ineligible. + /// The ancestor scheme that makes the declaring scheme ineligible. Any version in its + /// registered predecessor chain refers to the selected version. pub ancestor: SchemeId, /// Which children of the ancestor this rule applies to. pub children: ChildSelection, diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index fa42231d422..f00cde3abd9 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,17 +124,31 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// The serialized IDs this scheme may write its output under. + /// The serialized IDs this scheme may write its output under. Every ID must be permitted before + /// this scheme can be selected. /// /// Cascaded children are compressed by other schemes, which declare their own IDs, so only /// arrays constructed directly by [`compress`](Scheme::compress) belong here. Canonical /// arrays the scheme merely rearranges do not need to be declared. /// - /// An encoding with several wire formats lists every one of them, oldest first. The writer - /// keeps the scheme while any of them is permitted, and the scheme picks the newest - /// permitted one as its compression mode. + /// Alternative versions belong in the [`predecessor`](Scheme::predecessor) chain, rather than + /// in this list. Once selected, a scheme must produce output compatible with these IDs without + /// consulting the writer's configuration. fn produced_encodings(&self) -> Vec; + /// The preceding version of this scheme, used when this version's serialized IDs are unavailable. + /// + /// Register only the newest version. The compressor selects the first eligible version in + /// this chain during configuration, before matching, generating statistics, or estimating. + /// A predecessor is a compatibility fallback, not an alternative compression candidate. + /// + /// Versions must have distinct scheme IDs and form an acyclic chain. They must support the + /// same input types and preserve child indices, because exclusions and scheme dependencies + /// referring to any version in the registered chain apply to the selected version. + fn predecessor(&self) -> Option<&'static dyn Scheme> { + None + } + /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies /// every scheme. diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 13dac7ce7b5..4d3b7bf1a56 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1729,6 +1729,41 @@ async fn test_encoding_registered_after_write_options() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::sparse(PrimitiveArray::from_iter( + (0..4096i32).map(|i| if i % 100 == 0 { i + 1 } else { 0 }), +).into_array())] +#[case::fsst(VarBinViewArray::from_iter( + (0..4096).map(|i| Some(format!("this_is_a_common_prefix_with_some_variation_{i}_and_a_common_suffix_pattern"))), + DType::Utf8(Nullability::NonNullable), +).into_array())] +#[tokio::test] +async fn test_writer_excludes_schemes_with_unavailable_outputs( + #[case] array: ArrayRef, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::() + .with::(); + // Permit Constant and VarBin, but not the subsequently registered Sparse and FSST. + crate::enable_all_registered_array_encodings(&session); + crate::register_default_encodings(&session); + let mut buf = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buf, array.clone().to_array_stream()) + .await?; + let read = session + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + assert_arrays_eq!(read, array, &mut session.create_execution_ctx()); + Ok(()) +} + #[tokio::test] async fn test_writer_empty_chunks() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 79d485f9b1f..725405d3302 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -402,8 +402,6 @@ fn new_array_context( .registry() .read(|registry| registry.keys().copied().collect()) }; - // The compressor sees the same set: it keeps the schemes that can write one of these IDs, and - // an encoding with several wire formats produces the newest one permitted. let allowed_serialized_ids: HashSet = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { From b28f0b65aa41a87d8af0b4eb969797035762810d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 14:40:19 -0400 Subject: [PATCH 15/16] Fix optional scheme trait methods and documentation links Restore produced_encodings in feature-gated schemes and fix stale trait links. Document scheme ID resolution for Clippy. Signed-off-by: "Matt Katz" --- vortex-btrblocks/src/builder/mod.rs | 2 +- vortex-btrblocks/src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- vortex-btrblocks/src/schemes/string/zstd_buffers.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 3 ++- 6 files changed, 7 insertions(+), 6 deletions(-) diff --git a/vortex-btrblocks/src/builder/mod.rs b/vortex-btrblocks/src/builder/mod.rs index ea0ef849dd2..49cd98213af 100644 --- a/vortex-btrblocks/src/builder/mod.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -231,7 +231,7 @@ impl BtrBlocksCompressorBuilder { /// call. /// /// At build time, each scheme is replaced by the newest version in its predecessor chain - /// whose [`required_serialized_ids`](Scheme::required_serialized_ids) are all permitted. + /// whose [`produced_encodings`](Scheme::produced_encodings) are all permitted. /// Schemes with no eligible version are removed. This also applies to schemes added after /// this call. The file writer passes the serialized IDs its enabled editions permit. pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 5204e4e2478..3f06d65b061 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 86b69bba47e..46b2f1e302e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index dddaa349de1..a1bc8643775 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index 98d5feee2d9..cf691c70fcb 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 03733aca5e1..464768e8b36 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -101,7 +101,7 @@ impl CascadingCompressor { /// Selects the newest eligible version of each scheme, intersecting with any earlier call. /// - /// A version is eligible only when all of its [`Scheme::required_serialized_ids`] are allowed. + /// A version is eligible only when all of its [`Scheme::produced_encodings`] are allowed. /// Otherwise its predecessors are tried in order; the scheme is removed if none is eligible. /// Selection preserves registration order and happens before any compression or estimation. pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { @@ -154,6 +154,7 @@ impl CascadingCompressor { self.scheme_aliases.contains_key(&scheme) } + /// Resolves a registered version to the selected version, leaving unknown IDs unchanged. fn resolve_scheme_id(&self, scheme: SchemeId) -> SchemeId { self.scheme_aliases.get(&scheme).copied().unwrap_or(scheme) } From 52a1cb59ac1a6e1a9ebebecf63948a54f1c9dd66 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 16:34:19 -0400 Subject: [PATCH 16/16] Compress wide decimals when the writer permits the v2 format Register DecimalSchemeV2 with the frozen scheme as its predecessor, so format permissions select a scheme before decimal splitting. Cover compression fallback and writer round trips for both format versions. Signed-off-by: "Matt Katz" --- Cargo.lock | 1 + vortex-btrblocks/src/builder/mod.rs | 2 +- .../schemes/{decimal.rs => decimal/mod.rs} | 14 +- vortex-btrblocks/src/schemes/decimal/tests.rs | 272 ++++++++++++++++++ vortex-btrblocks/src/schemes/decimal/v2.rs | 103 +++++++ vortex-file/Cargo.toml | 1 + vortex-file/src/tests.rs | 156 ++++++++++ vortex/Cargo.toml | 1 + 8 files changed, 545 insertions(+), 5 deletions(-) rename vortex-btrblocks/src/schemes/{decimal.rs => decimal/mod.rs} (88%) create mode 100644 vortex-btrblocks/src/schemes/decimal/tests.rs create mode 100644 vortex-btrblocks/src/schemes/decimal/v2.rs diff --git a/Cargo.lock b/Cargo.lock index 639e06bf078..da8b437b3c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11132,6 +11132,7 @@ dependencies = [ "object_store", "parking_lot", "pin-project-lite", + "rand 0.10.2", "rstest", "tokio", "tracing", diff --git a/vortex-btrblocks/src/builder/mod.rs b/vortex-btrblocks/src/builder/mod.rs index 49cd98213af..0656c2006ec 100644 --- a/vortex-btrblocks/src/builder/mod.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -62,7 +62,7 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ //////////////////////////////////////////////////////////////////////////////////////////////// &binary::BinaryDictScheme, // Decimal schemes. - &decimal::DecimalScheme, + &decimal::DecimalSchemeV2, // Temporal schemes. &temporal::TemporalScheme, ]; diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal/mod.rs similarity index 88% rename from vortex-btrblocks/src/schemes/decimal.rs rename to vortex-btrblocks/src/schemes/decimal/mod.rs index 1dff2171f60..f92fdad4115 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal/mod.rs @@ -1,8 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Decimal compression scheme using byte-part decomposition. +//! Versioned decimal compression schemes using byte-part decomposition. +mod v2; +pub use v2::DecimalSchemeV2; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -27,7 +29,10 @@ use crate::SchemeExt; /// Compression scheme for decimal arrays via byte-part decomposition. /// /// Narrows the decimal to the smallest integer type, compresses the underlying primitive, and wraps -/// the result in a `DecimalBytePartsArray`. +/// the result in a `DecimalBytePartsArray` under the frozen single-part wire format. Values that +/// remain wider than 64 bits are left canonical. +/// +/// This is the compatibility predecessor of [`DecimalSchemeV2`]. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DecimalScheme; @@ -66,8 +71,6 @@ impl Scheme for DecimalScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): add support splitting i128/256 buffers into chunks of primitive values - // for compression. 2 for i128 and 4 for i256. let decimal = data.array().clone().execute::(exec_ctx)?; let decimal = narrowed_decimal(decimal); let validity = decimal.validity()?; @@ -85,3 +88,6 @@ impl Scheme for DecimalScheme { DecimalByteParts::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array()) } } + +#[cfg(test)] +mod tests; diff --git a/vortex-btrblocks/src/schemes/decimal/tests.rs b/vortex-btrblocks/src/schemes/decimal/tests.rs new file mode 100644 index 00000000000..5affbff30d8 --- /dev/null +++ b/vortex-btrblocks/src/schemes/decimal/tests.rs @@ -0,0 +1,272 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::iter; +use std::sync::LazyLock; + +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; +use rstest::rstest; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::i256; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; + +use super::DecimalScheme; +use super::DecimalSchemeV2; +use crate::BtrBlocksCompressor; +use crate::BtrBlocksCompressorBuilder; +use crate::SchemeExt; +use crate::SchemeId; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_decimal_byte_parts::initialize(&session); + session +}); + +/// Number of values per array: large enough for cascaded integer schemes to use sampling. +const N: usize = 16_384; + +fn ten_pow(exp: u32) -> i256 { + i256::from_i128(10).wrapping_pow(exp) +} + +/// Deterministic 24-bit noise, so the low part of each value is neither constant nor a +/// sequence — the realistic shape for a wide decimal column with a large fixed magnitude. +fn noise(seed: u64) -> impl Iterator { + let mut rng = StdRng::seed_from_u64(seed); + iter::repeat_with(move || i128::from(rng.random::() >> 8)) +} + +/// `i128`-backed values that need more than 64 bits, so the encoding must carry one lower +/// part. +fn wide_i128_array(validity: Validity) -> DecimalArray { + let base = 10i128.pow(25); + let values: Buffer = noise(7).take(N).map(|delta| base + delta).collect(); + DecimalArray::new(values, DecimalDType::new(38, 2), validity) +} + +/// `i256`-backed values that need more than 128 bits, so the encoding must carry three +/// lower parts. +fn wide_i256_array(validity: Validity) -> DecimalArray { + let base = ten_pow(40); + let values: Buffer = noise(11) + .take(N) + .map(|delta| base + i256::from_i128(delta)) + .collect(); + DecimalArray::new(values, DecimalDType::new(76, 2), validity) +} + +/// Compress with no restriction on serialized IDs, which selects the v2 scheme. +fn compress(array: &ArrayRef) -> VortexResult { + BtrBlocksCompressor::default().compress(array, &mut SESSION.create_execution_ctx()) +} + +/// Compress as a writer whose editions permit the frozen byte-parts format but not v2. +fn compress_v1_only(array: &ArrayRef) -> VortexResult { + let v1_only = HashSet::from([DecimalByteParts.id()]); + BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&v1_only) + .build() + .compress(array, &mut SESSION.create_execution_ctx()) +} + +fn byte_parts(array: &ArrayRef) -> &ArrayRef { + assert!( + array.is::(), + "expected DecimalByteParts, got {}", + array.encoding_id() + ); + array +} + +fn lower_part_count(array: &ArrayRef) -> usize { + byte_parts(array) + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len() +} + +/// When the writer may emit the v2 format, values too wide for a single signed part split into +/// lower parts: one for `i128` storage, three for `i256`. +#[rstest] +#[case::i128(wide_i128_array(Validity::NonNullable).into_array(), 1)] +#[case::i128_nullable(wide_i128_array(Validity::from_iter((0..N).map(|i| i % 3 != 0))).into_array(), 1)] +#[case::i256(wide_i256_array(Validity::NonNullable).into_array(), 3)] +#[case::i256_nullable(wide_i256_array(Validity::from_iter((0..N).map(|i| i % 5 != 0))).into_array(), 3)] +fn test_wide_decimals_split_when_v2_is_permitted( + #[case] array: ArrayRef, + #[case] expected_lower_parts: usize, + #[values(false, true)] explicit_ids: bool, +) -> VortexResult<()> { + let mut builder = BtrBlocksCompressorBuilder::default(); + if explicit_ids { + builder = builder.allow_serialized_ids(&HashSet::from([ + DecimalByteParts.id(), + decimal_byte_parts_v2_id(), + ])); + } + let compressed = builder + .build() + .compress(&array, &mut SESSION.create_execution_ctx())?; + assert_eq!(lower_part_count(&compressed), expected_lower_parts); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + + let serialization = SESSION + .array_serialize(&compressed)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialization.serialized_id, decimal_byte_parts_v2_id()); + Ok(()) +} + +/// A writer that may emit only the frozen format leaves wide values as the canonical decimal: +/// splitting them would need lower parts, for which no single-part form exists. +#[rstest] +#[case::i128(wide_i128_array(Validity::NonNullable).into_array())] +#[case::i128_nullable(wide_i128_array(Validity::from_iter((0..N).map(|i| i % 3 != 0))).into_array())] +#[case::i256(wide_i256_array(Validity::NonNullable).into_array())] +#[case::i256_nullable(wide_i256_array(Validity::from_iter((0..N).map(|i| i % 5 != 0))).into_array())] +fn test_wide_decimals_stay_canonical_without_v2(#[case] array: ArrayRef) -> VortexResult<()> { + let compressed = compress_v1_only(&array)?; + + assert!( + compressed.as_opt::().is_none(), + "expected the wide decimal to be left canonical, got {}", + compressed.encoding_id() + ); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[test] +fn test_i256_decimal_round_trips_extreme_values() -> VortexResult<()> { + // Every 64-bit window exercised, including the sign boundary of the most significant + // part. Bounded by the precision so the values are legal `Decimal(76, 0)` scalars. + let max = ten_pow(76) - i256::ONE; + let values: Buffer = (0..N) + .map(|i| match i % 8 { + 0 => i256::ZERO, + 1 => i256::ONE, + 2 => i256::ZERO - i256::ONE, + 3 => i256::from_parts(u128::MAX, 0), + 4 => i256::from_parts(0, 1), + 5 => i256::from_parts(0, -1), + 6 => max, + _ => i256::ZERO - max, + }) + .collect(); + let array = + DecimalArray::new(values, DecimalDType::new(76, 0), Validity::NonNullable).into_array(); + + let compressed = compress(&array)?; + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[rstest] +fn test_narrow_decimal_has_no_lower_parts( + #[values(false, true)] v1_only: bool, +) -> VortexResult<()> { + // Values that fit 64 bits are narrowed rather than split, even when the declared + // precision needs an i256. + let values: Buffer = (0..N as i128).map(|i| i256::from_i128(i * 3)).collect(); + let array = + DecimalArray::new(values, DecimalDType::new(76, 2), Validity::NonNullable).into_array(); + + let compressed = if v1_only { + compress_v1_only(&array)? + } else { + compress(&array)? + }; + assert_eq!(lower_part_count(&compressed), 0); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + + // Narrow values keep the frozen format even with the v2 scheme. + let serialization = SESSION + .array_serialize(&compressed)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialization.serialized_id, DecimalByteParts.id()); + Ok(()) +} + +#[rstest] +fn test_narrow_precision_with_wide_null_slot( + #[values(false, true)] v1_only: bool, +) -> VortexResult<()> { + let array = DecimalArray::new( + buffer![1i64, i64::MAX, 3], + DecimalDType::new(2, 0), + Validity::from_iter([true, false, true]), + ) + .into_array(); + let compressed = if v1_only { + compress_v1_only(&array)? + } else { + compress(&array)? + }; + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[rstest] +#[case::neither(vec![], false)] +#[case::v1(vec![DecimalByteParts.id()], true)] +#[case::v2_without_v1(vec![decimal_byte_parts_v2_id()], false)] +#[case::both(vec![DecimalByteParts.id(), decimal_byte_parts_v2_id()], true)] +fn test_decimal_scheme_requires_every_possible_wire_id( + #[case] allowed: Vec, + #[case] enabled: bool, +) { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&allowed.into_iter().collect()) + .build(); + assert_eq!(compressor.has_scheme(DecimalScheme.id()), enabled); + assert_eq!(compressor.has_scheme(DecimalSchemeV2.id()), enabled); +} + +#[rstest] +#[case::v1(DecimalScheme.id())] +#[case::v2(DecimalSchemeV2.id())] +fn test_excluding_either_decimal_version_removes_the_chain(#[case] excluded: SchemeId) { + let compressor = BtrBlocksCompressorBuilder::default() + .exclude_schemes([excluded]) + .build(); + assert!(!compressor.has_scheme(DecimalScheme.id())); + assert!(!compressor.has_scheme(DecimalSchemeV2.id())); +} + +#[test] +fn test_canonical_of_compressed_wide_decimal_keeps_storage_width() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + let array = wide_i128_array(Validity::NonNullable).into_array(); + let canonical = compress(&array)?.execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + + let array = wide_i256_array(Validity::NonNullable).into_array(); + let canonical = compress(&array)?.execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + Ok(()) +} diff --git a/vortex-btrblocks/src/schemes/decimal/v2.rs b/vortex-btrblocks/src/schemes/decimal/v2.rs new file mode 100644 index 00000000000..219d2af8a1e --- /dev/null +++ b/vortex-btrblocks/src/schemes/decimal/v2.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal compression with lower parts for wide values. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::decimal::narrowed_decimal; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsSlots; +use vortex_decimal_byte_parts::MAX_LOWER_PARTS; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_decimal_byte_parts::split_decimal; +use vortex_error::VortexResult; + +use super::DecimalScheme; +use crate::ArrayAndStats; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; +use crate::SchemeExt; + +/// Compression scheme for decimals with a signed most significant part and up to three lower parts. +/// +/// Both byte-parts wire IDs must be permitted: wide values serialize under v2, while values that +/// narrow to a single part retain the frozen format. The compressor falls back to [`DecimalScheme`] +/// when only the frozen format is available. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct DecimalSchemeV2; + +impl Scheme for DecimalSchemeV2 { + fn scheme_name(&self) -> &'static str { + "vortex.decimal.byte_parts_v2" + } + + fn matches(&self, canonical: &Canonical) -> bool { + DecimalScheme.matches(canonical) + } + + fn produced_encodings(&self) -> Vec { + vec![DecimalByteParts.id(), decimal_byte_parts_v2_id()] + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + Some(&DecimalScheme) + } + + /// Children: msp=0, then up to [`MAX_LOWER_PARTS`] lower parts. + fn num_children(&self) -> usize { + DecimalBytePartsSlots::FIXED_COUNT + MAX_LOWER_PARTS + } + + fn expected_compression_ratio( + &self, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + DecimalScheme.expected_compression_ratio(data, compress_ctx, exec_ctx) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let decimal = data.array().clone().execute::(exec_ctx)?; + let decimal = narrowed_decimal(decimal); + let parts = split_decimal(&decimal, exec_ctx)?; + + let msp = compressor.compress_child( + &parts.msp, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::MSP, + exec_ctx, + )?; + let lower_parts = parts + .lower_parts + .iter() + .enumerate() + .map(|(idx, part)| { + compressor.compress_child( + part, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx, + exec_ctx, + ) + }) + .collect::>>()?; + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal.decimal_dtype()) + .map(|d| d.into_array()) + } +} diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 626a3fabffe..25126511b7d 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -63,6 +63,7 @@ vortex-zstd = { workspace = true, optional = true } [dev-dependencies] allocator-api2 = { workspace = true } divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 4d3b7bf1a56..ae8e8956515 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -11,6 +11,9 @@ use flatbuffers::FlatBufferBuilder; use futures::StreamExt; use futures::TryStreamExt; use futures::pin_mut; +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -38,6 +41,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; +use vortex_array::dtype::i256; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::and; @@ -72,9 +76,16 @@ use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use vortex_decimal_byte_parts::split_decimal; +use vortex_edition::EDITION_DECLARATIONS; use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_err; use vortex_flatbuffers::footer as fb; use vortex_io::session::RuntimeSession; use vortex_layout::DynLayout; @@ -251,6 +262,73 @@ async fn test_round_trip_many_types() { assert_eq!(read.len(), 3); } +/// End-to-end check that decimals wider than 64 bits survive a write/read round trip. +/// +/// The test session permits both byte-parts wire formats, so wide values can be split and compressed. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_wide_decimal_round_trips_through_a_file() -> VortexResult<()> { + const N: usize = 16_384; + + /// Deterministic 24-bit noise, so the low bits of each value are neither constant nor a + /// sequence. + fn noise(seed: u64) -> impl Iterator { + let mut rng = StdRng::seed_from_u64(seed); + iter::repeat_with(move || i128::from(rng.random::() >> 8)) + } + + // Values that need more than 64 bits, so `i128` storage cannot be narrowed away. + let decimal_38 = DecimalArray::new( + noise(7) + .take(N) + .map(|delta| 10i128.pow(25) + delta) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array(); + + // Values that need more than 128 bits, so `i256` storage cannot be narrowed away. + let base = i256::from_i128(10).wrapping_pow(40); + let decimal_76 = DecimalArray::new( + noise(11) + .take(N) + .map(|delta| base + i256::from_i128(delta)) + .collect::>(), + DecimalDType::new(76, 4), + Validity::from_iter((0..N).map(|i| i % 9 != 0)), + ) + .into_array(); + + let st = StructArray::from_fields(&[ + ("decimal_38", decimal_38), + ("decimal_76_nullable", decimal_76), + ])? + .into_array(); + let dtype = st.dtype().clone(); + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, st.clone().to_array_stream()) + .await?; + + let chunks: Vec<_> = SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .try_collect() + .await?; + let read = ChunkedArray::try_new(chunks, dtype)?.into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(read.len(), N); + assert_arrays_eq!(st, read, &mut ctx); + + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_read_simple_with_spawn() { @@ -2875,3 +2953,81 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { assert_eq!(result.len(), 1); Ok(()) } + +/// The default writer selects the wide or single-part scheme from the enabled editions, +/// including when its input already carries lower parts. +#[rstest] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_default_writer_selects_decimal_scheme_version( + #[values(false, true)] v2_enabled: bool, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::(); + crate::register_default_encodings(&session); + if v2_enabled { + crate::enable_all_registered_array_encodings(&session); + } else { + for declaration in EDITION_DECLARATIONS { + session + .register_edition(declaration) + .map_err(|error| vortex_err!("{error}"))?; + } + session + .enable_edition(CORE_2026_08_3) + .map_err(|error| vortex_err!("{error}"))?; + } + + let decimal = DecimalArray::new( + (0..64i128) + .map(|i| (1i128 << 70) + i) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + + // Building the encoded array is allowed; only getting it into a file is restricted. + let parts = split_decimal(&decimal, &mut session.create_execution_ctx())?; + assert_eq!(parts.lower_parts.len(), 1, "expected a wide split"); + let encoded = DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + )? + .into_array(); + + let st = StructArray::from_fields(&[("wide", encoded)])?.into_array(); + let mut buf = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buf, st.clone().to_array_stream()) + .await?; + + let chunks: Vec<_> = session + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .try_collect() + .await?; + + let lower_part_counts: Vec = chunks + .iter() + .flat_map(|chunk| chunk.depth_first_traversal()) + .filter_map(|node| { + node.as_opt::() + .map(|array| array.lower_parts().len()) + }) + .collect(); + assert_eq!(lower_part_counts.is_empty(), !v2_enabled); + assert!( + lower_part_counts.iter().all(|count| *count == 1), + "expected one lower part per array, got {lower_part_counts:?}" + ); + + let mut ctx = session.create_execution_ctx(); + let read = ChunkedArray::try_new(chunks, st.dtype().clone())?.into_array(); + assert_arrays_eq!(st, read, &mut ctx); + Ok(()) +} diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index e90f6f6e8e7..c97a0e3360e 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -69,6 +69,7 @@ tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } vortex = { path = ".", features = ["tokio"] } +vortex-array = { workspace = true, features = ["_test-harness"] } [features] default = ["files", "wasm-bindgen", "zstd"]