diff --git a/Cargo.lock b/Cargo.lock index f7d3103c9d5..04baecf6504 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10921,8 +10921,10 @@ dependencies = [ name = "vortex-decimal-byte-parts" version = "0.1.0" dependencies = [ + "codspeed-divan-compat", "num-traits", "prost 0.14.4", + "rand 0.10.2", "rstest", "vortex-array", "vortex-buffer", diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 4934ec4fa27..9f2e387a4da 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -26,5 +26,15 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } + +[[bench]] +name = "dbp_assemble" +harness = false + +[[bench]] +name = "dbp_split" +harness = false diff --git a/encodings/decimal-byte-parts/benches/common/mod.rs b/encodings/decimal-byte-parts/benches/common/mod.rs new file mode 100644 index 00000000000..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 d5b0024f5b7..a7f63bfa082 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(