From 469a1e8a43e36b00198422a36fce3e9165d9b1d5 Mon Sep 17 00:00:00 2001 From: Matthew Katz <87445739+mhk197@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:15:20 -0400 Subject: [PATCH 1/5] Add wide decimal byte-part splitting and assembly logic (#9808) Signed-off-by: "Matt Katz" 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 | 47 +++ .../src/decimal_byte_parts/limbs/mod.rs | 324 ++++++++++++++++++ .../src/decimal_byte_parts/limbs/tests.rs | 226 ++++++++++++ .../src/decimal_byte_parts/mod.rs | 9 + vortex-array/src/dtype/bigint/mod.rs | 1 + 9 files changed, 718 insertions(+) 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 create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 275db7ef8ee..1343e6dda00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10983,8 +10983,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..3e0e60730e4 --- /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].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..327899fa640 --- /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 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"); + 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..ba716d1ba5b --- /dev/null +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -0,0 +1,47 @@ +// 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 dbp_split_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Validity::AllValid); +} + +#[divan::bench(args = cases())] +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); +} + +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 new file mode 100644 index 00000000000..1e561b149fe --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -0,0 +1,324 @@ +// 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, +} + +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(), + } + } + + 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(), + } + } +} + +/// 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 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 +/// +/// 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::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_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_wide(&decimal.buffer::(), &mask, i256_to_parts); + DecimalParts::new(msp, lower, validity) + } + }) +} + +/// Split wide integers into a signed MSP and `N` unsigned lower parts. +/// +/// `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.map(BufferMut::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]); + + 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); + } + } + } + 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"), + } + + // 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)) +} + +/// Extract the high signed word and low unsigned word of an `i128`. +#[inline] +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] +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], + ) +} + +/// 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 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_wide::(msp, [first]), + decimal_dtype, + validity, + ), + [first, second] => DecimalArray::new( + assemble_wide::(msp, [first, second]), + decimal_dtype, + validity, + ), + [first, second, third] => DecimalArray::new( + assemble_wide::(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 `K` unsigned lower parts into wide integers. +/// +/// 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 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| { + 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() +} + +#[cfg(test)] +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..3e3de06c44e --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs @@ -0,0 +1,226 @@ +// 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)))] +#[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, 63, 64, 65, 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(()) +} + +#[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(()) +} 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 e1d081ddaca..951a893a785 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -9,9 +9,18 @@ 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; +#[doc(hidden)] +pub mod _benchmarking { + pub use super::limbs::assemble_decimal; +} + use prost::Message as _; use vortex_array::ArrayEq; use vortex_array::ArrayHash; 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( From c04e404b56fc4f097614794adf6752b6963414cc Mon Sep 17 00:00:00 2001 From: Matthew Katz <87445739+mhk197@users.noreply.github.com> Date: Tue, 15 Sep 2026 09:05:12 -0400 Subject: [PATCH 2/5] Support wide decimals in `DecimalBytePartsArray` and kernels (#9809) `DecimalBytePartsArray` previously stored the entire unscaled decimal value in one signed integer child, limiting it to values that fit in 64 bits. It now supports wide decimals by representing each value as integer parts that can be compressed independently, while preserving the decimal's logical precision, scale, and nullability. The array has a signed most significant part (MSP) and up to three unsigned 64-bit lower parts, ordered most significant first. Splitting canonical decimal storage produces: | Decimal storage | Children | | --- | --- | | `i8` / `i16` / `i32` / `i64` | Signed MSP only; shares the original value buffer | | `i128` | `i64` MSP + one `u64` lower part | | `i256` | `i64` MSP + three `u64` lower parts | Only the MSP carries validity. Every lower part must be a non-nullable `u64` array with the same length, and splitting wide decimals zeroes the parts at null positions. All children remain `ArrayRef`s, so their individual encodings are independent of the decimal representation. `execute::` reassembles a `DecimalArray` from the MSP and lower parts children. `take` with nullable indices is not yet supported by the DecimalByteParts kernel for arrays with lower parts; it falls back to canonical execution. Taking each part directly would make the lower parts nullable, violating the representation's invariant. The frozen serializer also continues to reject arrays with lower parts. This PR also refines the splitting and assembly modules. * Assembly takes `ArrayRef`s instead of `PrimitiveArray`s so that in the future, we can add special fast paths for constant arrays * Assembly casts lower parts to `u64`s, allowing for assembly of narrowed lower parts. * Assembly loop is optimized such that it vectorizes for `i256` assembly on local runs. --------- Signed-off-by: Matt Katz --- .gitignore | 2 + Cargo.lock | 85 +- Cargo.toml | 1 + encodings/decimal-byte-parts/Cargo.toml | 2 + .../decimal-byte-parts/benches/common/mod.rs | 80 +- .../benches/dbp_assemble.rs | 181 +++- .../decimal-byte-parts/benches/dbp_split.rs | 105 +- .../src/decimal_byte_parts/array.rs | 922 ++++++++++++++++++ .../src/decimal_byte_parts/assemble.rs | 464 +++++++++ .../src/decimal_byte_parts/compute/cast.rs | 13 +- .../src/decimal_byte_parts/compute/compare.rs | 50 + .../src/decimal_byte_parts/compute/filter.rs | 46 +- .../decimal_byte_parts/compute/is_constant.rs | 27 +- .../src/decimal_byte_parts/compute/kernel.rs | 8 - .../src/decimal_byte_parts/compute/mask.rs | 15 +- .../src/decimal_byte_parts/compute/mod.rs | 34 + .../decimal_byte_parts/{ => compute}/slice.rs | 16 +- .../src/decimal_byte_parts/compute/take.rs | 109 ++- .../src/decimal_byte_parts/limbs/mod.rs | 324 ------ .../src/decimal_byte_parts/limbs/tests.rs | 226 ----- .../src/decimal_byte_parts/mod.rs | 414 +------- .../src/decimal_byte_parts/prop_tests.rs | 183 ++++ .../src/decimal_byte_parts/rules.rs | 42 +- .../src/decimal_byte_parts/split.rs | 310 ++++++ .../src/decimal_byte_parts/testing.rs | 47 + vortex-array/src/dtype/bigint/mod.rs | 14 + vortex-btrblocks/src/trace_tests.rs | 5 +- .../kernel/encodings/decimal_byte_parts.rs | 7 + 28 files changed, 2626 insertions(+), 1106 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs rename encodings/decimal-byte-parts/src/decimal_byte_parts/{ => compute}/slice.rs (53%) delete mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs delete mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/testing.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 1343e6dda00..ca9cd626f1b 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", ] @@ -2229,6 +2229,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" @@ -4071,7 +4090,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.9.1", "zlib-rs", ] @@ -4678,6 +4697,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" @@ -6592,6 +6656,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" @@ -6898,6 +6971,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" @@ -10984,11 +11063,13 @@ name = "vortex-decimal-byte-parts" version = "0.1.0" dependencies = [ "codspeed-divan-compat", + "hegeltest", "num-traits", "prost 0.14.4", "rand 0.10.2", "rstest", "vortex-array", + "vortex-bench-support", "vortex-buffer", "vortex-error", "vortex-mask", diff --git a/Cargo.toml b/Cargo.toml index 6a94cd1eb56..4d73870beb4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -176,6 +176,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..e5d15c2c4f7 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -27,9 +27,11 @@ 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"] } +vortex-bench-support = { workspace = true } [[bench]] name = "dbp_assemble" diff --git a/encodings/decimal-byte-parts/benches/common/mod.rs b/encodings/decimal-byte-parts/benches/common/mod.rs index 3e0e60730e4..252431a58a3 100644 --- a/encodings/decimal-byte-parts/benches/common/mod.rs +++ b/encodings/decimal-byte-parts/benches/common/mod.rs @@ -1,51 +1,77 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Shared decimal inputs for splitting and assembly benchmarks. +//! Shared 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] + [DecimalType::I128, DecimalType::I256] .into_iter() .flat_map(|values_type| [1_024, 8_192].map(|len| (values_type, len))) .collect() } -pub(super) fn decimal_array( - values_type: DecimalType, - len: usize, - validity: Validity, -) -> DecimalArray { +pub(super) fn i128_values(len: usize) -> Buffer { let mut rng = StdRng::seed_from_u64(42); + let max = 10i128.pow(38) - 1; + (0..len).map(|_| rng.random_range(-max..=max)).collect() +} + +pub(super) fn i256_values(len: usize) -> Buffer { + let mut rng = StdRng::seed_from_u64(42); + // Keep the magnitude below 10^76 while exercising all four signed/unsigned words. + (0..len) + .map(|_| i256::from_parts(rng.random(), rng.random::() >> 4)) + .collect() +} + +#[cfg(not(codspeed))] +pub(super) mod arrays { + 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::validity::Validity; + use vortex_buffer::Buffer; + use vortex_error::vortex_panic; + + use super::i128_values; + use super::i256_values; - 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) - }}; + pub(crate) fn cases() -> Vec<(DecimalType, usize)> { + [DecimalType::I64, DecimalType::I128, DecimalType::I256] + .into_iter() + .flat_map(|values_type| [1_024, 8_192].map(|len| (values_type, len))) + .collect() } - 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) + pub(crate) fn decimal_array( + values_type: DecimalType, + len: usize, + validity: Validity, + ) -> DecimalArray { + match values_type { + DecimalType::I64 => { + let mut rng = StdRng::seed_from_u64(42); + let max = 10i64.pow(18) - 1; + let values: Buffer = (0..len).map(|_| rng.random_range(-max..=max)).collect(); + DecimalArray::new(values, DecimalDType::new(18, 2), validity) + } + DecimalType::I128 => { + DecimalArray::new(i128_values(len), DecimalDType::new(38, 2), validity) + } + DecimalType::I256 => { + DecimalArray::new(i256_values(len), DecimalDType::new(76, 2), validity) + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), } - _ => 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 index 327899fa640..48161a2c664 100644 --- a/encodings/decimal-byte-parts/benches/dbp_assemble.rs +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -1,48 +1,169 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Reassembling primitive decimal parts across storage widths and lengths. +//! Reassembling integer slices, including output-buffer allocation. +//! Array benchmarks also include validation, part execution, and casts. 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 vortex_array::dtype::i256; +use vortex_buffer::buffer; +use vortex_decimal_byte_parts::_benchmarking::assemble_wide_decimal; +use vortex_decimal_byte_parts::_benchmarking::i128_to_parts; +use vortex_decimal_byte_parts::_benchmarking::i256_to_parts; +use vortex_decimal_byte_parts::_benchmarking::split_wide; +use vortex_error::vortex_panic; +use vortex_mask::Mask; use crate::common::cases; -use crate::common::decimal_array; +use crate::common::i128_values; +use crate::common::i256_values; fn main() { divan::main(); } +#[vortex_bench_support::cpu_features] #[divan::bench(args = cases())] -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"); - 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") - }); +fn dbp_assemble_kernel(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let validity = Mask::new_true(len); + match values_type { + DecimalType::I128 => { + let values = i128_values(len); + let (msp, [lower]) = split_wide(values.as_slice(), &validity, i128_to_parts); + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(lower.as_slice()).iter().map(|&word| [word]), + ) + }); + } + DecimalType::I256 => { + let values = i256_values(len); + let (msp, [first, second, third]) = + split_wide(values.as_slice(), &validity, i256_to_parts); + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(first.as_slice()) + .iter() + .zip(black_box(second.as_slice())) + .zip(black_box(third.as_slice())) + .map(|((&a, &b), &c)| [a, b, c]), + ) + }); + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } +} + +// The kernel widens the MSP while assembling; lower parts are already `u64`. +#[vortex_bench_support::cpu_features] +#[divan::bench(args = cases())] +fn dbp_assemble_kernel_narrow_msp(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let msp = buffer![-1i8; len]; + match values_type { + DecimalType::I128 => { + let lower = buffer![u64::from(u8::MAX); len]; + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(lower.as_slice()).iter().map(|&word| [word]), + ) + }); + } + DecimalType::I256 => { + let first = buffer![u64::from(u8::MAX); len]; + let second = buffer![u64::from(u16::MAX); len]; + let third = buffer![u64::from(u32::MAX); len]; + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(first.as_slice()) + .iter() + .zip(black_box(second.as_slice())) + .zip(black_box(third.as_slice())) + .map(|((&a, &b), &c)| [a, b, c]), + ) + }); + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } +} + +/// Benchmarks of full end-to-end dbp assembly, including kernel. These +/// are more volatile and should not be run in CI due to noise. +#[cfg(not(codspeed))] +mod arrays { + use divan::Bencher; + use divan::black_box; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::DecimalType; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_decimal_byte_parts::_benchmarking::assemble_decimal; + use vortex_decimal_byte_parts::split_decimal; + use vortex_error::VortexExpect; + use vortex_error::vortex_panic; + + use crate::common::arrays::cases; + use crate::common::arrays::decimal_array; + + #[divan::bench(args = cases())] + 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"); + bench_parts( + bencher, + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ); + } + + #[divan::bench(args = cases())] + fn dbp_assemble_narrowed(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let (precision, lower_parts) = match values_type { + DecimalType::I64 => (18, vec![]), + DecimalType::I128 => (38, vec![buffer![u8::MAX; len].into_array()]), + DecimalType::I256 => ( + 76, + vec![ + buffer![u8::MAX; len].into_array(), + buffer![u16::MAX; len].into_array(), + buffer![u32::MAX; len].into_array(), + ], + ), + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + }; + bench_parts( + bencher, + buffer![-1i8; len].into_array(), + lower_parts, + DecimalDType::new(precision, 2), + ); + } + + fn bench_parts( + bencher: Bencher, + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) { + let session = array_session(); + bencher + .with_inputs(|| session.create_execution_ctx()) + .bench_refs(|ctx| { + assemble_decimal(black_box(&msp), black_box(&lower_parts), decimal_dtype, ctx) + .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 index ba716d1ba5b..dd5b9abcf13 100644 --- a/encodings/decimal-byte-parts/benches/dbp_split.rs +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Splitting decimal arrays across storage widths, lengths, and validity paths. +//! Splitting integer slices, including output-buffer allocation. +//! Array benchmarks also include validity execution and array construction. mod common; @@ -10,38 +11,104 @@ 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 vortex_decimal_byte_parts::_benchmarking::i128_to_parts; +use vortex_decimal_byte_parts::_benchmarking::i256_to_parts; +use vortex_decimal_byte_parts::_benchmarking::split_wide; +use vortex_error::vortex_panic; +use vortex_mask::Mask; use crate::common::cases; -use crate::common::decimal_array; +use crate::common::i128_values; +use crate::common::i256_values; fn main() { divan::main(); } +#[vortex_bench_support::cpu_features] #[divan::bench(args = cases())] -fn dbp_split_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { - bench_split(bencher, values_type, len, Validity::AllValid); +fn dbp_split_kernel_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Mask::new_true(len)); } +#[vortex_bench_support::cpu_features] #[divan::bench(args = cases())] -fn dbp_split_mixed_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { +fn dbp_split_kernel_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))); + let validity = Mask::from_iter((0..len).map(|_| rng.random_bool(0.5))); 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") - }); +fn bench_split(bencher: Bencher, values_type: DecimalType, len: usize, validity: Mask) { + match values_type { + DecimalType::I128 => { + let values = i128_values(len); + bencher.bench(|| { + split_wide( + black_box(values.as_slice()), + black_box(&validity), + i128_to_parts, + ) + }); + } + DecimalType::I256 => { + let values = i256_values(len); + bencher.bench(|| { + split_wide( + black_box(values.as_slice()), + black_box(&validity), + i256_to_parts, + ) + }); + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } +} + +/// Benchmarks of full end-to-end dbp assembly, including kernel. These +/// are more volatile and should not be run in CI due to noise. +#[cfg(not(codspeed))] +mod arrays { + 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::arrays::cases; + use crate::common::arrays::decimal_array; + + #[divan::bench(args = cases())] + 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 dbp_split_all_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Validity::AllInvalid); + } + + #[divan::bench(args = cases())] + 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); + } + + 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/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs new file mode 100644 index 00000000000..bcdb3579f53 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -0,0 +1,922 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal byte-parts array types, validation, and VTable implementations. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::PType; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::serde::ArrayChildren; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityChild; +use vortex_array::vtable::ValidityVTableFromChild; +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 super::LOWER_PART_DTYPE; +use super::MAX_LOWER_PARTS; +use super::assemble::assemble_decimal; +use super::assemble::assemble_wide_decimal_value; +use super::rules::PARENT_RULES; + +/// A [`DecimalByteParts`]-encoded Vortex array. +pub type DecimalBytePartsArray = Array; + +#[derive(Clone, prost::Message)] +pub struct DecimalBytesPartsMetadata { + #[prost(enumeration = "PType", tag = "1")] + zeroth_child_ptype: i32, + #[prost(uint32, tag = "2")] + 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) + } +} + +/// 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 unsigned integer holding a 64-bit window of the value. +/// Parts may have narrower integer dtypes when their values fit; their positions stay 64 bits apart. +/// +/// e.g. for a decimal i128 \[ 127..64 | 63..0 \] msp = 127..64 and lower_part\[0\] = 63..0 +/// +/// All parts live in slots, so the array carries no additional data. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsData; + +impl Display for DecimalBytePartsData { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl ArrayHash for DecimalBytePartsData { + fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} +} + +impl ArrayEq for DecimalBytePartsData { + fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { + true + } +} + +impl DecimalBytePartsData { + /// 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 unsigned integer array of length `len`. + pub fn validate<'a>( + msp: &ArrayRef, + lower_parts: impl ExactSizeIterator, + decimal_dtype: DecimalDType, + dtype: &DType, + len: usize, + ) -> VortexResult<()> { + if !msp.dtype().is_signed_int() { + vortex_bail!("msp must be a signed integer array") + } + + let expected_dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + vortex_ensure!( + dtype == &expected_dtype, + "expected dtype {expected_dtype}, got {dtype}" + ); + vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); + + let lower_part_count = lower_parts.len(); + + vortex_ensure!( + lower_part_count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" + ); + for (idx, part) in lower_parts.enumerate() { + vortex_ensure!( + part.dtype().is_unsigned_int() && !part.dtype().is_nullable(), + "lower part {idx} must have a non-nullable unsigned integer dtype, got {}", + part.dtype() + ); + vortex_ensure!( + part.len() == len, + "lower part {idx} has len {}, expected {len}", + part.len() + ); + } + Ok(()) + } +} + +#[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 unsigned integer + /// array of the same length as the MSP. See [`super::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 { + 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), + ) + } + + /// Construct a [`DecimalBytePartsArray`] from parts whose invariants are already established. + /// + /// # Safety + /// + /// The MSP must have a signed integer dtype (`i8`, `i16`, `i32`, or `i64`). There must be + /// at most [`MAX_LOWER_PARTS`] lower parts, each a non-nullable unsigned integer array with the same + /// length as the MSP. Lower parts are ordered most significant first. + pub(super) unsafe fn new_unchecked( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> DecimalBytePartsArray { + let len = msp.len(); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + // SAFETY: the caller guarantees the part types, lengths, and count. The slot builder + // fills every required slot, and the length and nullability come from the MSP. + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData) + .with_slots(slots), + ) + } + } +} + +impl VTable for DecimalByteParts { + type TypedArrayData = DecimalBytePartsData; + + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let Some(decimal_dtype) = dtype.as_decimal_opt() else { + vortex_bail!("expected decimal dtype, got {}", dtype) + }; + + let min_slots = DecimalBytePartsSlots::FIXED_COUNT; + let max_slots = min_slots + MAX_LOWER_PARTS; + vortex_ensure!( + (min_slots..=max_slots).contains(&slots.len()), + "expected {min_slots}..={max_slots} slots, got {}", + slots.len() + ); + for (idx, slot) in slots.iter().enumerate() { + vortex_ensure!(slot.is_some(), "missing required slot {idx}"); + } + + let slots = DecimalBytePartsSlotsView::from_slots(slots); + DecimalBytePartsData::validate( + slots.msp, + slots.lower_parts.iter(), + *decimal_dtype, + dtype, + len, + ) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("DecimalBytePartsArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + vortex_ensure!( + array.lower_parts().is_empty(), + "serializing DecimalByteParts with lower parts is not supported" + ); + Ok(Some( + DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + _buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = DecimalBytesPartsMetadata::decode(metadata)?; + vortex_ensure!( + metadata.lower_part_count()? == 0, + "vortex.decimal_byte_parts must not carry lower parts" + ); + metadata.into_array_parts(dtype, len, children) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + DecimalBytePartsSlots::slot_name(idx) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let lower_parts = array.lower_parts().to_vec(); + let assembled = assemble_decimal(array.msp(), &lower_parts, array.decimal_dtype(), ctx)?; + + Ok(ExecutionResult::done(assembled)) + } +} + +#[array_slots(DecimalByteParts)] +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, +} + +pub(crate) trait DecimalBytePartsArrayExt: DecimalBytePartsArraySlotsExt { + /// The decimal dtype of this array. + fn decimal_dtype(&self) -> DecimalDType { + *self + .as_ref() + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype") + } + + /// Rebuild the array by applying `f` to the MSP and every lower part, in slot order. + /// + /// 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 { + 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()) + } + + /// Rebuild the array with a replacement MSP, preserving its lower parts, precision and scale. + /// + /// 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(), + ) + } +} + +impl> DecimalBytePartsArrayExt for T {} + +impl OperationsVTable for DecimalByteParts { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, DecimalByteParts>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let scalar = array.msp().execute_scalar(index, ctx)?; + + // 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"); + + // Zero-extend each narrowed lower value to its 64-bit window. + 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 = match lower_parts.as_slice() { + [] => DecimalValue::I64(msp), + [first] => DecimalValue::I128(assemble_wide_decimal_value(msp, [*first])), + [first, second] => { + DecimalValue::I256(assemble_wide_decimal_value(msp, [*first, *second])) + } + [first, second, third] => { + DecimalValue::I256(assemble_wide_decimal_value(msp, [*first, *second, *third])) + } + _ => vortex_bail!( + "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", + lower_parts.len() + ), + }; + + Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value))) + } +} + +impl ValidityChild for DecimalByteParts { + fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { + // validity stored in 0th child + array.msp().clone() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::Array; + use vortex_array::ArrayParts; + 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::builtins::ArrayBuiltins; + 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; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; + use vortex_array::vtable::VTable; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::DecimalByteParts; + use super::DecimalBytePartsArray; + use super::DecimalBytePartsArraySlotsExt; + use super::DecimalBytePartsData; + use crate::decimal_byte_parts::LOWER_PART_DTYPE; + use crate::decimal_byte_parts::MAX_LOWER_PARTS; + 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() { + let decimal_dtype = DecimalDType::new(8, 2); + let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); + let array = DecimalByteParts::try_new( + PrimitiveArray::new( + buffer![100i32, 200i32, 400i32], + Validity::Array(BoolArray::from_iter(vec![false, true, true]).into_array()), + ) + .into_array(), + decimal_dtype, + ) + .unwrap() + .into_array(); + + assert_eq!( + Scalar::null(dtype.clone()), + array + .execute_scalar(0, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Decimal(DecimalValue::I64(200))) + ) + .unwrap(), + array + .execute_scalar(1, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), + array + .execute_scalar(2, &mut array_session().create_execution_ctx()) + .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(()) + } + + #[rstest] + #[case([PType::U8, PType::U16, PType::U32])] + #[case([PType::U16, PType::U32, PType::U64])] + #[case([PType::U32, PType::U64, PType::U8])] + #[case([PType::U64, PType::U8, PType::U16])] + fn test_independently_narrowed_parts( + #[case] lower_ptypes: [PType; 3], + #[values(PType::I8, PType::I16, PType::I32, PType::I64)] msp_ptype: PType, + #[values(1, 2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let validity = Validity::from_iter([false, true, true]); + let msp = PrimitiveArray::new(buffer![0i64, 1, -1], validity.clone()) + .into_array() + .cast(DType::Primitive(msp_ptype, Nullability::Nullable))?; + // Set the unsigned type's highest bit to catch accidental sign extension. + let words = lower_ptypes.map(|ptype| 1u64 << (ptype.byte_width() * 8 - 1)); + let lower = lower_ptypes + .into_iter() + .zip(words) + .take(lower_count) + .map(|(ptype, word)| { + buffer![0u64, word, word] + .into_array() + .cast(DType::Primitive(ptype, Nullability::NonNullable)) + }) + .collect::>>()?; + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + let encoded = DecimalByteParts::try_new_with_lower_parts(msp, lower, dtype)?; + let expected = match lower_count { + 1 => DecimalArray::new( + buffer![ + 0i128, + (1i128 << 64) | i128::from(words[0]), + (-1i128 << 64) | i128::from(words[0]) + ], + dtype, + validity, + ), + 2 => { + let low = (u128::from(words[0]) << 64) | u128::from(words[1]); + DecimalArray::new( + buffer![ + i256::ZERO, + i256::from_parts(low, 1), + i256::from_parts(low, -1) + ], + dtype, + validity, + ) + } + _ => { + let low = (u128::from(words[1]) << 64) | u128::from(words[2]); + DecimalArray::new( + buffer![ + i256::ZERO, + i256::from_parts(low, (1i128 << 64) | i128::from(words[0])), + i256::from_parts(low, (-1i128 << 64) | i128::from(words[0])), + ], + dtype, + validity, + ) + } + }; + let actual = encoded + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(expected.clone(), actual, &mut ctx); + for row in 0..expected.len() { + assert_eq!( + encoded.execute_scalar(row, &mut ctx)?, + expected.execute_scalar(row, &mut ctx)? + ); + } + 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() + ); + } + + #[rstest] + #[case::no_slots(vec![])] + #[case::missing_msp(vec![None])] + #[case::missing_lower(vec![Some(msp()), None])] + #[case::gap_in_lower(vec![Some(msp()), None, Some(lower_part())])] + fn test_rejects_missing_slots(#[case] slots: Vec>) { + let parts = ArrayParts::new( + DecimalByteParts, + DType::Decimal(DecimalDType::new(76, 2), Nullability::NonNullable), + 3, + DecimalBytePartsData, + ) + .with_slots(slots.into_iter().collect()); + assert!(Array::try_from_parts(parts).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/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs new file mode 100644 index 00000000000..35d93d4da1e --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs @@ -0,0 +1,464 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reassembling decimal arrays and values from their parts. + +use std::ops::BitOr; +use std::ops::Shl; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::i256; +use vortex_array::match_each_signed_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::trusted_len::TrustedLen; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use super::LOWER_PART_BITS; +use super::LOWER_PART_DTYPE; +use super::MAX_LOWER_PARTS; + +/// Reassemble decimal parts into a decimal array. +/// +/// The MSP must have a signed integer dtype, and every lower part must have a non-nullable +/// unsigned integer dtype. All parts 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`. +/// +/// If there are lower parts, each part (including the MSP) is widened into a 64 bit array. +/// For example, parts consisting of a `i8` MSP and a single `u32` lower part is assembled into +/// a decimal array with `64 + 64 = 128` bit storage. +/// +/// # Errors +/// +/// Returns an error for invalid part dtypes, lengths, or counts, or if executing a part fails. +pub fn assemble_decimal( + msp: &ArrayRef, + lower_parts: &[ArrayRef], + decimal_dtype: DecimalDType, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + vortex_ensure!( + msp.dtype().is_signed_int(), + "MSP must have a signed integer dtype" + ); + + let validity = msp.validity()?; + + if lower_parts.is_empty() { + return assemble_narrow_decimal(msp, validity, decimal_dtype, exec_ctx); + } + + vortex_ensure!( + lower_parts.len() <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", + lower_parts.len() + ); + let len = msp.len(); + for (idx, part) in lower_parts.iter().enumerate() { + vortex_ensure!( + part.dtype().is_unsigned_int() && !part.dtype().is_nullable(), + "lower part {idx} must have a non-nullable unsigned integer dtype, got {}", + part.dtype() + ); + vortex_ensure!( + part.len() == len, + "lower part {idx} has len {}, expected {len}", + part.len() + ); + } + + assemble_wide_decimal_from_arrays(msp, lower_parts, validity, decimal_dtype, exec_ctx) +} + +fn assemble_narrow_decimal( + msp: &ArrayRef, + validity: Validity, + decimal_dtype: DecimalDType, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + // TODO(mk): Broadcast a constant MSP directly instead of materializing its buffer. + let msp = msp.clone().execute::(exec_ctx)?; + Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { + DecimalArray::new(msp.to_buffer::

(), decimal_dtype, validity).into_array() + })) +} + +/// Execute the MSP at its signed integer width and cast lower parts to `u64` before assembly. +/// The number of lower parts determines the decimal storage type: one produces `i128`, while +/// two or three produce `i256`. +fn assemble_wide_decimal_from_arrays( + msp: &ArrayRef, + lower_parts: &[ArrayRef], + validity: Validity, + decimal_dtype: DecimalDType, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + // TODO(mk): Broadcast constant parts directly instead of materializing their buffers. + let msp = msp.clone().execute::(exec_ctx)?; + // TODO(mk): Revisit dispatching on lower-part dtypes and widening values during assembly. + // Casting narrowed parts allocates temporary buffers and adds passes over the data. + // Nested dtype dispatch is significantly in benchmarks, but adds code and generic instantiations. + let lower = lower_parts + .iter() + .map(|part| { + part.cast(LOWER_PART_DTYPE)? + .execute::(exec_ctx) + }) + .collect::>>()?; + + Ok(match_each_signed_integer_ptype!(msp.ptype(), |Msp| { + let msp = msp.as_slice::(); + match lower.as_slice() { + [first] => DecimalArray::new( + assemble_wide_decimal::( + msp, + first.as_slice::().iter().map(|&word| [word]), + ), + decimal_dtype, + validity, + ) + .into_array(), + [first, second] => DecimalArray::new( + assemble_wide_decimal::( + msp, + first + .as_slice::() + .iter() + .zip(second.as_slice::()) + .map(|(&a, &b)| [a, b]), + ), + decimal_dtype, + validity, + ) + .into_array(), + [first, second, third] => DecimalArray::new( + assemble_wide_decimal::( + msp, + first + .as_slice::() + .iter() + .zip(second.as_slice::()) + .zip(third.as_slice::()) + .map(|((&a, &b), &c)| [a, b, c]), + ), + decimal_dtype, + validity, + ) + .into_array(), + _ => vortex_bail!("expected between one and {MAX_LOWER_PARTS} lower parts"), + } + })) +} + +/// Assemble a signed MSP slice and one array of `K` lower words per row. +/// +/// The caller zips the `u64` lower-part slices into rows. The iterator must have the same +/// length as the MSP slice. MSP values are widened to `i64` as they are read. +pub fn assemble_wide_decimal( + msp: &[Msp], + lower: impl TrustedLen, +) -> Buffer +where + T: NativeDecimalType + From + From + Shl + BitOr, + Msp: Copy + Into, +{ + let mut out = BufferMut::::with_capacity(msp.len()); + out.extend_trusted( + msp.iter() + .zip(lower) + .map(|(&value, parts)| assemble_wide_decimal_value(value.into(), parts)), + ); + out.freeze() +} + +/// Reassemble a decimal's unscaled integer from its signed MSP and `K` lower words. +/// +/// Sign-extend the MSP to `T`, then append each lower word by shifting left 64 bits and +/// filling the low bits. Lower words are ordered most significant first. Callers select +/// `i128` for one lower word and `i256` for two or three. +#[inline] +pub(crate) fn assemble_wide_decimal_value(msp: i64, lower: [u64; K]) -> T +where + T: NativeDecimalType + From + From + Shl + BitOr, +{ + let mut value: T = msp.into(); + for part in lower { + value = (value << LOWER_PART_BITS) | part.into(); + } + value +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Constant; + 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::DecimalType; + use vortex_array::dtype::NativeDecimalType; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::match_each_decimal_value_type; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::assemble_decimal; + use crate::decimal_byte_parts::split_decimal; + + #[rstest] + #[case::empty_non_nullable(0, Validity::NonNullable)] + #[case::empty_nullable(0, Validity::AllValid)] + #[case::empty_all_null(0, Validity::AllInvalid)] + #[case::all_null(3, Validity::AllInvalid)] + #[case::all_null_array(3, Validity::Array(BoolArray::from_iter([false; 3]).into_array()))] + fn test_split_without_valid_rows( + #[case] len: usize, + #[case] validity: Validity, + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + values_type: DecimalType, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = match_each_decimal_value_type!(values_type, |T| { + DecimalArray::new( + Buffer::::zeroed(len), + DecimalDType::new(T::MAX_PRECISION, 0), + validity, + ) + }); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.msp.is::()); + assert!(parts.lower_parts.iter().all(|part| part.is::())); + assert_eq!(parts.msp.len(), len); + assert_eq!( + parts.msp.dtype().nullability(), + decimal.dtype().nullability() + ); + let round_tripped = round_trip(decimal.clone())?; + assert_eq!(round_tripped.values_type(), values_type); + assert_arrays_eq!(decimal, round_tripped, &mut ctx); + Ok(()) + } + + #[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)))] + #[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, 63, 64, 65, 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 mask = decimal.validity()?.execute_mask(len, &mut ctx)?; + let expected = PrimitiveArray::new( + mask.iter() + .map(|valid| if valid { u64::MAX } else { 0 }) + .collect::>(), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), if wide_256 { 3 } else { 1 }); + assert_eq!( + parts.msp.dtype(), + &DType::Primitive(PType::I64, decimal.dtype().nullability()) + ); + 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)?; + assemble_decimal( + &parts.msp, + &parts.lower_parts, + decimal.decimal_dtype(), + &mut ctx, + )? + .execute::(&mut ctx) + } + + #[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_reuses_values( + #[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![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(()) + } + + #[rstest] + #[case::signed(PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable))] + #[case::float(PrimitiveArray::new(buffer![0f32; 2], Validity::NonNullable))] + #[case::nullable_all_valid(PrimitiveArray::new(buffer![0u64; 2], Validity::AllValid))] + #[case::nullable_all_null(PrimitiveArray::new(buffer![0u64; 2], Validity::AllInvalid))] + #[case::nullable_mixed(PrimitiveArray::new(buffer![0u64; 2], Validity::from_iter([true, false])))] + fn test_assemble_rejects_invalid_lower_dtype( + #[case] invalid_lower: PrimitiveArray, + #[values(1, 2, 3)] lower_count: 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] = invalid_lower; + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + let lower = lower + .into_iter() + .map(IntoArray::into_array) + .collect::>(); + let mut ctx = array_session().create_execution_ctx(); + assert!(assemble_decimal(&msp.into_array(), &lower, dtype, &mut ctx).is_err()); + } + + #[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); + let lower = lower + .into_iter() + .map(IntoArray::into_array) + .collect::>(); + let mut ctx = array_session().create_execution_ctx(); + assert!(assemble_decimal(&msp.into_array(), &lower, dtype, &mut ctx).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| buffer![word; 2].into_array()); + let dtype = DecimalDType::new(76, 0); + let actual = assemble_decimal( + &msp.into_array(), + &lower[3 - lower_count..], + dtype, + &mut ctx, + )?; + 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(()) + } +} 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..e4d81c274da 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,17 +10,18 @@ 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; impl CastReduce for DecimalByteParts { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { - // Check if this is just a nullability change + // Check if this is just a nullability change. + // TODO(mk): Support non-nullability changes as well, e.g. precision. if !dtype.eq_ignore_nullability(array.dtype()) { 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 { - // Cannot cast decimal to non-decimal types - delegate to canonical form + // DecimalBytePartsArray can only have Decimal dtype. + let DType::Decimal(_, target_nullability) = dtype else { return Ok(None); }; @@ -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(), - )) + array.with_msp(new_msp).map(|a| Some(a.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..65933865780 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,15 @@ 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. + // + // TODO(mk): Compare the signed MSP and then the unsigned lower parts in significance + // order to avoid canonicalizing wide decimals. + 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 +167,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 +231,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..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 @@ -5,22 +5,17 @@ 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::DecimalBytePartsArrayExt; + 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())) + array + .map_parts(|part| part.filter(mask.clone())) + .map(|d| Some(d.into_array())) } } @@ -32,9 +27,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 +58,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..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 @@ -6,24 +6,17 @@ 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::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; 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(), - )) + 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 6c2d0dabb31..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)] @@ -19,10 +20,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 +97,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 +121,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/slice.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs similarity index 53% 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 14807421c73..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 @@ -7,23 +7,15 @@ 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::DecimalBytePartsArrayExt; 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(), - )) + 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 7a18f7bf91b..5b07af47252 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,103 @@ 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::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; -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 { + fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + // Taking with nullable indices makes every taken part nullable, but lower parts must + // stay non-nullable — validity belongs to the MSP alone. Fall back to the + // canonical path rather than rebuilding parts we would have to strip nullability from. + // + // TODO(mk): Support lower parts using fill_null for nullable indices. + if indices.dtype().is_nullable() && !array.lower_parts().is_empty() { + return Ok(None); + } + + array + .map_parts(|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 deleted file mode 100644 index 1e561b149fe..00000000000 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ /dev/null @@ -1,324 +0,0 @@ -// 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, -} - -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(), - } - } - - 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(), - } - } -} - -/// 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 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 -/// -/// 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::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_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_wide(&decimal.buffer::(), &mask, i256_to_parts); - DecimalParts::new(msp, lower, validity) - } - }) -} - -/// Split wide integers into a signed MSP and `N` unsigned lower parts. -/// -/// `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.map(BufferMut::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]); - - 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); - } - } - } - 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"), - } - - // 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)) -} - -/// Extract the high signed word and low unsigned word of an `i128`. -#[inline] -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] -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], - ) -} - -/// 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 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_wide::(msp, [first]), - decimal_dtype, - validity, - ), - [first, second] => DecimalArray::new( - assemble_wide::(msp, [first, second]), - decimal_dtype, - validity, - ), - [first, second, third] => DecimalArray::new( - assemble_wide::(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 `K` unsigned lower parts into wide integers. -/// -/// 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 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| { - 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() -} - -#[cfg(test)] -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 deleted file mode 100644 index 3e3de06c44e..00000000000 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs +++ /dev/null @@ -1,226 +0,0 @@ -// 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)))] -#[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, 63, 64, 65, 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(()) -} - -#[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(()) -} 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 951a893a785..919b7bb44a3 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -1,385 +1,59 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt::Display; -use std::fmt::Formatter; -use std::hash::Hasher; +//! Decimal byte-parts encoding. +//! +//! 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; - -impl ArrayHash for DecimalBytePartsData { - fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} -} - -impl ArrayEq for DecimalBytePartsData { - fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { - true - } -} - -#[derive(Clone, prost::Message)] -pub struct DecimalBytesPartsMetadata { - #[prost(enumeration = "PType", tag = "1")] - zeroth_child_ptype: i32, - #[prost(uint32, tag = "2")] - lower_part_count: u32, -} - -impl VTable for DecimalByteParts { - type TypedArrayData = DecimalBytePartsData; - - type OperationsVTable = Self; - type ValidityVTable = ValidityVTableFromChild; - - fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); - *ID - } - - fn validate( - &self, - _data: &Self::TypedArrayData, - dtype: &DType, - len: usize, - slots: &[Option], - ) -> VortexResult<()> { - 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) - } - - fn nbuffers(_array: ArrayView<'_, Self>) -> usize { - 0 - } - - fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - vortex_panic!("DecimalBytePartsArray buffer index {idx} out of bounds") - } - - fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") - } - - fn with_buffers( - &self, - array: ArrayView<'_, Self>, - buffers: &[BufferHandle], - ) -> VortexResult> { - vortex_array::vtable::with_empty_buffers(self, array, buffers) - } - - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DecimalBytesPartsMetadata { - zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, - lower_part_count: 0, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _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" - ); - - 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)) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - DecimalBytePartsSlots::NAMES[idx].to_string() - } - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - to_canonical_decimal(&array, ctx).map(ExecutionResult::done) - } -} - -#[array_slots(DecimalByteParts)] -pub struct DecimalBytePartsSlots { - /// The most significant parts of the decimal values. - #[slot(0)] - pub msp: ArrayRef, -} - -/// This array encodes decimals as between 1-4 columns of primitive typed children. -/// The most significant part (msp) sorting the most significant decimal bits. -/// This array must be signed and is nullable iff the decimal is nullable. -/// -/// e.g. for a decimal i128 \[ 127..64 | 64..0 \] msp = 127..64 and lower_part\[0\] = 64..0 -#[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, -} - -impl Display for DecimalBytePartsData { - fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { - Ok(()) - } -} - -pub struct DecimalBytePartsDataParts { - pub msp: ArrayRef, -} - -impl DecimalBytePartsData { - pub fn validate( - msp: &ArrayRef, - decimal_dtype: DecimalDType, - dtype: &DType, - len: usize, - ) -> VortexResult<()> { - if !msp.dtype().is_signed_int() { - vortex_bail!("decimal bytes parts, first part must be a signed array") - } - - let expected_dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); - vortex_ensure!( - dtype == &expected_dtype, - "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()); - vortex_ensure!( - msp_dtype.is_signed_int(), - "decimal bytes parts, first part must be a signed array" - ); - let _ = msp_len; - drop(expected_dtype); - Ok(Self { - _lower_parts: Vec::new(), - }) - } -} - -#[derive(Clone, Debug)] -pub struct DecimalByteParts; - -impl DecimalByteParts { - /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. - pub fn try_new( - msp: ArrayRef, - decimal_dtype: DecimalDType, - ) -> VortexResult { - 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), - ) - }) - } -} - -/// 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() - })) -} - -impl OperationsVTable for DecimalByteParts { - type ProbeState = (); +mod array; +mod assemble; +pub(crate) mod compute; +#[cfg(test)] +mod prop_tests; +mod rules; +mod split; +#[cfg(test)] +mod testing; - fn scalar_at( - array: ArrayView<'_, DecimalByteParts>, - index: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - // TODO(joe): support parts len != 1 - let scalar = array.msp().execute_scalar(index, ctx)?; +pub use array::*; +pub use split::DecimalParts; +pub use split::dbp_encode; +pub use split::split_decimal; - // 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))), - ) - } +#[doc(hidden)] +pub mod _benchmarking { + pub use super::assemble::assemble_decimal; + pub use super::assemble::assemble_wide_decimal; + pub use super::split::i128_to_parts; + pub use super::split::i256_to_parts; + pub use super::split::split_wide; } -impl ValidityChild for DecimalByteParts { - fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { - // validity stored in 0th child - array.msp().clone() - } -} +/// The maximum number of 64-bit lower parts an encoded `i128` decimal can carry. +const MAX_I128_LOWER_PARTS: usize = 1; -#[cfg(test)] -mod tests { - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::array_session; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::dtype::DType; - use vortex_array::dtype::DecimalDType; - use vortex_array::dtype::Nullability; - use vortex_array::scalar::DecimalValue; - use vortex_array::scalar::Scalar; - use vortex_array::scalar::ScalarValue; - use vortex_array::validity::Validity; - use vortex_buffer::buffer; +/// The maximum number of 64-bit lower parts an encoded `i256` decimal can carry. +const MAX_I256_LOWER_PARTS: usize = 3; - use crate::DecimalByteParts; +/// The maximum number of 64-bit lower parts an encoded decimal can carry. +const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS; - #[test] - fn test_scalar_at_decimal_parts() { - let decimal_dtype = DecimalDType::new(8, 2); - let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); - let array = DecimalByteParts::try_new( - PrimitiveArray::new( - buffer![100i32, 200i32, 400i32], - Validity::Array(BoolArray::from_iter(vec![false, true, true]).into_array()), - ) - .into_array(), - decimal_dtype, - ) - .unwrap() - .into_array(); +/// Number of bits stored in each lower part. +const LOWER_PART_BITS: usize = 64; - assert_eq!( - Scalar::null(dtype.clone()), - array - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new( - dtype.clone(), - Some(ScalarValue::Decimal(DecimalValue::I64(200))) - ) - .unwrap(), - array - .execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), - array - .execute_scalar(2, &mut array_session().create_execution_ctx()) - .unwrap() - ); - } -} +/// Dtype of the lower parts produced by splitting, before any narrowing. +const LOWER_PART_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs new file mode 100644 index 00000000000..5630d1706a8 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs @@ -0,0 +1,183 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Property tests for decimal byte-parts round trips. + +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_error::VortexExpect; + +use super::DecimalByteParts; +use super::DecimalBytePartsArray; +use super::testing::encode; + +/// 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(); + crate::initialize(&session); + session.create_execution_ctx() +} + +/// 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).vortex_expect("encode").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).vortex_expect("encode").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/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/split.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs new file mode 100644 index 00000000000..a3aa1b4afdc --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -0,0 +1,310 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splitting canonical decimal arrays into signed and unsigned parts. + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::DecimalByteParts; +use super::DecimalBytePartsArray; +use super::LOWER_PART_BITS; +use super::MAX_I128_LOWER_PARTS; +use super::MAX_I256_LOWER_PARTS; + +/// Create a [`DecimalBytePartsArray`] from a [`DecimalArray`] by splitting it into parts. +/// +/// # Errors +/// +/// Returns an error if the decimal cannot be split. +pub fn dbp_encode( + decimal: &DecimalArray, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let parts = split_decimal(decimal, exec_ctx)?; + // SAFETY: splitting produces a signed MSP and zero, one, or three non-nullable u64 lower + // parts, all with the decimal's length and in most-significant-first order. This also holds + // for the constant parts used for empty and all-null inputs. The decimal dtype is preserved. + Ok(unsafe { + DecimalByteParts::new_unchecked(parts.msp, parts.lower_parts, decimal.decimal_dtype()) + }) +} + +/// A decimal array decomposed into byte parts. +pub struct DecimalParts { + /// The signed most significant part. This carries the validity of the whole array. + pub msp: ArrayRef, + /// The unsigned 64-bit lower parts, most significant first. + pub lower_parts: Vec, +} + +impl DecimalParts { + /// Construct decimal parts from the MSP buffer constituting a narrow decimal (`i64` or narrower). + /// Narrow decimals have an MSP at most as wide as `i64` and no lower parts. + fn from_narrow(values: Buffer, validity: Validity) -> Self { + Self { + msp: PrimitiveArray::new(values, validity).into_array(), + lower_parts: Vec::new(), + } + } + + /// Construct decimal parts arrays from the buffers constituting a wide decimal (`i128` or `i256`). + /// Wide decimals have an `i64` MSP and up to [`super::MAX_LOWER_PARTS`] `u64` lower parts. + fn from_wide( + 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(), + } + } +} + +/// 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 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. +/// Empty and all-null arrays use constant parts, preserving the part types and MSP's nullability. +/// +/// # Errors +/// +/// 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()?; + let len = decimal.len(); + let mask = validity.execute_mask(len, ctx)?; + + if mask.all_false() || decimal.is_empty() { + return Ok(split_no_valid_row(decimal, &validity)); + } + + Ok(match decimal.values_type() { + DecimalType::I8 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I16 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I32 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I64 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I128 => { + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i128_to_parts); + DecimalParts::from_wide(msp, lower, validity) + } + DecimalType::I256 => { + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i256_to_parts); + DecimalParts::from_wide(msp, lower, validity) + } + }) +} + +/// Splits decimals with no valid rows (all null or empty) into constant decimal parts with the +/// corresponding nullability. +fn split_no_valid_row(decimal: &DecimalArray, validity: &Validity) -> DecimalParts { + let (msp_ptype, lower_part_count) = match decimal.values_type() { + DecimalType::I8 => (PType::I8, 0), + DecimalType::I16 => (PType::I16, 0), + DecimalType::I32 => (PType::I32, 0), + DecimalType::I64 => (PType::I64, 0), + DecimalType::I128 => (PType::I64, MAX_I128_LOWER_PARTS), + DecimalType::I256 => (PType::I64, MAX_I256_LOWER_PARTS), + }; + // Empty masks are also all-false. The default scalar is null for nullable inputs + // and zero for non-nullable empty inputs, preserving the MSP's nullability. + let msp = Scalar::default_value(&DType::Primitive(msp_ptype, validity.nullability())); + let len = decimal.len(); + DecimalParts { + msp: ConstantArray::new(msp, len).into_array(), + lower_parts: vec![ConstantArray::new(0u64, len).into_array(); lower_part_count], + } +} + +/// Split wide integers into a signed MSP and `N` unsigned lower parts. +/// +/// `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. The caller handles empty +/// and all-null arrays before calling this function. +pub fn split_wide( + values: &[T], + 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)); + + // 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]); + + 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); + } + } + } + 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!("all-null arrays are handled by split_decimal"), + } + + // 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)) +} + +/// Extract the high signed word and low unsigned word of an `i128`. +#[inline] +pub const fn i128_to_parts(value: i128) -> (i64, [u64; MAX_I128_LOWER_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" + )] + ((value >> LOWER_PART_BITS) as i64, [value as u64]) +} + +/// Extract the signed MSP and three unsigned lower words of an `i256`. +#[inline] +pub const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_I256_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], + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + 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::PType; + use vortex_array::dtype::i256; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::split_decimal; + use crate::decimal_byte_parts::LOWER_PART_DTYPE; + use crate::decimal_byte_parts::MAX_LOWER_PARTS; + + #[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(()) + } + + #[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(()) + } +} 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..2dfe2a55b3c --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared fixtures for decimal byte-parts tests. + +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 super::DecimalBytePartsArray; +use super::dbp_encode; + +/// Encode a canonical decimal array as byte parts, splitting wide values into lower parts. +pub(crate) fn encode(decimal: &DecimalArray) -> VortexResult { + dbp_encode(decimal, &mut array_session().create_execution_ctx()) +} + +/// 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(super) fn i256_of(high: i128, low: u128) -> i256 { + i256::from_parts(low, high) +} diff --git a/vortex-array/src/dtype/bigint/mod.rs b/vortex-array/src/dtype/bigint/mod.rs index 47195526b1f..b03bdef43c2 100644 --- a/vortex-array/src/dtype/bigint/mod.rs +++ b/vortex-array/src/dtype/bigint/mod.rs @@ -130,6 +130,20 @@ impl From for i256 { } } +impl From for i256 { + #[inline] + fn from(value: i64) -> Self { + Self::from_i128(value.into()) + } +} + +impl From for i256 { + #[inline] + fn from(value: u64) -> Self { + Self::from_i128(value.into()) + } +} + impl Display for i256 { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.0) diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 21e440a7436..5d3f610afc7 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -403,7 +403,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 @@ -439,6 +439,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 c16771eec8a1a8e4242d2bef2d19f3b2d4c013d8 Mon Sep 17 00:00:00 2001 From: Matthew Katz <87445739+mhk197@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:04:50 -0400 Subject: [PATCH 3/5] Add `ArrayPlugin` to enable serde for new `DecimalBytePartsArray` format (#9810) Add serialization for `DecimalBytePartsArray` with lower parts while preserving the frozen v1 file format. `DecimalBytePartsPlugin` owns both formats and selects the serialized ID from the array's layout: | Array layout | Serialized ID | | --- | --- | | Signed most significant part (MSP) only | `vortex.decimal_byte_parts` | | Signed MSP plus one to three unsigned lower parts | `vortex.decimal_byte_parts_v2` | Both formats deserialize into the same in-memory array representation. Each lower part represents a 64-bit window, but its storage dtype may be narrowed to `u8`, `u16`, or `u32` when its values fit. - Record the MSP's physical integer type, lower-part count, and each lower part's physical integer type in metadata. Deserialization validates these types and restores each child with its recorded dtype. - Keep the frozen format restricted to zero lower parts and require at least one for v2. Frozen metadata remains byte-identical. - Make DBP's VTable `serialize` and `deserialize` methods return errors directing callers to `DecimalBytePartsPlugin`. Keep DBP metadata and serde helpers in `plugin.rs` so the dedicated plugin owns both formats. - Add a separate v2 compatibility fixture covering positive and negative wide `i128` values and nullable `i256` values. The frozen fixture is unchanged. **Breaking:** directly registering `DecimalByteParts`, or calling its VTable serde methods, now errors for both serialization and deserialization, including the frozen v1 format. Replace: ```rust session.arrays().register(DecimalByteParts); ``` with: ```rust session.arrays().register(DecimalBytePartsPlugin); ``` `vortex_decimal_byte_parts::initialize(&session)` already registers the plugin. Existing v1 files remain readable through it; no file migration is required. --------- Signed-off-by: Matt Katz --- docs/specs/editions.md | 4 +- .../src/decimal_byte_parts/array.rs | 205 +---------- .../src/decimal_byte_parts/compute/filter.rs | 12 +- .../src/decimal_byte_parts/compute/mod.rs | 11 +- .../src/decimal_byte_parts/compute/take.rs | 13 +- .../src/decimal_byte_parts/mod.rs | 6 +- .../src/decimal_byte_parts/plugin/mod.rs | 98 ++++++ .../src/decimal_byte_parts/plugin/tests.rs | 324 ++++++++++++++++++ .../src/decimal_byte_parts/plugin/v1.rs | 84 +++++ .../src/decimal_byte_parts/plugin/v2.rs | 114 ++++++ .../src/decimal_byte_parts/prop_tests.rs | 13 +- .../src/decimal_byte_parts/split.rs | 33 +- .../src/decimal_byte_parts/testing.rs | 31 +- encodings/decimal-byte-parts/src/lib.rs | 4 +- vortex-btrblocks/src/trace_tests.rs | 10 +- .../golden__compact__decimal_prices.snap | 2 +- .../golden__regular__decimal_prices.snap | 2 +- .../encodings/decimal_byte_parts_v2.rs | 105 ++++++ .../arrays/synthetic/encodings/mod.rs | 2 + 19 files changed, 809 insertions(+), 264 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs create mode 100644 vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs diff --git a/docs/specs/editions.md b/docs/specs/editions.md index ee5f494861f..1aa3cfe4e11 100644 --- a/docs/specs/editions.md +++ b/docs/specs/editions.md @@ -172,9 +172,9 @@ representation gains support for wide decimals, represented by a signed most-sig represented that way, it emits `vortex.decimal_byte_parts` with `lower_part_count = 0`, even if the current in-memory array has lower-part children. - An array that cannot be collapsed into that old form losslessly uses the new - `vortex.decimal_byte_parts_v2` component, initially staged in a draft edition. + `vortex.decimal_byte_parts.v2` component, initially staged in a draft edition. - A new reader deserializes both IDs into the same in-memory representation. An older reader reports - `vortex.decimal_byte_parts_v2` as unknown instead of trying to decode a wire format it does not support. + `vortex.decimal_byte_parts.v2` as unknown instead of trying to decode a wire format it does not support. - When targeting an edition that permits only the old ID, serializing a value that can be collapsed succeeds; an irreducibly multi-part value fails because no lossless downgrade exists. diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs index bcdb3579f53..7059fd09838 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -7,14 +7,12 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; use vortex_array::ArrayRef; -use vortex_array::ArraySlots; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; @@ -24,7 +22,6 @@ use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::PType; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -37,92 +34,18 @@ 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 super::LOWER_PART_DTYPE; use super::MAX_LOWER_PARTS; use super::assemble::assemble_decimal; use super::assemble::assemble_wide_decimal_value; +use super::decimal_byte_parts_v2_id; use super::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. pub type DecimalBytePartsArray = Array; -#[derive(Clone, prost::Message)] -pub struct DecimalBytesPartsMetadata { - #[prost(enumeration = "PType", tag = "1")] - zeroth_child_ptype: i32, - #[prost(uint32, tag = "2")] - 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) - } -} - /// 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 @@ -222,7 +145,7 @@ impl DecimalByteParts { /// /// Lower parts are ordered most significant first and must each be a non-nullable unsigned integer /// array of the same length as the MSP. See [`super::split_decimal`] for producing them from a - /// canonical decimal array. + /// decimal array. /// /// # Errors /// @@ -274,8 +197,7 @@ impl VTable for DecimalByteParts { type ValidityVTable = ValidityVTableFromChild; fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); - *ID + decimal_byte_parts_v2_id() } fn validate( @@ -331,33 +253,22 @@ impl VTable for DecimalByteParts { } fn serialize( - array: ArrayView<'_, Self>, + _array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - vortex_ensure!( - array.lower_parts().is_empty(), - "serializing DecimalByteParts with lower parts is not supported" - ); - Ok(Some( - DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), - )) + vortex_bail!("DecimalByteParts serialization requires DecimalBytePartsPlugin") } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], + _dtype: &DType, + _len: usize, + _metadata: &[u8], _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + _children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - let metadata = DecimalBytesPartsMetadata::decode(metadata)?; - vortex_ensure!( - metadata.lower_part_count()? == 0, - "vortex.decimal_byte_parts must not carry lower parts" - ); - metadata.into_array_parts(dtype, len, children) + vortex_bail!("DecimalByteParts deserialization requires DecimalBytePartsPlugin") } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -510,18 +421,15 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; - use vortex_array::vtable::VTable; use vortex_buffer::buffer; use vortex_error::VortexResult; use super::DecimalByteParts; - use super::DecimalBytePartsArray; use super::DecimalBytePartsArraySlotsExt; use super::DecimalBytePartsData; use crate::decimal_byte_parts::LOWER_PART_DTYPE; use crate::decimal_byte_parts::MAX_LOWER_PARTS; 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] @@ -563,62 +471,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))] - 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); @@ -639,7 +491,7 @@ mod tests { #[test] fn test_lower_part_layout_i256() -> VortexResult<()> { let array = i256_parts( - vec![i256_of((5i128 << 64) | 6, (7u128 << 64) | 8)], + vec![i256::from_parts((7u128 << 64) | 8, (5i128 << 64) | 6)], Validity::NonNullable, ); assert_eq!(array.lower_parts().len(), MAX_LOWER_PARTS); @@ -654,27 +506,6 @@ mod tests { 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, @@ -860,7 +691,7 @@ mod tests { 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 i256_array = i256_parts(vec![i256::from_parts(0, 1 << 100)], Validity::NonNullable); let canonical = i256_array.into_array().execute::(&mut ctx)?; assert_eq!(canonical.values_type(), DecimalType::I256); @@ -885,7 +716,10 @@ mod tests { )?; 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)]); + assert_eq!( + canonical.buffer::().as_slice(), + &[i256::from_parts(9, 1)] + ); Ok(()) } @@ -912,11 +746,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/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index 49c4021dd18..6921e5dfda4 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 @@ -27,12 +27,12 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::i256; 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] @@ -72,11 +72,11 @@ mod test { 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), + i256::from_parts(0, 1), + i256::from_parts(5, -1), + i256::from_parts(u128::MAX, 0), + i256::from_parts(7, 1 << 64), + i256::from_parts(0, 0), ], Validity::from_iter([true, false, true, true, false]), ); 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 f9848e1b2e7..c8385c2d6d6 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 @@ -27,7 +27,6 @@ mod tests { 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. @@ -43,11 +42,11 @@ mod tests { 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), + i256::from_parts(0, 1), + i256::from_parts(0, -1), + i256::from_parts(u128::MAX, 0), + i256::from_parts(7, 1 << 64), + i256::from_parts(0, 0), ] } 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 5b07af47252..a74915e9f22 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 @@ -38,14 +38,13 @@ mod tests { 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_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 @@ -61,7 +60,9 @@ mod tests { Validity::NonNullable, ); let indices = buffer![0u64, 2].into_array(); - let taken = encode(&decimal)?.into_array().take(indices)?; + let taken = DecimalByteParts::encode(&decimal, &mut session.create_execution_ctx())? + .into_array() + .take(indices)?; assert!( taken.is::(), @@ -80,7 +81,7 @@ mod tests { Validity::NonNullable, ))] #[case::three_lower_parts(DecimalArray::new( - Buffer::from(vec![i256_of(1, 1 << 70), i256_of(0, 2), i256_of(0, 3)]), + Buffer::from(vec![i256::from_parts(1 << 70, 1), i256::from_parts(2, 0), i256::from_parts(3, 0)]), DecimalDType::new(76, 2), Validity::NonNullable, ))] @@ -96,7 +97,9 @@ mod tests { .take(indices.clone())? .execute::(&mut ctx)?; - let taken = encode(&decimal)?.into_array().take(indices)?; + let taken = DecimalByteParts::encode(&decimal, &mut ctx)? + .into_array() + .take(indices)?; let actual = taken.execute::(&mut ctx)?; assert_arrays_eq!(expected, actual, &mut ctx); 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 919b7bb44a3..65bb8222f7a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -22,6 +22,7 @@ use vortex_array::dtype::PType; mod array; mod assemble; pub(crate) mod compute; +mod plugin; #[cfg(test)] mod prop_tests; mod rules; @@ -30,8 +31,11 @@ mod split; mod testing; pub use array::*; +pub use plugin::DecimalBytePartsPlugin; +pub use plugin::DecimalBytePartsV2Metadata; +pub use plugin::decimal_byte_parts_v1_id; +pub use plugin::decimal_byte_parts_v2_id; pub use split::DecimalParts; -pub use split::dbp_encode; pub use split::split_decimal; #[doc(hidden)] diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs new file mode 100644 index 00000000000..3896cd51c96 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! ArrayPlugin implementation for DBP that handles different wire formats. + +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; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::DecimalByteParts; +use super::DecimalBytePartsArraySlotsExt; + +#[cfg(test)] +mod tests; + +mod v1; +mod v2; + +pub use v2::DecimalBytePartsV2Metadata; + +/// The frozen single-child DBP serialized ID. +pub fn decimal_byte_parts_v1_id() -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); + *ID +} + +/// The current in-memory DBP ID and serialized ID for arrays with lower parts. +pub fn decimal_byte_parts_v2_id() -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts.v2"); + *ID +} + +/// Serde for the [`DecimalByteParts`] array using the frozen v1 and v2 wire formats. +/// +/// Each version owns its metadata schema and serde functions. The plugin writes v1 whenever an +/// array has no lower parts, so such arrays stay readable by older readers, and v2 otherwise. +/// The v2 format itself accepts any lower part count up to the maximum. +/// +/// Register this plugin, or call [`crate::initialize`], to enable both formats. Direct registration +/// of [`DecimalByteParts`] does not support serde. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsPlugin; + +impl ArrayPlugin for DecimalBytePartsPlugin { + fn id(&self) -> ArrayId { + VTable::id(&DecimalByteParts) + } + + fn serialized_ids(&self) -> Vec { + vec![decimal_byte_parts_v1_id(), 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 = if view.lower_parts().is_empty() { + v1::serialize(view)? + } else { + v2::serialize(view)? + }; + Ok(Some(serialized)) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let array = if parts.serialized_id == decimal_byte_parts_v1_id() { + v1::deserialize(parts)? + } else if parts.serialized_id == decimal_byte_parts_v2_id() { + v2::deserialize(parts)? + } else { + vortex_bail!( + "DecimalByteParts plugin does not recognize serialized ID {}", + parts.serialized_id + ) + }; + Ok(array.into_array()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs new file mode 100644 index 00000000000..92c9d1edbb7 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs @@ -0,0 +1,324 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message as _; +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayVTable; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +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::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::MAX_LOWER_PARTS; + +#[rstest] +#[case::no_lower_parts(DecimalByteParts::try_new( + buffer![1i32, 2, 3].into_array(), DecimalDType::new(9, 2), +))] +#[case::one_lower_part(DecimalByteParts::try_new_with_lower_parts( + msp(), vec![lower_part()], DecimalDType::new(38, 2), +))] +#[case::wider_i64_storage(DecimalByteParts::encode( + &DecimalArray::new(buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable), + &mut array_session().create_execution_ctx(), +))] +#[case::wider_i128_storage(DecimalByteParts::encode( + &DecimalArray::new(buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable), + &mut array_session().create_execution_ctx(), +))] +#[case::wider_i256_storage(DecimalByteParts::encode( + &DecimalArray::new( + buffer![i256::from_i128(-99), i256::ZERO, i256::from_i128(99)], + DecimalDType::new(2, 0), Validity::NonNullable, + ), + &mut array_session().create_execution_ctx(), +))] +#[case::redundant_lower_parts(DecimalByteParts::try_new_with_lower_parts( + buffer![0i64; 3].into_array(), + vec![buffer![0u64; 3].into_array(), buffer![0u64; 3].into_array(), lower_part()], + DecimalDType::new(38, 2), +))] +#[case::narrowed_lower_parts(DecimalByteParts::try_new_with_lower_parts( + buffer![-1i16, 0, 1].into_array(), + vec![ + buffer![u8::MAX, 128, 0].into_array(), + ConstantArray::new(u16::MAX, 3).into_array(), + buffer![0u32, 1 << 31, u32::MAX].into_array(), + ], + DecimalDType::new(76, 2), +))] +#[case::nullable_mixed_lower_parts(DecimalByteParts::try_new_with_lower_parts( + PrimitiveArray::new( + buffer![-1i8, 0, 1], Validity::from_iter([true, false, true]), + ).into_array(), + vec![ + buffer![u64::MAX, 1 << 63, 0].into_array(), + buffer![0u8, 128, u8::MAX].into_array(), + buffer![u32::MAX, 1 << 31, 0].into_array(), + ], + DecimalDType::new(76, 2), +))] +fn serde_round_trip(#[case] array: VortexResult) -> VortexResult<()> { + let session = session(); + let array = array?; + let lower_part_count = array.lower_parts().len(); + let lower_part_dtypes: Vec<_> = array + .lower_parts() + .iter() + .map(|part| part.dtype().clone()) + .collect(); + let array = array.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + + let expected_id = if lower_part_count == 0 { + decimal_byte_parts_v1_id() + } 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() + .iter() + .map(|part| part.dtype().clone()) + .collect::>(), + lower_part_dtypes, + "lower-part dtypes and order must survive serde" + ); + + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) +} + +#[test] +fn v1_metadata_is_unchanged() -> VortexResult<()> { + let session = session(); + let array = DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2))?.into_array(); + let serialized = session + .array_serialize(&array)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialized.serialized_id, decimal_byte_parts_v1_id()); + // v1 metadata for an i64 MSP: field 1 = 7, with no lower-part fields emitted. + assert_eq!(serialized.metadata, [8, 7]); + Ok(()) +} + +#[rstest] +#[case::v1_lower_part_count( + decimal_byte_parts_v1_id(), + v1_metadata(1), + vec![msp(), lower_part()], + "must not carry lower parts" +)] +#[case::v1_extra_child( + decimal_byte_parts_v1_id(), + v1_metadata(0), + vec![msp(), lower_part()], + "exactly one child" +)] +#[case::v2_missing_child( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::U64 as i32]), + vec![msp()], + "expected 2 children, got 1" +)] +#[case::v2_extra_child( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::U64 as i32]), + vec![msp(), lower_part(), lower_part()], + "expected 2 children, got 3" +)] +#[case::v2_too_many_lower_parts( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::U64 as i32; MAX_LOWER_PARTS + 1]), + vec![msp(), lower_part(), lower_part(), lower_part(), lower_part()], + "lower parts, got 4" +)] +#[case::v2_signed_lower_part( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::I64 as i32]), + vec![msp(), lower_part()], + "unsigned integer dtype" +)] +#[case::v2_unknown_ptype( + decimal_byte_parts_v2_id(), + v2_metadata(vec![i32::MAX]), + vec![msp(), lower_part()], + "invalid PType" +)] +fn decoder_rejects_malformed_payloads( + #[case] serialized_id: ArrayId, + #[case] metadata: Vec, + #[case] children: Vec, + #[case] expected_error: &str, +) { + let result = deserialize_with(serialized_id, &metadata, children); + assert!( + result + .as_ref() + .is_err_and(|err| err.to_string().contains(expected_error)), + "expected {expected_error}, got {result:?}" + ); +} + +#[test] +fn serialization_requires_v2_permission() -> VortexResult<()> { + let session = session(); + 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( + [decimal_byte_parts_v1_id(), ArrayVTable::id(&Primitive)] + .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( + [ + decimal_byte_parts_v1_id(), + decimal_byte_parts_v2_id(), + ArrayVTable::id(&Primitive), + ] + .into_iter() + .collect(), + ); + array.serialize(&permissive, &session, &SerializeOptions::default())?; + 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_serde() -> VortexResult<()> { + let session = array_session(); + session.arrays().register(DecimalByteParts); + let msp = msp(); + let array = DecimalByteParts::try_new(msp.clone(), DecimalDType::new(19, 2))?.into_array(); + let result = session.array_serialize(&array); + assert!( + result.as_ref().is_err_and(|err| err + .to_string() + .contains("DecimalByteParts serialization requires DecimalBytePartsPlugin")), + "expected unsupported VTable serialization, got {result:?}" + ); + + let id = VTable::id(&DecimalByteParts); + let plugin = session + .arrays() + .registry() + .get(&id) + .vortex_expect("registered"); + let children = vec![msp]; + let result = plugin.deserialize( + ArrayDeserialization::new(id, array.dtype(), array.len(), &[8, 7], &[], &children), + &session, + ); + assert!( + result.as_ref().is_err_and(|err| err + .to_string() + .contains("DecimalByteParts deserialization requires DecimalBytePartsPlugin")), + "expected unsupported VTable deserialization, got {result:?}" + ); + Ok(()) +} + +fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() +} + +fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() +} + +/// v1 metadata for an i64 MSP: field 1 = 7, then field 2 only when the count is non-zero, as +/// proto3 omits default values. +fn v1_metadata(lower_part_count: u8) -> Vec { + let mut metadata = vec![8, 7]; + if lower_part_count > 0 { + metadata.extend([16, lower_part_count]); + } + metadata +} + +fn v2_metadata(lower_part_ptypes: Vec) -> Vec { + DecimalBytePartsV2Metadata { + msp_ptype: PType::I64 as i32, + lower_part_ptypes, + } + .encode_to_vec() +} + +fn session() -> VortexSession { + let session = array_session(); + crate::initialize(&session); + session +} + +/// Decode a hand-built payload of three rows through the plugin. +fn deserialize_with( + serialized_id: ArrayId, + metadata: &[u8], + children: Vec, +) -> VortexResult { + let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); + DecimalBytePartsPlugin.deserialize( + ArrayDeserialization::new(serialized_id, &dtype, 3, metadata, &[], &children), + &array_session(), + ) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs new file mode 100644 index 00000000000..ec816c62736 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serde for the frozen single-child DBP wire format. + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayParts; +use vortex_array::ArraySerialization; +use vortex_array::ArrayView; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::decimal_byte_parts_v1_id; +use crate::DecimalByteParts; +use crate::DecimalBytePartsArray; +use crate::DecimalBytePartsArraySlotsExt; +use crate::DecimalBytePartsData; + +#[derive(Clone, prost::Message)] +struct DecimalBytePartsMetadata { + #[prost(enumeration = "PType", tag = "1")] + zeroth_child_ptype: i32, + #[prost(uint32, tag = "2")] + lower_part_count: u32, +} + +pub(super) fn serialize( + array: ArrayView<'_, DecimalByteParts>, +) -> VortexResult { + vortex_ensure!( + array.lower_parts().is_empty(), + "v1 must not carry lower parts" + ); + let msp = array.msp(); + let metadata = DecimalBytePartsMetadata { + zeroth_child_ptype: PType::try_from(msp.dtype())? as i32, + lower_part_count: 0, + } + .encode_to_vec(); + Ok(ArraySerialization::new( + decimal_byte_parts_v1_id(), + metadata, + vec![], + vec![msp.clone()], + )) +} + +pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult { + vortex_ensure!( + parts.serialized_id == decimal_byte_parts_v1_id(), + "expected the v1 format" + ); + let metadata = DecimalBytePartsMetadata::decode(parts.metadata)?; + vortex_ensure!( + parts.dtype.as_decimal_opt().is_some(), + "expected a decimal dtype" + ); + vortex_ensure!( + metadata.lower_part_count == 0, + "v1 must not carry lower parts" + ); + vortex_ensure!(parts.children.len() == 1, "v1 must carry exactly one child"); + let ptype = PType::try_from(metadata.zeroth_child_ptype)?; + vortex_ensure!( + ptype.is_signed_int(), + "MSP must have a signed integer dtype" + ); + let encoded_dtype = DType::Primitive(ptype, parts.dtype.nullability()); + let msp = parts.children.get(0, &encoded_dtype, parts.len)?; + Array::try_from_parts( + ArrayParts::new( + DecimalByteParts, + parts.dtype.clone(), + parts.len, + DecimalBytePartsData, + ) + .with_slots(smallvec![Some(msp)]), + ) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs new file mode 100644 index 00000000000..c82db1cbce3 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serde for DBP values with unsigned lower parts. + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayParts; +use vortex_array::ArraySerialization; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use super::decimal_byte_parts_v2_id; +use crate::DecimalByteParts; +use crate::DecimalBytePartsArray; +use crate::DecimalBytePartsArraySlotsExt; +use crate::DecimalBytePartsData; +use crate::decimal_byte_parts::MAX_LOWER_PARTS; + +/// Metadata for decimal byte parts with lower parts. +#[derive(Clone, prost::Message)] +pub struct DecimalBytePartsV2Metadata { + /// Ptype of the most significant part. + #[prost(enumeration = "PType", tag = "1")] + pub(super) msp_ptype: i32, + /// Ptypes of the lower parts, ordered most significant first. + #[prost(enumeration = "PType", repeated, tag = "2")] + pub(super) lower_part_ptypes: Vec, +} + +pub(super) fn serialize( + array: ArrayView<'_, DecimalByteParts>, +) -> VortexResult { + let lower_parts = array.lower_parts(); + + let metadata = DecimalBytePartsV2Metadata { + msp_ptype: PType::try_from(array.msp().dtype())? as i32, + lower_part_ptypes: lower_parts + .iter() + .map(|part| PType::try_from(part.dtype()).map(|ptype| ptype as i32)) + .collect::>()?, + } + .encode_to_vec(); + + let mut children = Vec::with_capacity(1 + lower_parts.len()); + children.push(array.msp().clone()); + children.extend(lower_parts.iter().cloned()); + + Ok(ArraySerialization::new( + decimal_byte_parts_v2_id(), + metadata, + vec![], + children, + )) +} + +pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult { + let metadata = DecimalBytePartsV2Metadata::decode(parts.metadata)?; + vortex_ensure!( + parts.dtype.as_decimal_opt().is_some(), + "expected a decimal dtype" + ); + + let lower_part_count = metadata.lower_part_ptypes.len(); + vortex_ensure!( + lower_part_count <= MAX_LOWER_PARTS, + "v2 carries at most {MAX_LOWER_PARTS} lower parts, got {lower_part_count}" + ); + vortex_ensure!( + parts.children.len() == 1 + lower_part_count, + "expected {} children, got {}", + 1 + lower_part_count, + parts.children.len() + ); + + let msp_ptype = PType::try_from(metadata.msp_ptype)?; + vortex_ensure!( + msp_ptype.is_signed_int(), + "MSP must have a signed integer dtype, got {msp_ptype}" + ); + let msp_dtype = DType::Primitive(msp_ptype, parts.dtype.nullability()); + + let mut slots = ArraySlots::with_capacity(parts.children.len()); + slots.push(Some(parts.children.get(0, &msp_dtype, parts.len)?)); + for (idx, raw_ptype) in metadata.lower_part_ptypes.into_iter().enumerate() { + let ptype = PType::try_from(raw_ptype) + .map_err(|_| vortex_err!("invalid PType {raw_ptype} for lower part {idx}"))?; + vortex_ensure!( + ptype.is_unsigned_int(), + "lower part {idx} must have an unsigned integer dtype, got {ptype}" + ); + slots.push(Some(parts.children.get( + 1 + idx, + &DType::Primitive(ptype, Nullability::NonNullable), + parts.len, + )?)); + } + Array::try_from_parts( + ArrayParts::new( + DecimalByteParts, + parts.dtype.clone(), + parts.len, + DecimalBytePartsData, + ) + .with_slots(slots), + ) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs index 5630d1706a8..b1625485ee7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs @@ -21,7 +21,6 @@ use vortex_error::VortexExpect; use super::DecimalByteParts; use super::DecimalBytePartsArray; -use super::testing::encode; /// Largest magnitude a `Decimal(38, _)` can hold: 38 nines. const MAX_I128: i128 = 10i128.pow(38) - 1; @@ -139,10 +138,8 @@ fn decoded_survives_encode_then_decode(tc: TestCase) { let decimal = draw_decimal(&tc); let mut ctx = ctx(); - let round_tripped = canonicalize( - encode(&decimal).vortex_expect("encode").into_array(), - &mut ctx, - ); + let encoded = DecimalByteParts::encode(&decimal, &mut ctx).vortex_expect("encode"); + let round_tripped = canonicalize(encoded.into_array(), &mut ctx); assert_eq!(round_tripped.values_type(), decimal.values_type()); assert_arrays_eq!(decimal, round_tripped, &mut ctx); @@ -160,10 +157,8 @@ fn encoded_survives_decode_then_encode(tc: TestCase) { let mut ctx = ctx(); let decoded = canonicalize(array.into_array(), &mut ctx); - let re_decoded = canonicalize( - encode(&decoded).vortex_expect("encode").into_array(), - &mut ctx, - ); + let re_encoded = DecimalByteParts::encode(&decoded, &mut ctx).vortex_expect("encode"); + let re_decoded = canonicalize(re_encoded.into_array(), &mut ctx); assert_arrays_eq!(decoded, re_decoded, &mut ctx); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs index a3aa1b4afdc..8d8818924e9 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -27,22 +27,23 @@ use super::LOWER_PART_BITS; use super::MAX_I128_LOWER_PARTS; use super::MAX_I256_LOWER_PARTS; -/// Create a [`DecimalBytePartsArray`] from a [`DecimalArray`] by splitting it into parts. -/// -/// # Errors -/// -/// Returns an error if the decimal cannot be split. -pub fn dbp_encode( - decimal: &DecimalArray, - exec_ctx: &mut ExecutionCtx, -) -> VortexResult { - let parts = split_decimal(decimal, exec_ctx)?; - // SAFETY: splitting produces a signed MSP and zero, one, or three non-nullable u64 lower - // parts, all with the decimal's length and in most-significant-first order. This also holds - // for the constant parts used for empty and all-null inputs. The decimal dtype is preserved. - Ok(unsafe { - DecimalByteParts::new_unchecked(parts.msp, parts.lower_parts, decimal.decimal_dtype()) - }) +impl DecimalByteParts { + /// Encode a [`DecimalArray`] as byte parts, splitting wide values into lower parts. + /// + /// # Errors + /// + /// Returns an error if the decimal cannot be split. + pub fn encode( + decimal: &DecimalArray, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let parts = split_decimal(decimal, exec_ctx)?; + // SAFETY: splitting produces a signed MSP and zero, one, or three non-nullable u64 lower + // parts, all with the decimal's length and in most-significant-first order. This also + // holds for the constant parts used for empty and all-null inputs. The decimal dtype is + // preserved. + Ok(unsafe { Self::new_unchecked(parts.msp, parts.lower_parts, decimal.decimal_dtype()) }) + } } /// A decimal array decomposed into byte parts. 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 2dfe2a55b3c..af270bad407 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -11,37 +11,24 @@ use vortex_array::dtype::i256; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; -use vortex_error::VortexResult; +use super::DecimalByteParts; use super::DecimalBytePartsArray; -use super::dbp_encode; - -/// Encode a canonical decimal array as byte parts, splitting wide values into lower parts. -pub(crate) fn encode(decimal: &DecimalArray) -> VortexResult { - dbp_encode(decimal, &mut array_session().create_execution_ctx()) -} /// 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, - )) + DecimalByteParts::encode( + &DecimalArray::new(Buffer::from(values), DecimalDType::new(38, 2), validity), + &mut array_session().create_execution_ctx(), + ) .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, - )) + DecimalByteParts::encode( + &DecimalArray::new(Buffer::from(values), DecimalDType::new(76, 2), validity), + &mut array_session().create_execution_ctx(), + ) .vortex_expect("valid decimal byte parts") } - -/// Build an `i256` from a signed high `i128` and unsigned low `u128`. -pub(super) fn i256_of(high: i128, low: u128) -> i256 { - i256::from_parts(low, high) -} 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/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 5d3f610afc7..fa6dca51bb4 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -250,7 +250,7 @@ fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { optimize root=vortex.binary(bool, len=4096) session=false reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=4096) child=vortex.dict(i16, len=4096) -> vortex.dict(bool, len=4096) done output=vortex.dict(bool, len=4096) - child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.dict(bool, len=4096) + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.decimal_byte_parts.v2(decimal(15,2), len=4096) -> vortex.dict(bool, len=4096) iter 1 current=vortex.dict(bool, len=4096) builder_active=false ExecuteSlot slot=0 parent=vortex.dict(bool, len=4096) child=fastlanes.bitpacked(u8, len=4096) iter 2 current=fastlanes.bitpacked(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false @@ -403,8 +403,8 @@ 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: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) + reduce_parent static:FilterReduceAdaptor(DecimalByteParts) slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts.v2(decimal(15,2), len=4096) -> vortex.decimal_byte_parts.v2(decimal(15,2), len=43) + done output=vortex.decimal_byte_parts.v2(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 reduce_parent static:FoRFilterPushDownRule slot=0 parent=vortex.filter(i32, len=43) child=fastlanes.for(i32, len=4096) -> fastlanes.for(i32, len=43) @@ -440,8 +440,8 @@ 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) + reduce_parent static:TakeReduceAdaptor(DecimalByteParts) slot=1 parent=vortex.dict(decimal(15,2), len=64) child=vortex.decimal_byte_parts.v2(decimal(15,2), len=4096) -> vortex.decimal_byte_parts.v2(decimal(15,2), len=64) + done output=vortex.decimal_byte_parts.v2(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-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap index 780ef30a6b0..fad1b4d2b4e 100644 --- a/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: decimal(12,2), len=16384, nbytes=131072 -root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=47666 +root: vortex.decimal_byte_parts.v2(decimal(12,2), len=16384) nbytes=47666 metadata: msp: vortex.pco(i32, len=16384) nbytes=47666 metadata: ptype: i32, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap index f669755e4b1..6eb4ccfed8d 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: decimal(12,2), len=16384, nbytes=131072 -root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=49152 +root: vortex.decimal_byte_parts.v2(decimal(12,2), len=16384) nbytes=49152 metadata: msp: fastlanes.bitpacked(i32, len=16384) nbytes=49152 metadata: bit_width: 24, offset: 0 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..166d84f5ac6 --- /dev/null +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `DecimalByteParts` fixture for wide decimal values that need lower parts. + +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..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,6 +12,7 @@ mod bytebool; mod constant; mod datetimeparts; mod decimal_byte_parts; +mod decimal_byte_parts_v2; mod delta; mod dict; mod for_; @@ -38,6 +39,7 @@ pub fn fixtures() -> Vec> { 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), From 2c44ac0ce35bb29b84e43fb341944b59ebd75ca1 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 16 Sep 2026 12:02:35 -0400 Subject: [PATCH 4/5] Declare the frozen v1 wire ID from the decimal compression scheme Signed-off-by: Matt Katz --- vortex-btrblocks/src/schemes/decimal.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 1dff2171f60..f77a77d8c50 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -8,7 +8,6 @@ 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::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; @@ -16,6 +15,7 @@ use vortex_array::dtype::DecimalType; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -41,7 +41,9 @@ impl Scheme for DecimalScheme { } fn produced_encodings(&self) -> Vec { - vec![DecimalByteParts.id()] + // This scheme only builds single-part arrays, which serialize under the frozen v1 ID. + // The in-memory ID is the v2 wire ID, which no edition permits yet. + vec![decimal_byte_parts_v1_id()] } /// Children: primitive=0. From 026bb26441b2f5af6254795dc032b2cfddf3200b Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 16 Sep 2026 13:51:18 -0400 Subject: [PATCH 5/5] Do not link the private lower-part limit from public docs Signed-off-by: Matt Katz --- encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs | 4 ++-- encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs index 7059fd09838..c1826ae46d4 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -82,7 +82,7 @@ impl DecimalBytePartsData { /// # 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`] + /// 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 unsigned integer array of length `len`. pub fn validate<'a>( msp: &ArrayRef, @@ -169,7 +169,7 @@ impl DecimalByteParts { /// # Safety /// /// The MSP must have a signed integer dtype (`i8`, `i16`, `i32`, or `i64`). There must be - /// at most [`MAX_LOWER_PARTS`] lower parts, each a non-nullable unsigned integer array with the same + /// at most `MAX_LOWER_PARTS` lower parts, each a non-nullable unsigned integer array with the same /// length as the MSP. Lower parts are ordered most significant first. pub(super) unsafe fn new_unchecked( msp: ArrayRef, diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs index 8d8818924e9..f33101e4a3a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -65,7 +65,7 @@ impl DecimalParts { } /// Construct decimal parts arrays from the buffers constituting a wide decimal (`i128` or `i256`). - /// Wide decimals have an `i64` MSP and up to [`super::MAX_LOWER_PARTS`] `u64` lower parts. + /// Wide decimals have an `i64` MSP and up to `MAX_LOWER_PARTS` `u64` lower parts. fn from_wide( msp: Buffer, lower_parts: impl IntoIterator>,