From 3877465a7a7485f9ac4b7ea4240309a833ce2b53 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 16:23:13 -0400 Subject: [PATCH 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 6/7] 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 41fbd809ea2e32fac1ab5886427bf218bd00e2ed Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 15:07:50 -0400 Subject: [PATCH 7/7] address comments DCO Remediation Commit for Matt Katz I, Matt Katz , hereby add my Signed-off-by to this commit: 3877465a7a7485f9ac4b7ea4240309a833ce2b53 I, Matt Katz , hereby add my Signed-off-by to this commit: 966977b9a25252d403df68f6cc621a4b48a6445f I, Matt Katz , hereby add my Signed-off-by to this commit: 643f4b1b9132a89d668ce8b2849fe8d1a6a43ead I, Matt Katz , hereby add my Signed-off-by to this commit: d93444d6224a6f049148e74ed942c157dd920bec I, Matt Katz , hereby add my Signed-off-by to this commit: 1f2143a8a896fe6bd73f40c34c26138b9b718765 Signed-off-by: Matt Katz --- .../decimal-byte-parts/benches/common/mod.rs | 2 +- .../benches/dbp_assemble.rs | 2 +- .../decimal-byte-parts/benches/dbp_split.rs | 16 +-- .../src/decimal_byte_parts/limbs/mod.rs | 106 ++++++++---------- vortex-array/src/dtype/bigint/mod.rs | 1 + 5 files changed, 54 insertions(+), 73 deletions(-) diff --git a/encodings/decimal-byte-parts/benches/common/mod.rs b/encodings/decimal-byte-parts/benches/common/mod.rs index eed00b98c1a..3e0e60730e4 100644 --- a/encodings/decimal-byte-parts/benches/common/mod.rs +++ b/encodings/decimal-byte-parts/benches/common/mod.rs @@ -17,7 +17,7 @@ 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))) + .flat_map(|values_type| [1_024, 8_192].map(|len| (values_type, len))) .collect() } diff --git a/encodings/decimal-byte-parts/benches/dbp_assemble.rs b/encodings/decimal-byte-parts/benches/dbp_assemble.rs index 74a66a133f5..327899fa640 100644 --- a/encodings/decimal-byte-parts/benches/dbp_assemble.rs +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -25,7 +25,7 @@ fn main() { } #[divan::bench(args = cases())] -fn assemble(bencher: Bencher, (values_type, len): (DecimalType, usize)) { +fn dbp_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"); diff --git a/encodings/decimal-byte-parts/benches/dbp_split.rs b/encodings/decimal-byte-parts/benches/dbp_split.rs index 8258c50931b..ba716d1ba5b 100644 --- a/encodings/decimal-byte-parts/benches/dbp_split.rs +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -25,29 +25,17 @@ fn main() { } #[divan::bench(args = cases())] -fn all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { +fn dbp_split_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)) { +fn dbp_split_mixed_null(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(); 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..1e561b149fe 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 @@ -14,6 +14,9 @@ //! This is exactly the two's complement bit pattern of the decimal value cut on 64-bit //! boundaries. +use std::ops::BitOr; +use std::ops::Shl; + use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -22,6 +25,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativeDecimalType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -30,6 +34,7 @@ use vortex_array::match_each_signed_integer_ptype; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -192,24 +197,24 @@ fn split_wide( /// Extract the high signed word and low unsigned word of an `i128`. #[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 i128_to_parts(value: i128) -> (i64, [u64; 1]) { + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "each cast preserves a 64-bit window of the original two's complement bits" + )] ((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, + clippy::cast_sign_loss, + reason = "each cast preserves a 64-bit window of the original two's complement bits" + )] ( (high >> LOWER_PART_BITS) as i64, [high as u64, (low >> LOWER_PART_BITS) as u64, low as u64], @@ -263,12 +268,18 @@ pub fn assemble_decimal( .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] => DecimalArray::new( + assemble_wide::(msp, [first]), + decimal_dtype, + validity, + ), + [first, second] => DecimalArray::new( + assemble_wide::(msp, [first, second]), + decimal_dtype, + validity, + ), [first, second, third] => DecimalArray::new( - assemble_i256(msp, [first, second, third]), + assemble_wide::(msp, [first, second, third]), decimal_dtype, validity, ), @@ -279,51 +290,32 @@ pub fn assemble_decimal( }) } -/// Reassemble a signed MSP and one `u64` lower part into `i128` values. +/// Reassemble a signed MSP and `K` unsigned lower parts into wide integers. /// -/// 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 { - let mut out = BufferMut::::with_capacity(msp.len()); - match_each_signed_integer_ptype!(msp.ptype(), |P| { - 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. - (i128::from(i64::from(*value)) << LOWER_PART_BITS) | i128::from(*part) - })); - }); - out.freeze() -} - -/// Reassemble a signed MSP and two or three `u64` lower parts into `i256` values. +/// Each row starts with the MSP sign-extended to `T`. Appending a lower word shifts the +/// accumulated value left by 64 bits and fills the low bits with that word. Lower parts +/// are appended most significant first. /// -/// 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" -)] -fn assemble_i256(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer { - let mut out = BufferMut::::with_capacity(msp.len()); +/// The callers select `i128` for one lower part and `i256` for two or three. Since `K` +/// is constant, the compiler can unroll the loop that appends the lower words. +fn assemble_wide(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer +where + T: NativeDecimalType + Shl + BitOr, +{ + 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 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.extend_trusted(msp.as_slice::

().iter().enumerate().map(|(row, value)| { + #[allow( + clippy::useless_conversion, + reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" + )] + let mut value = T::from(i64::from(*value)).vortex_expect("MSP fits in the output type"); + for part in lower { + value = (value << LOWER_PART_BITS) + | T::from(part[row]).vortex_expect("lower word fits in the output type"); + } + value + })); }); out.freeze() } diff --git a/vortex-array/src/dtype/bigint/mod.rs b/vortex-array/src/dtype/bigint/mod.rs index 3ebf01425d6..47195526b1f 100644 --- a/vortex-array/src/dtype/bigint/mod.rs +++ b/vortex-array/src/dtype/bigint/mod.rs @@ -259,6 +259,7 @@ impl Shr for i256 { impl Shl for i256 { type Output = Self; + #[inline] fn shl(self, rhs: usize) -> Self::Output { use num_traits::ToPrimitive; Self(