diff --git a/encodings/decimal-byte-parts/benches/dbp_assemble.rs b/encodings/decimal-byte-parts/benches/dbp_assemble.rs index 582d4bb52dc..9a263826e0e 100644 --- a/encodings/decimal-byte-parts/benches/dbp_assemble.rs +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -2,7 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Reassembling integer slices, including output-buffer allocation. -//! Array benchmarks also include validation, part execution, and casts. +//! +//! Array benchmarks also include validation and part execution. mod common; @@ -146,15 +147,8 @@ mod arrays { 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(), - ], - ), + DecimalType::I128 => (38, vec![buffer![u64::MAX; len].into_array()]), + DecimalType::I256 => (76, vec![buffer![u64::MAX; len].into_array(); 3]), _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), }; bench_parts( diff --git a/encodings/decimal-byte-parts/benches/dbp_compute.rs b/encodings/decimal-byte-parts/benches/dbp_compute.rs index 2bbe20fed68..7e431aa5677 100644 --- a/encodings/decimal-byte-parts/benches/dbp_compute.rs +++ b/encodings/decimal-byte-parts/benches/dbp_compute.rs @@ -3,11 +3,12 @@ //! Decimal compute on DecimalByteParts arrays whose MSP holds the whole value. //! -//! Every case has no lower parts. Most cases store the values in an MSP narrower than the -//! canonical width of the dtype, and `msp_i64_p18_canonical` is the control whose MSP has the -//! canonical width. Sum and min-max read the decoded decimal, and the filter filters the MSP and -//! then decodes the result, so their cost follows the width that decoding produces. Comparisons -//! against a constant run on the MSP. +//! Every case has no lower parts. Most cases store the values in an MSP narrower than the canonical +//! width of the dtype, and `msp_i64_p18_canonical` is the control whose MSP has the canonical +//! width. `decode` and `filter` stop at the form that decoding produces, a [`Widen`] or the +//! canonical array, and sum and min-max read that form, so their cost follows its width. +//! `canonicalize` also pays the widening copy that an export pays, so it omits the control, where +//! decoding already returns the canonical array. Comparisons against a constant run on the MSP. use std::fmt; use std::sync::LazyLock; @@ -20,7 +21,6 @@ use rand::distr::uniform::SampleRange; use rand::distr::uniform::SampleUniform; use rand::rngs::StdRng; use vortex_array::ArrayRef; -use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -31,6 +31,7 @@ use vortex_array::arrays::BoolArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Widen; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::NativePType; @@ -253,11 +254,25 @@ fn bench_case( fn decode(bencher: Bencher, case: Case) { bench_case(bencher, case, |array, ctx| { array - .execute::(ctx) + .execute_until::(ctx) .vortex_expect("decode DecimalByteParts") }); } +#[divan::bench(args = [ + Case::MspI32P18, + Case::MspI32P18Nullable, + Case::GapMspI16P10, + Case::MspI64P38, +])] +fn canonicalize(bencher: Bencher, case: Case) { + bench_case(bencher, case, |array, ctx| { + array + .execute::(ctx) + .vortex_expect("canonicalize DecimalByteParts") + }); +} + /// Compares a new array of `case` with `Gte` against a constant of `value` in each iteration. fn bench_compare(bencher: Bencher, case: Case, value: DecimalValue) { let constant = case.constant(value); @@ -302,7 +317,7 @@ fn filter(bencher: Bencher, case: Case) { array .filter(mask) .vortex_expect("filter DecimalByteParts") - .execute::(&mut ctx) + .execute_until::(&mut ctx) .vortex_expect("execute the filter") }); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs deleted file mode 100644 index 796789d5799..00000000000 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ /dev/null @@ -1,766 +0,0 @@ -// 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 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::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::arrays::Primitive; -use vortex_array::buffer::BufferHandle; -use vortex_array::dtype::DType; -use vortex_array::dtype::DecimalDType; -use vortex_array::require_child; -use vortex_array::scalar::DecimalValue; -use vortex_array::scalar::Scalar; -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_ensure_eq; -use vortex_error::vortex_panic; -use vortex_session::VortexSession; - -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; - -/// 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_eq!(dtype, &expected_dtype); - vortex_ensure_eq!(msp.len(), 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_eq!(part.len(), len, "lower part {idx} length mismatch"); - } - 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 - /// 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, - 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, - slots, - )) - } - } -} - -impl VTable for DecimalByteParts { - type TypedArrayData = DecimalBytePartsData; - - type OperationsVTable = Self; - type ValidityVTable = ValidityVTableFromChild; - - fn id(&self) -> ArrayId { - decimal_byte_parts_v2_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>> { - Ok(None) - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("DecimalByteParts deserialization requires DecimalBytePartsPlugin") - } - - 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 mut array = require_child!(array, array.msp(), DecimalBytePartsSlots::MSP => Primitive); - for idx in 0..array.lower_parts().len() { - array = require_child!( - array, - array.lower_parts()[idx], - DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx => Primitive - ); - } - 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_decimal(value, array.decimal_dtype(), array.dtype().nullability()) - } -} - -impl ValidityChild for DecimalByteParts { - fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { - // validity stored in 0th child - array.msp().clone() - } -} - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - 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::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_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use super::DecimalByteParts; - 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_parts; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[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 SESSION.create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new( - dtype.clone(), - Some(ScalarValue::Decimal(DecimalValue::I64(200))) - ) - .unwrap(), - array - .execute_scalar(1, &mut SESSION.create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), - array - .execute_scalar(2, &mut SESSION.create_execution_ctx()) - .unwrap() - ); - } - - #[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 = 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::from_parts((7u128 << 64) | 8, (5i128 << 64) | 6)], - Validity::NonNullable, - ); - assert_eq!(array.lower_parts().len(), MAX_LOWER_PARTS); - - let mut ctx = 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] - 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 = 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 = 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 = 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, - ) - } - }; - assert_arrays_eq!(encoded, expected, &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, - slots.into_iter().collect(), - ); - assert!(Array::try_from_parts(parts).is_err()); - } - - #[test] - fn test_wide_decimal_buffer_types() -> VortexResult<()> { - let mut ctx = 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::from_parts(0, 1 << 100)], 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::from_parts(9, 1)] - ); - 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 = SESSION.create_execution_ctx(); - assert_eq!( - array.execute_scalar(0, &mut ctx)?, - Scalar::null(array.dtype().clone()) - ); - assert_arrays_eq!( - array, - DecimalArray::new( - buffer![0i128, 0, 11], - DecimalDType::new(38, 2), - Validity::from_iter([false, false, true]), - ), - &mut ctx - ); - Ok(()) - } -} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array/mod.rs new file mode 100644 index 00000000000..91b6a1de84d --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array/mod.rs @@ -0,0 +1,412 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The DecimalByteParts encoding and its validated part layout. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +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::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::arrays::Primitive; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::require_child; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +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_ensure_eq; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; + +use super::LOWER_PART_DTYPE; +use super::MAX_I128_LOWER_PARTS; +use super::MAX_I256_LOWER_PARTS; +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; + +/// 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 a signed integer +/// of at most 64 bits, and it is nullable iff the decimal is nullable. Only the MSP may use a +/// narrower dtype when its values fit. +/// +/// Every lower part is a non-nullable `u64` array holding a 64-bit window of the value. The lower +/// parts never assemble wider than the canonical width of the decimal dtype: precision up to 18 +/// allows no lower parts, precision 19 to 38 allows at most one, and precision 39 to 76 allows +/// up to three. +/// +/// Executing the array decodes it into the canonical decimal array or a decimal [`Widen`], as +/// [`assemble_decimal`] describes. +/// +/// 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. +/// +/// [`Widen`]: vortex_array::arrays::Widen +#[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 the lower parts assemble wider than the canonical + /// width of `decimal_dtype`, or if any lower part is not a non-nullable `u64` array of length + /// `len`. + pub fn validate<'a>( + msp: &ArrayRef, + lower_parts: impl ExactSizeIterator, + decimal_dtype: DecimalDType, + dtype: &DType, + len: usize, + ) -> 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_eq!(dtype, &expected_dtype); + vortex_ensure_eq!(msp.len(), len); + + let lower_part_count = lower_parts.len(); + let max_lower_parts = max_lower_parts(decimal_dtype); + vortex_ensure!( + lower_part_count <= max_lower_parts, + "{decimal_dtype} allows at most {max_lower_parts} lower parts, got {lower_part_count}" + ); + + for (idx, part) in lower_parts.enumerate() { + vortex_ensure!( + part.dtype() == &LOWER_PART_DTYPE, + "lower part {idx} must have dtype {LOWER_PART_DTYPE}, got {}", + part.dtype() + ); + vortex_ensure_eq!(part.len(), len, "lower part {idx} length mismatch"); + } + + Ok(()) + } +} + +/// The largest number of lower parts that assemble to at most the canonical width of +/// `decimal_dtype`. +/// +/// One lower part assembles to `i128`, and two or three assemble to `i256`. An MSP is at most +/// `i64`, so a canonical width of `i64` or less allows no lower parts. With this bound, decoding +/// never narrows values that it assembled from lower parts. +fn max_lower_parts(decimal_dtype: DecimalDType) -> usize { + match decimal_dtype.canonical_values_type() { + DecimalType::I8 | DecimalType::I16 | DecimalType::I32 | DecimalType::I64 => 0, + DecimalType::I128 => MAX_I128_LOWER_PARTS, + DecimalType::I256 => MAX_I256_LOWER_PARTS, + } +} + +/// The encoding vtable for decimal byte parts. +#[derive(Clone, Debug)] +pub struct DecimalByteParts; + +impl DecimalByteParts { + /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array. + pub fn try_new( + msp: ArrayRef, + decimal_dtype: DecimalDType, + ) -> VortexResult { + Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + } + + /// Construct a new [`DecimalBytePartsArray`] from an MSP array, its lower parts, and a + /// decimal dtype. + /// + /// Lower parts are ordered most significant first and must each be a non-nullable `u64` array + /// of the same length as the MSP. See [`super::split_decimal`] for producing them from a + /// 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, + slots, + )) + } +} + +impl VTable for DecimalByteParts { + type TypedArrayData = DecimalBytePartsData; + + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + decimal_byte_parts_v2_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>> { + Ok(None) + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("DecimalByteParts deserialization requires DecimalBytePartsPlugin") + } + + 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 mut array = require_child!(array, array.msp(), DecimalBytePartsSlots::MSP => Primitive); + for idx in 0..array.lower_parts().len() { + array = require_child!( + array, + array.lower_parts()[idx], + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx => Primitive + ); + } + 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"); + + 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_decimal(value, array.decimal_dtype(), array.dtype().nullability()) + } +} + +impl ValidityChild for DecimalByteParts { + fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { + // validity stored in 0th child + array.msp().clone() + } +} + +#[cfg(test)] +mod tests; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array/tests.rs new file mode 100644 index 00000000000..af2c81f0ccc --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array/tests.rs @@ -0,0 +1,286 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests of part validation and execution. + +use std::sync::LazyLock; + +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::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::widen::WidenArraySlotsExt; +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::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use super::DecimalByteParts; +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_parts; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() +} + +fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() +} + +#[test] +fn test_scalar_at_decimal_parts() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(8, 2); + let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); + let array = DecimalByteParts::try_new( + PrimitiveArray::new( + buffer![100i32, 200, 400], + Validity::from_iter([false, true, true]), + ) + .into_array(), + decimal_dtype, + )? + .into_array(); + let expected = [ + Scalar::null(dtype), + Scalar::decimal(DecimalValue::I64(200), decimal_dtype, Nullability::Nullable), + Scalar::decimal(DecimalValue::I64(400), decimal_dtype, Nullability::Nullable), + ]; + + for (row, expected) in expected.into_iter().enumerate() { + assert_eq!(array.execute_scalar(row, &mut ctx)?, expected); + } + + 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 = 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::from_parts((7u128 << 64) | 8, (5i128 << 64) | 6)], + Validity::NonNullable, + ); + assert_eq!(array.lower_parts().len(), MAX_LOWER_PARTS); + + let mut ctx = 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(()) +} + +/// `scalar_at` assembles one row without decoding the array, so it must agree with decoding. +#[rstest] +fn test_scalar_at_matches_decoded_rows( + #[values(PType::I8, PType::I16, PType::I32, PType::I64)] msp_ptype: PType, + #[values(0, 1, 2, 3)] lower_count: usize, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let msp = PrimitiveArray::new( + buffer![0i64, 1, -1], + Validity::from_iter([false, true, true]), + ) + .into_array() + .cast(DType::Primitive(msp_ptype, Nullability::Nullable))?; + // Set the highest bit of each word to catch accidental sign extension. The low bits differ, so + // a swapped word shows up too. + let lower = [(1u64 << 63) | 1, (1 << 63) | 2, (1 << 63) | 3] + .into_iter() + .take(lower_count) + .map(|word| buffer![0u64, word, word].into_array()) + .collect(); + let dtype = DecimalDType::new(if lower_count <= 1 { 38 } else { 76 }, 0); + let array = DecimalByteParts::try_new_with_lower_parts(msp, lower, dtype)?; + + let decoded = array + .clone() + .into_array() + .execute::(&mut ctx)?; + + for row in 0..array.len() { + assert_eq!( + array.execute_scalar(row, &mut ctx)?, + decoded.execute_scalar(row, &mut ctx)? + ); + } + + Ok(()) +} + +#[rstest] +#[case::signed_lower_part( + vec![buffer![1i64, 2, 3].into_array()], + DecimalDType::new(38, 2), + "must have dtype u64, got i64" +)] +#[case::nullable_lower_part( + vec![PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid).into_array()], + DecimalDType::new(38, 2), + "must have dtype u64, got u64?" +)] +#[case::narrow_lower_part( + vec![buffer![1u32, 2, 3].into_array()], + DecimalDType::new(38, 2), + "must have dtype u64, got u32" +)] +#[case::mismatched_length( + vec![buffer![1u64, 2].into_array()], + DecimalDType::new(38, 2), + "length mismatch" +)] +#[case::too_many_parts( + vec![lower_part(), lower_part(), lower_part(), lower_part()], + DecimalDType::new(76, 2), + "expected 1..=4 slots, got 5" +)] +#[case::one_lower_part_at_p18( + vec![lower_part()], + DecimalDType::new(18, 2), + "allows at most 0 lower parts, got 1" +)] +#[case::three_lower_parts_at_p38( + vec![lower_part(), lower_part(), lower_part()], + DecimalDType::new(38, 2), + "allows at most 1 lower parts, got 3" +)] +fn test_rejects_invalid_parts( + #[case] lower_parts: Vec, + #[case] decimal_dtype: DecimalDType, + #[case] expected_error: &str, +) { + let result = DecimalByteParts::try_new_with_lower_parts(msp(), lower_parts, decimal_dtype); + assert!( + result + .as_ref() + .is_err_and(|err| err.to_string().contains(expected_error)), + "expected {expected_error}, got {result:?}" + ); +} + +#[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, + slots.into_iter().collect(), + ); + assert!(Array::try_from_parts(parts).is_err()); +} + +#[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 = SESSION.create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + assert_arrays_eq!( + array, + DecimalArray::new( + buffer![0i128, 0, 11], + DecimalDType::new(38, 2), + Validity::from_iter([false, false, true]), + ), + &mut ctx + ); + + Ok(()) +} + +/// Planted exact statistics select the child without scanning the encoded MSP's values. +/// A scan of the actual values would pick `Decimal(1, 0)` instead. +#[rstest] +fn test_execute_keeps_stats_of_encoded_msp( + #[values(false, true)] single_step: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let msp = ConstantArray::new(5i32, 3).into_array(); + msp.statistics().set(Stat::Min, Precision::exact(-999i32)); + msp.statistics().set(Stat::Max, Precision::exact(999i32)); + let array = DecimalByteParts::try_new(msp, DecimalDType::new(18, 0))?.into_array(); + + let decoded = if single_step { + let mut decoded = array; + // One step executes the MSP, and the next one decodes. + for _ in 0..2 { + decoded = decoded.execute::(&mut ctx)?; + } + decoded + } else { + array.execute_until::(&mut ctx)? + }; + + let widen = decoded + .as_opt::() + .unwrap_or_else(|| panic!("expected a Widen, got {}", decoded.encoding_id())); + assert_eq!( + widen.values().dtype(), + &DType::Decimal(DecimalDType::new(3, 0), Nullability::NonNullable) + ); + + 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 deleted file mode 100644 index 1cdf27c15bd..00000000000 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs +++ /dev/null @@ -1,467 +0,0 @@ -// 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 vortex_error::vortex_ensure_eq; - -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_eq!(part.len(), len, "lower part {idx} length mismatch"); - } - - 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 std::sync::LazyLock; - - use rstest::rstest; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - 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 vortex_session::VortexSession; - - use super::assemble_decimal; - use crate::decimal_byte_parts::split_decimal; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[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 = 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!(round_tripped, decimal, &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 = 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!(lower, expected.clone(), &mut ctx); - } - assert_arrays_eq!(round_trip(decimal.clone())?, decimal, &mut ctx); - Ok(()) - } - - fn round_trip(decimal: DecimalArray) -> VortexResult { - let mut ctx = 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 = 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!(round_trip(decimal.clone())?, 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 = 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 = 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 = 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!( - actual, - DecimalArray::new(expected, dtype, Validity::NonNullable), - &mut ctx - ); - Ok(()) - } -} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble/mod.rs new file mode 100644 index 00000000000..900e116bc22 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble/mod.rs @@ -0,0 +1,225 @@ +// 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::RangeInclusive; +use std::ops::Shl; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min_max::min_max; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::WidenArray; +use vortex_array::arrays::decimal::DecimalDataParts; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::i256; +use vortex_array::expr::stats::Precision; +use vortex_array::match_each_signed_integer_ptype; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::trusted_len::TrustedLen; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; + +use super::DecimalBytePartsData; +use super::LOWER_PART_BITS; +use super::MAX_LOWER_PARTS; + +/// Reassemble decimal parts into the canonical array of `decimal_dtype`, or into a decimal +/// [`Widen`] over a narrower child. +/// +/// The parts are checked with [`DecimalBytePartsData::validate`], which owns the rules for part +/// dtypes, lengths, and counts. +/// +/// The parts first assemble at their _natural width_: the MSP width with no lower parts, `i128` +/// with one lower part, and `i256` with two or three. Because the validation bounds the lower +/// parts, only a lone MSP can be wider than the canonical width. Then +/// [`WidenArray::from_stored_decimal`] picks the result from the range of the valid values. A +/// narrow MSP becomes the child of a `Widen` without a copy when its width is canonical for the +/// selected child precision. +/// +/// A lone MSP that is not at the canonical width takes the range from its [`min_max`], which reuses +/// exact cached statistics. Values assembled from lower parts are scanned only when their natural +/// width is not canonical, which happens for one lower part under a precision above 38. +/// +/// # Errors +/// +/// Returns an error if [`DecimalBytePartsData::validate`] rejects the parts, or if executing a part +/// fails. When the natural width is not the canonical width, it also returns an error if a valid +/// value is outside the precision of `decimal_dtype`. +/// +/// [`Widen`]: vortex_array::arrays::Widen +pub fn assemble_decimal( + msp: &ArrayRef, + lower_parts: &[ArrayRef], + decimal_dtype: DecimalDType, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + DecimalBytePartsData::validate(msp, lower_parts.iter(), decimal_dtype, &dtype, msp.len())?; + + // TODO(connor): Broadcast constant parts directly instead of materializing their buffers. + let msp = msp.clone().execute::(exec_ctx)?; + let validity = msp.validity()?; + + if lower_parts.is_empty() { + let values_type = DecimalType::try_from(msp.ptype())?; + + // At the canonical width, `from_stored_decimal` reuses the buffer without reading the + // values, so a range pass would be wasted. A `None` range means that no value is valid, and + // the scan in `from_stored_decimal` then finds no valid value either. + let known_range = if values_type == decimal_dtype.canonical_values_type() { + Precision::Absent + } else { + msp_range(&msp, exec_ctx)?.map_or(Precision::Absent, Precision::Exact) + }; + + let parts = DecimalDataParts { + decimal_dtype, + values: msp.buffer_handle().clone(), + values_type, + validity, + }; + + return WidenArray::from_stored_decimal(parts, known_range, exec_ctx); + } + + let lower_parts = lower_parts + .iter() + .map(|part| part.clone().execute::(exec_ctx)) + .collect::>>()?; + let (values, values_type) = assemble_natural_width(&msp, &lower_parts); + + let parts = DecimalDataParts { + decimal_dtype, + values, + values_type, + validity, + }; + + WidenArray::from_stored_decimal(parts, Precision::Absent, exec_ctx) +} + +/// Returns the range of the valid values of `msp`, or `None` when no value is valid. +/// +/// [`min_max`] returns exact cached statistics without reading the values. Otherwise it makes one +/// vectorized pass over them and caches the result on `msp`. The decimal [`valid_range`] scan is +/// not vectorized, so it is not used here. +/// +/// [`valid_range`]: vortex_array::aggregate_fn::fns::min_max::valid_range +fn msp_range( + msp: &PrimitiveArray, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult>> { + let Some(extrema) = min_max(msp.as_ref(), exec_ctx, NumericalAggregateOpts::default())? else { + return Ok(None); + }; + + let min = i64::try_from(&extrema.min)?; + let max = i64::try_from(&extrema.max)?; + + Ok(Some(i256::from(min)..=i256::from(max))) +} + +/// Assemble the MSP and one to three `u64` lower parts at their natural width: `i128` for one +/// lower part, and `i256` for two or three. +fn assemble_natural_width( + msp: &PrimitiveArray, + lower_parts: &[PrimitiveArray], +) -> (BufferHandle, DecimalType) { + let (values, values_type) = match_each_signed_integer_ptype!(msp.ptype(), |Msp| { + let msp = msp.as_slice::(); + match lower_parts { + [first] => ( + assemble_wide_decimal::( + msp, + first.as_slice::().iter().map(|&word| [word]), + ) + .into_byte_buffer(), + DecimalType::I128, + ), + [first, second] => ( + assemble_wide_decimal::( + msp, + first + .as_slice::() + .iter() + .zip(second.as_slice::()) + .map(|(&a, &b)| [a, b]), + ) + .into_byte_buffer(), + DecimalType::I256, + ), + [first, second, third] => ( + assemble_wide_decimal::( + msp, + first + .as_slice::() + .iter() + .zip(second.as_slice::()) + .zip(third.as_slice::()) + .map(|((&a, &b), &c)| [a, b, c]), + ) + .into_byte_buffer(), + DecimalType::I256, + ), + _ => vortex_panic!( + "DecimalBytePartsData::validate bounds the lower parts to at most \ + {MAX_LOWER_PARTS}, and assemble_decimal passes at least one, got {}", + lower_parts.len() + ), + } + }); + + (BufferHandle::new_host(values), values_type) +} + +/// 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; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble/tests.rs new file mode 100644 index 00000000000..9597cb6559f --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble/tests.rs @@ -0,0 +1,473 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Tests of assembly widths, buffer reuse, and cached extrema. + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::Decimal; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::widen::WidenArraySlotsExt; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::BigCast; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::expr::stats::Precision; +use vortex_array::expr::stats::Stat; +use vortex_array::expr::stats::StatsProvider; +use vortex_array::match_each_decimal_value_type; +use vortex_array::match_each_signed_integer_ptype; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use super::assemble_decimal; +use crate::decimal_byte_parts::split_decimal; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +/// Decode an MSP with no lower parts, without executing the result. +fn decode_msp(msp: PrimitiveArray, decimal_dtype: DecimalDType) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + assemble_decimal(&msp.into_array(), &[], decimal_dtype, &mut ctx) +} + +/// Returns the rows of a fixture at `nullability`. A `None` is a null row, and a +/// non-nullable fixture leaves the `None` rows out. +fn rows_at(values: &[Option], nullability: Nullability) -> Vec> { + match nullability { + Nullability::NonNullable => values.iter().flatten().map(|&value| Some(value)).collect(), + Nullability::Nullable => values.to_vec(), + } +} + +/// Returns the validity of `rows` at `nullability`. +fn validity_of(rows: &[Option], nullability: Nullability) -> Validity { + match nullability { + Nullability::NonNullable => Validity::NonNullable, + Nullability::Nullable => Validity::from_iter(rows.iter().map(Option::is_some)), + } +} + +/// Returns an MSP of `ptype` with the given rows. Each null row holds the largest value of +/// `ptype`, so a range that counted it would pick a wider child. +fn msp_of(ptype: PType, rows: &[Option], validity: Validity) -> PrimitiveArray { + match_each_signed_integer_ptype!(ptype, |P| { + let values = rows + .iter() + .map(|row| match row { + Some(value) => num_traits::cast::(*value) + .vortex_expect("fixture values fit the MSP ptype"), + None => P::MAX, + }) + .collect::>(); + PrimitiveArray::new(values, validity) + }) +} + +/// Returns `rows` as the canonical array of `decimal_dtype`, with zero in each null slot. +fn canonical_decimal( + rows: &[Option], + decimal_dtype: DecimalDType, + validity: Validity, +) -> DecimalArray { + match_each_decimal_value_type!(decimal_dtype.canonical_values_type(), |T| { + let values = rows + .iter() + .map(|row| { + ::from(row.unwrap_or_default()) + .vortex_expect("fixture values fit the canonical width") + }) + .collect::>(); + DecimalArray::new(values, decimal_dtype, validity) + }) +} + +/// Returns the decimal array that holds the values of `decoded`: the child of a `Widen`, or +/// `decoded` itself. +fn decimal_leaf(decoded: &ArrayRef) -> DecimalArray { + let leaf = match decoded.as_opt::() { + Some(widen) => widen.values().clone(), + None => decoded.clone(), + }; + leaf.downcast::() +} + +/// Asserts that `decoded` is a `Widen` over a child at `child_precision`, or the canonical +/// array of `decimal_dtype` when `child_precision` is `None`. Either way, the decimal leaf +/// has the canonical width of its own dtype. +#[track_caller] +fn assert_form(decoded: &ArrayRef, decimal_dtype: DecimalDType, child_precision: Option) { + let leaf_dtype = match child_precision { + Some(precision) => DecimalDType::new(precision, decimal_dtype.scale()), + None => decimal_dtype, + }; + let leaf = decimal_leaf(decoded); + + assert_eq!(decoded.is::(), child_precision.is_some()); + assert_eq!(leaf.decimal_dtype(), leaf_dtype); + assert_eq!(leaf.values_type(), leaf_dtype.canonical_values_type()); +} + +fn round_trip(decimal: DecimalArray) -> VortexResult { + let mut ctx = 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::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 = 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!(round_tripped, decimal, &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 = 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!(lower, expected.clone(), &mut ctx); + } + assert_arrays_eq!(round_trip(decimal.clone())?, decimal, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::i128_zero(0i128)] +#[case::i128_one(1i128)] +#[case::i128_minus_one(-1i128)] +#[case::i128_limb_boundary(1i128 << 64)] +#[case::i128_just_below_limb_boundary((1i128 << 64) - 1)] +#[case::i128_negative_limb_boundary(-(1i128 << 64))] +#[case::i128_max(i128::MAX)] +#[case::i128_min(i128::MIN)] +#[case::i256_zero(i256::ZERO)] +#[case::i256_one(i256::ONE)] +#[case::i256_minus_one(i256::ZERO - i256::ONE)] +#[case::i256_max(i256::MAX)] +#[case::i256_min(i256::MIN)] +#[case::i256_word_1(i256::from_parts(1u128 << 64, 0))] +#[case::i256_word_2(i256::from_parts(0, 1))] +#[case::i256_word_3(i256::from_parts(0, 1i128 << 64))] +#[case::i256_mixed(i256::from_parts(u128::MAX, -3))] +fn test_split_assemble(#[case] value: T) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(T::MAX_PRECISION, 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 = SESSION.create_execution_ctx(); + let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(9, 0), validity); + let values_ptr = decimal.buffer::().as_ptr(); + + 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.clone().execute::(&mut ctx)?; + assert_eq!(msp.as_slice::().as_ptr(), values_ptr); + + let decoded = assemble_decimal( + &parts.msp, + &parts.lower_parts, + decimal.decimal_dtype(), + &mut ctx, + )? + .execute::(&mut ctx)?; + assert_eq!(decoded.buffer::().as_ptr(), values_ptr); + assert_arrays_eq!(decoded, decimal, &mut ctx); + + Ok(()) +} + +#[test] +fn test_assemble_validates_parts() { + let msp = buffer![0i64; 2].into_array(); + let lower = [buffer![0u32; 2].into_array()]; + + let result = assemble_decimal( + &msp, + &lower, + DecimalDType::new(38, 0), + &mut SESSION.create_execution_ctx(), + ); + + assert!( + result + .as_ref() + .is_err_and(|err| err.to_string().contains("must have dtype u64")), + "expected the lower-part dtype check to reject a u32 part, got {result:?}" + ); +} + +/// The tests of `WidenArray::from_stored_decimal` cover how a range picks the child, so these +/// cases cover how the MSP supplies that range. +#[rstest] +#[case::canonical_trusts_precision( + PType::I32, + vec![Some(2_000_000_000), None], + DecimalDType::new(9, 2), + None, + true, +)] +#[case::narrow_reuses( + PType::I32, + vec![Some(-999_999_999), None, Some(999_999_999)], + DecimalDType::new(18, 2), + Some(9), + true, +)] +#[case::negative_min_counts( + PType::I8, + vec![Some(-128), None, Some(5)], + DecimalDType::new(10, 0), + Some(3), + false, +)] +#[case::wider_to_widen( + PType::I64, + vec![Some(-99), None, Some(99)], + DecimalDType::new(9, 2), + Some(2), + false, +)] +#[case::all_null(PType::I32, vec![None, None], DecimalDType::new(18, 2), Some(2), false)] +#[case::empty(PType::I32, vec![], DecimalDType::new(18, 2), Some(2), false)] +fn test_assemble_result_form( + #[case] msp_ptype: PType, + #[case] values: Vec>, + #[case] decimal_dtype: DecimalDType, + #[case] child_precision: Option, + #[case] reuses_msp: bool, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let rows = rows_at(&values, nullability); + let validity = validity_of(&rows, nullability); + let msp = msp_of(msp_ptype, &rows, validity.clone()); + let msp_ptr = msp.buffer_handle().as_host().as_ptr(); + + let decoded = decode_msp(msp, decimal_dtype)?; + + assert_form(&decoded, decimal_dtype, child_precision); + if reuses_msp { + let leaf = decimal_leaf(&decoded); + assert_eq!(leaf.buffer_handle().as_host().as_ptr(), msp_ptr); + } + assert_arrays_eq!( + decoded, + canonical_decimal(&rows, decimal_dtype, validity), + &mut ctx + ); + + Ok(()) +} + +#[test] +fn test_assemble_narrowing_copy_rejects_value_outside_precision() { + // 100 fits `i8`, but not precision 2. + let msp = PrimitiveArray::new(buffer![1i64, 100], Validity::NonNullable); + + let result = decode_msp(msp, DecimalDType::new(2, 0)); + + assert!( + result + .as_ref() + .is_err_and(|err| err.to_string().contains("Decimal range must fit")), + "expected the range of the MSP to reject 100, got {result:?}" + ); +} + +#[rstest] +#[case::one_lower_part_at_p38(DecimalArray::new( + buffer![(3i128 << 64) | 2, (-3i128 << 64) | 2], + DecimalDType::new(38, 0), + Validity::NonNullable, +), 1, None)] +#[case::one_lower_part_at_p40(DecimalArray::new( + buffer![i256::from_i128((3i128 << 64) | 2), i256::from_i128((-3i128 << 64) | 2)], + DecimalDType::new(40, 0), + Validity::NonNullable, +), 1, Some(20))] +#[case::two_lower_parts_at_p76(DecimalArray::new( + buffer![ + i256::from_parts((1u128 << 64) | 2, 3), + i256::from_parts((1u128 << 64) | 2, -3), + ], + DecimalDType::new(76, 0), + Validity::NonNullable, +), 2, None)] +#[case::three_lower_parts_at_p76(DecimalArray::new( + buffer![ + i256::from_parts((1u128 << 64) | 2, (3i128 << 64) | 4), + i256::from_parts((1u128 << 64) | 2, (-3i128 << 64) | 4), + ], + DecimalDType::new(76, 0), + Validity::NonNullable, +), 3, None)] +fn test_assemble_lower_parts_result_form( + #[case] expected: DecimalArray, + #[case] lower_count: usize, + #[case] child_precision: Option, + #[values(false, true)] narrow_msp: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = expected.decimal_dtype(); + let msp = if narrow_msp { + buffer![3i8, -3].into_array() + } else { + buffer![3i64, -3].into_array() + }; + let lower = [4u64, 1, 2].map(|word| buffer![word; 2].into_array()); + + let decoded = assemble_decimal(&msp, &lower[3 - lower_count..], decimal_dtype, &mut ctx)?; + + assert_form(&decoded, decimal_dtype, child_precision); + assert_arrays_eq!(decoded, expected, &mut ctx); + + Ok(()) +} + +/// The range from `min_max` on the MSP is exact, so the child records it as exact extrema, and +/// a later `min_max` on the decoded array reads them without a scan. +#[test] +fn test_assemble_records_exact_child_extrema() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(18, 0); + let msp = PrimitiveArray::new(buffer![-5i32, 7], Validity::NonNullable); + + let decoded = decode_msp(msp, decimal_dtype)?; + + assert_form(&decoded, decimal_dtype, Some(1)); + let child_dtype = DecimalDType::new(1, 0); + let extremum = |value: i8| { + Scalar::decimal( + DecimalValue::I8(value), + child_dtype, + Nullability::NonNullable, + ) + }; + let leaf = decimal_leaf(&decoded); + let statistics = leaf.statistics(); + assert_eq!(statistics.get(Stat::Min), Precision::Exact(extremum(-5))); + assert_eq!(statistics.get(Stat::Max), Precision::Exact(extremum(7))); + + Ok(()) +} + +#[test] +fn test_assemble_ignores_inexact_stats() -> VortexResult<()> { + let decimal_dtype = DecimalDType::new(18, 0); + let msp = PrimitiveArray::new(buffer![-5i32, 7], Validity::NonNullable); + msp.statistics().set(Stat::Min, Precision::inexact(-999i32)); + msp.statistics().set(Stat::Max, Precision::inexact(999i32)); + + let decoded = decode_msp(msp, decimal_dtype)?; + + assert_form(&decoded, decimal_dtype, Some(1)); + + Ok(()) +} 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 949b3d048d5..2e382e48b51 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 @@ -612,15 +612,14 @@ mod tests { ) } - /// Narrow parts cannot hold every constant of their precision: a constant above the parts' - /// range answers by its sign, and a word above a lower part's range sits above every value - /// in that part. + /// A constant beyond the narrow MSP's range answers by its sign. Constants with matching + /// high words still compare their lower words. #[test] - fn narrow_parts_match_decimal() -> VortexResult<()> { + fn narrow_msp_matches_decimal() -> VortexResult<()> { let dtype = DecimalDType::new(38, 2); let validity = Validity::from_iter([true, false, true, true]); let msp = PrimitiveArray::new(buffer![1i16, 0, -1, 3], validity.clone()).into_array(); - let lower = buffer![5u16, 7, 9, u16::MAX].into_array(); + let lower = buffer![5u64, 7, 9, u64::from(u16::MAX)].into_array(); let array = DecimalByteParts::try_new_with_lower_parts(msp, vec![lower], dtype)?; let values: Vec = vec![ @@ -650,15 +649,14 @@ mod tests { ) } - /// Middle lower parts can also be narrow. A word above a middle part's range decides the - /// answer before the less significant parts. + /// A differing middle word decides the answer before the less significant parts. #[test] - fn narrow_middle_parts_match_decimal() -> VortexResult<()> { + fn middle_parts_match_decimal() -> VortexResult<()> { let dtype = DecimalDType::new(76, 2); let validity = Validity::from_iter([true, false, true, true]); let msp = PrimitiveArray::new(buffer![1i8, 0, -1, 3], validity.clone()).into_array(); - let upper = buffer![5u8, 7, u8::MAX, 0].into_array(); - let middle = buffer![9u16, u16::MAX, 0, 3].into_array(); + let upper = buffer![5u64, 7, u64::from(u8::MAX), 0].into_array(); + let middle = buffer![9u64, u64::from(u16::MAX), 0, 3].into_array(); let lower = buffer![2u64, u64::MAX, 4, 0].into_array(); let array = DecimalByteParts::try_new_with_lower_parts(msp, vec![upper, middle, lower], dtype)?; 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 65bb8222f7a..5df1b622c53 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -12,8 +12,10 @@ //! ``` //! //! This is exactly the two's complement bit pattern of the decimal value cut on 64-bit -//! boundaries. -//! Each part may use a narrower integer dtype when its values fit; the word boundaries stay fixed. +//! boundaries. [`DecimalBytePartsData`] owns the rules for the part dtypes and count, and +//! [`assemble_decimal`] owns the rules that pick the decoded form. +//! +//! [`assemble_decimal`]: assemble::assemble_decimal use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -59,5 +61,5 @@ const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS; /// Number of bits stored in each lower part. const LOWER_PART_BITS: usize = 64; -/// Dtype of the lower parts produced by splitting, before any narrowing. +/// The dtype of every lower part. const LOWER_PART_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable); 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 index f74c3bf1c55..69ac1bf6052 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs @@ -18,7 +18,6 @@ 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; @@ -47,39 +46,28 @@ static SESSION: LazyLock = LazyLock::new(VortexSession::empty); &DecimalArray::new(buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable), &mut SESSION.create_execution_ctx(), ))] -#[case::wider_i128_storage(DecimalByteParts::encode( - &DecimalArray::new(buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable), - &mut 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 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), + DecimalDType::new(76, 2), ))] -#[case::narrowed_lower_parts(DecimalByteParts::try_new_with_lower_parts( +#[case::narrow_msp_with_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(), + buffer![u64::from(u8::MAX), 128, 0].into_array(), + ConstantArray::new(u64::from(u16::MAX), 3).into_array(), + buffer![0u64, 1 << 31, u64::from(u32::MAX)].into_array(), ], DecimalDType::new(76, 2), ))] -#[case::nullable_mixed_lower_parts(DecimalByteParts::try_new_with_lower_parts( +#[case::nullable_narrow_msp_with_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(), + buffer![0u64, 128, u64::from(u8::MAX)].into_array(), + buffer![u64::from(u32::MAX), 1 << 31, 0].into_array(), ], DecimalDType::new(76, 2), ))] @@ -183,7 +171,13 @@ fn v1_metadata_is_unchanged() -> VortexResult<()> { decimal_byte_parts_v2_id(), v2_metadata(vec![PType::I64 as i32]), vec![msp(), lower_part()], - "unsigned integer dtype" + "must have dtype u64, got i64" +)] +#[case::v2_narrow_lower_part( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::U32 as i32]), + vec![msp(), lower_part()], + "must have dtype u64, got u32" )] #[case::v2_unknown_ptype( decimal_byte_parts_v2_id(), 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 index ca4ae4d6138..3a83b94a691 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs @@ -31,7 +31,7 @@ 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. + /// Ptypes of the lower parts, ordered most significant first. Readers accept only `U64`. #[prost(enumeration = "PType", repeated, tag = "2")] pub(super) lower_part_ptypes: Vec, } @@ -93,8 +93,8 @@ pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult 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()) }) + Self::try_new_with_lower_parts(parts.msp, parts.lower_parts, decimal.decimal_dtype()) } } @@ -254,6 +252,7 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; + use super::DecimalByteParts; use super::split_decimal; use crate::decimal_byte_parts::LOWER_PART_DTYPE; use crate::decimal_byte_parts::MAX_LOWER_PARTS; @@ -312,4 +311,23 @@ mod tests { } Ok(()) } + + /// Storage wider than the canonical width splits into more lower parts than the dtype allows. + #[rstest] + #[case::wider_i128_storage(DecimalArray::new( + buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, + ))] + #[case::wider_i256_storage(DecimalArray::new( + buffer![i256::from_i128(-99), i256::ZERO, i256::from_i128(99)], + DecimalDType::new(2, 0), Validity::NonNullable, + ))] + fn test_encode_rejects_lower_parts_past_canonical_width(#[case] decimal: DecimalArray) { + let result = DecimalByteParts::encode(&decimal, &mut SESSION.create_execution_ctx()); + assert!( + result + .as_ref() + .is_err_and(|err| err.to_string().contains("allows at most 0 lower parts")), + "expected the lower-part bound to reject the parts, got {result:?}" + ); + } } diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 2557555eac8..8390d3787ad 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -1,18 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -mod decimal_byte_parts; +//! An encoding that compresses decimals with integer compression schemes. +//! +//! A decimal is split into a signed most significant part (MSP) plus up to three `u64` lower +//! parts, and each part can be compressed independently. Only the MSP narrows: it can use the +//! smallest signed integer dtype that holds its values, while every lower part stays `u64`. +//! Decoding returns the canonical decimal array, or a decimal [`Widen`] over a narrower child that +//! can reuse a narrow MSP. +//! +//! [`Widen`]: vortex_array::arrays::Widen -use decimal_byte_parts::compute::is_constant::DecimalBytePartsIsConstantKernel; -/// This encoding allow compression of decimals using integer compression schemes. -/// Decimals can be compressed by narrowing the signed decimal value into the smallest signed value, -/// then integer compression if that is a value `ptype`, otherwise the decimal can be split into -/// parts. -/// These parts can be individually compressed. -/// This encoding will compress large signed decimals by removing the leading zeroes (after the sign) -/// an i128 decimal could be converted into a [i64, u64] with further narrowing applied to either -/// value. -pub use decimal_byte_parts::*; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; @@ -20,6 +18,11 @@ use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::session::ArraySessionExt; use vortex_session::VortexSession; +use crate::decimal_byte_parts::compute::is_constant::DecimalBytePartsIsConstantKernel; + +mod decimal_byte_parts; +pub use decimal_byte_parts::*; + /// Initialize decimal-byte-parts encoding in the given session. pub fn initialize(session: &VortexSession) { // One plugin owns both serialized formats: registering it reads either ID and writes the diff --git a/scripts/codspeed_bench_counts.json b/scripts/codspeed_bench_counts.json index e93ed45edfc..c393ac0c5fe 100644 --- a/scripts/codspeed_bench_counts.json +++ b/scripts/codspeed_bench_counts.json @@ -105,7 +105,7 @@ }, "vortex-decimal-byte-parts": { "dbp_assemble": 0, - "dbp_compute": 28, + "dbp_compute": 32, "dbp_split": 0 }, "vortex-elias-fano": { diff --git a/vortex-cuda/kernels/src/decimal_cast.cu b/vortex-cuda/kernels/src/decimal_cast.cu index 45c8c75e2a0..f55569a4a70 100644 --- a/vortex-cuda/kernels/src/decimal_cast.cu +++ b/vortex-cuda/kernels/src/decimal_cast.cu @@ -6,17 +6,16 @@ #include #include -// Arrow decimal schemas fix the physical values buffer width: -// - Decimal32: 4 bytes per value. -// - Decimal64: 8 bytes per value. -// - Decimal128: 16 bytes per value. -// - Decimal256: 32 bytes per value. +// Widen decimal values buffers to a wider physical width. Two callers use these kernels: +// - The Arrow export casts to the width that the Arrow schema fixes: 4 bytes for Decimal32, +// 8 bytes for Decimal64, 16 bytes for Decimal128, and 32 bytes for Decimal256. +// - The DecimalByteParts decoder casts a narrow MSP to the canonical width of its decimal +// dtype, from 2 to 32 bytes. // -// Vortex storage width can differ, so export casts to the schema-implied width. -// Rust-side export rejects narrowing casts because detecting overflow on-device -// would require synchronizing an overflow flag back to the host. +// The Rust side rejects narrowing casts, because detecting overflow on the device would require +// synchronizing an overflow flag back to the host. -// Low 64-bit conversion for Decimal32/64 outputs. +// Low 64-bit conversion for 16-, 32-, and 64-bit outputs. template __device__ __forceinline__ int64_t decimal_to_i64(Input value) { if constexpr (std::is_same_v) { @@ -42,10 +41,12 @@ __device__ __forceinline__ int128_t decimal_to_i128(Input value) { } } -// Convert one value to the Arrow schema's physical width. +// Convert one value to the output physical width. template __device__ __forceinline__ Output decimal_cast_value(Input value) { - if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return static_cast(decimal_to_i64(value)); + } else if constexpr (std::is_same_v) { return static_cast(decimal_to_i64(value)); } else if constexpr (std::is_same_v) { return decimal_to_i64(value); @@ -63,7 +64,7 @@ __device__ __forceinline__ Output decimal_cast_value(Input value) { } } -// Cast a contiguous values buffer to the Arrow schema's physical width. +// Cast a contiguous values buffer to the output physical width. template __device__ void decimal_cast_device(const Input *__restrict input, Output *__restrict output, uint64_t array_len) { @@ -76,8 +77,13 @@ decimal_cast_device(const Input *__restrict input, Output *__restrict output, ui } } -// Generate Decimal32/64/128/256 cast kernels for one input storage type. +// Generate the i16, i32, i64, i128, and i256 cast kernels for one input storage type. #define GENERATE_DECIMAL_CAST_KERNELS(input_suffix, InputType) \ + extern "C" __global__ void decimal_cast_##input_suffix##_i16(const InputType *__restrict input, \ + int16_t *__restrict output, \ + uint64_t array_len) { \ + decimal_cast_device(input, output, array_len); \ + } \ extern "C" __global__ void decimal_cast_##input_suffix##_i32(const InputType *__restrict input, \ int32_t *__restrict output, \ uint64_t array_len) { \ diff --git a/vortex-cuda/src/arrow/canonical.rs b/vortex-cuda/src/arrow/canonical.rs index 4bb7714b100..40ebdba697b 100644 --- a/vortex-cuda/src/arrow/canonical.rs +++ b/vortex-cuda/src/arrow/canonical.rs @@ -43,15 +43,12 @@ use vortex::array::arrays::varbinview::VarBinViewDataParts; use vortex::array::buffer::BufferHandle; use vortex::array::buffer::DeviceBuffer; use vortex::array::builtins::ArrayBuiltins; -use vortex::array::match_each_decimal_value_type; use vortex::array::validity::Validity; use vortex::buffer::Buffer; use vortex::dtype::DType; use vortex::dtype::DecimalType; -use vortex::dtype::NativeDecimalType; use vortex::dtype::Nullability; use vortex::dtype::PType; -use vortex::dtype::i256; use vortex::encodings::fsst::FSST; use vortex::error::VortexResult; use vortex::error::vortex_bail; @@ -60,7 +57,6 @@ use vortex::error::vortex_ensure_eq; use vortex::error::vortex_err; use vortex::extension::datetime::AnyTemporal; use vortex_cuda_driver::CudaSlice; -use vortex_cuda_driver::DeviceRepr; use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::PushKernelArg; use vortex_onpair::OnPair; @@ -89,6 +85,7 @@ use crate::executor::execute_validity_cuda; use crate::kernel::DecodedVarBin; use crate::kernel::decode_fsst_varbin; use crate::kernel::decode_onpair_varbin; +use crate::kernel::widen_decimal_values; use crate::stream::zero_padding; const BITMAP_THREADS_PER_BLOCK: u32 = 256; @@ -472,73 +469,11 @@ async fn export_decimal_values( "cannot export decimal values from {values_type} storage to Arrow {target_type}: narrowing would require a device-to-host overflow check", ); } - let values = ctx.ensure_on_device(values).await?; if values_type == target_type { - return Ok(values); + return ctx.ensure_on_device(values).await; } - match_each_decimal_value_type!(values_type, |S| { - export_decimal_values_from::(values, target_type, len, ctx).await - }) -} - -/// Dispatch from a concrete storage type `S` to the Arrow-required output width. -async fn export_decimal_values_from( - values: BufferHandle, - target_type: DecimalType, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, -{ - match target_type { - DecimalType::I32 => decimal_cast::(values, len, ctx).await, - DecimalType::I64 => decimal_cast::(values, len, ctx).await, - DecimalType::I128 => decimal_cast::(values, len, ctx).await, - DecimalType::I256 => decimal_cast::(values, len, ctx).await, - target_type => { - vortex_bail!("cannot export DecimalArray as Arrow decimal value type {target_type}") - } - } -} - -/// Launches the CUDA kernel that casts from Vortex storage type `S` to Arrow output type `D`. -/// -/// The caller must ensure this is not a narrowing cast. -async fn decimal_cast( - values: BufferHandle, - len: usize, - ctx: &mut CudaExecutionCtx, -) -> VortexResult -where - S: NativeDecimalType + DeviceRepr, - D: NativeDecimalType + DeviceRepr, -{ - if len == 0 { - return ctx - .ensure_on_device(BufferHandle::new_host( - Buffer::::empty().into_byte_buffer(), - )) - .await; - } - - let mut output_buffer = ctx.device_alloc::(len)?; - - let values_view = values.cuda_view::()?; - let len_u64 = len as u64; - let cuda_function = ctx.load_function_with_suffixes( - "decimal_cast", - &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], - )?; - - ctx.launch_kernel(&cuda_function, len, |args| { - args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); - })?; - - Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( - output_buffer, - )))) + widen_decimal_values(values, values_type, target_type, len, ctx).await } /// Export Vortex binary views as an Arrow Device array with `Utf8View`/`BinaryView` layout. diff --git a/vortex-cuda/src/kernel/decimal_cast.rs b/vortex-cuda/src/kernel/decimal_cast.rs new file mode 100644 index 00000000000..7a517a2ee06 --- /dev/null +++ b/vortex-cuda/src/kernel/decimal_cast.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Widening casts of decimal values buffers on the device. +//! +//! The Arrow export widens decimal values to the width that the Arrow schema fixes, and the +//! DecimalByteParts decoder widens a narrow MSP to the canonical width of its decimal dtype. Both +//! launch the `decimal_cast` kernels through this module. +//! +//! These casts only widen. A narrowing cast can overflow, and to detect that on the device, the +//! kernel would need an overflow flag that synchronizes with the host. So this module rejects a +//! narrowing cast, and each caller decides how to handle values that are wider than its target. + +use std::sync::Arc; + +use vortex::array::buffer::BufferHandle; +use vortex::array::match_each_decimal_value_type; +use vortex::buffer::Buffer; +use vortex::dtype::DecimalType; +use vortex::dtype::NativeDecimalType; +use vortex::dtype::i256; +use vortex::error::VortexResult; +use vortex::error::vortex_ensure; +use vortex::error::vortex_panic; +use vortex_cuda_driver::DeviceRepr; +use vortex_cuda_driver::PushKernelArg; + +use crate::CudaBufferExt; +use crate::CudaDeviceBuffer; +use crate::CudaExecutionCtx; + +/// Widen `len` decimal values from `values_type` to the wider `target_type` on the device. +/// +/// The values can be on the host or on the device, and the result is on the device. Each value is +/// sign-extended, so the result holds the same values as the input. +/// +/// # Errors +/// +/// Returns an error if `target_type` is not wider than `values_type`, or if the copy to the device +/// or the kernel launch fails. +pub(crate) async fn widen_decimal_values( + values: BufferHandle, + values_type: DecimalType, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult { + vortex_ensure!( + values_type < target_type, + "a decimal cast on the GPU must widen to a type wider than {values_type}, got \ + {target_type}" + ); + let values = ctx.ensure_on_device(values).await?; + + match_each_decimal_value_type!(values_type, |S| { + decimal_cast_to::(values, target_type, len, ctx).await + }) +} + +/// Dispatch [`decimal_cast`] from `S` on the runtime `target_type`. +async fn decimal_cast_to( + values: BufferHandle, + target_type: DecimalType, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, +{ + match target_type { + DecimalType::I8 => vortex_panic!( + "widen_decimal_values ensures a target wider than the values type, and no decimal \ + type is narrower than i8, got {target_type}" + ), + DecimalType::I16 => decimal_cast::(values, len, ctx).await, + DecimalType::I32 => decimal_cast::(values, len, ctx).await, + DecimalType::I64 => decimal_cast::(values, len, ctx).await, + DecimalType::I128 => decimal_cast::(values, len, ctx).await, + DecimalType::I256 => decimal_cast::(values, len, ctx).await, + } +} + +/// Launch the `decimal_cast__` kernel, which casts device values of type `S` to type `D`. +/// +/// The cast must widen. A narrowing cast truncates each value with no overflow check. +async fn decimal_cast( + values: BufferHandle, + len: usize, + ctx: &mut CudaExecutionCtx, +) -> VortexResult +where + S: NativeDecimalType + DeviceRepr, + D: NativeDecimalType + DeviceRepr, +{ + if len == 0 { + return ctx + .ensure_on_device(BufferHandle::new_host( + Buffer::::empty().into_byte_buffer(), + )) + .await; + } + + let mut output_buffer = ctx.device_alloc::(len)?; + + let values_view = values.cuda_view::()?; + let len_u64 = len as u64; + let cuda_function = ctx.load_function_with_suffixes( + "decimal_cast", + &[&S::DECIMAL_TYPE.to_string(), &D::DECIMAL_TYPE.to_string()], + )?; + + ctx.launch_kernel(&cuda_function, len, |args| { + args.arg(&values_view).arg(&mut output_buffer).arg(&len_u64); + })?; + + Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new( + output_buffer, + )))) +} diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index ee3814ca9fc..7b93b7a4930 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -9,15 +9,18 @@ use vortex::array::ArrayRef; use vortex::array::Canonical; use vortex::array::arrays::DecimalArray; use vortex::array::arrays::primitive::PrimitiveDataParts; +use vortex::dtype::DecimalType; use vortex::encodings::decimal_byte_parts::DecimalByteParts; use vortex::encodings::decimal_byte_parts::DecimalBytePartsArraySlotsExt; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex::error::vortex_ensure; use crate::CudaExecutionCtx; use crate::executor::CudaArrayExt; use crate::executor::CudaExecute; +use crate::kernel::widen_decimal_values; // See `DecimalBytePartsArray` #[derive(Debug)] @@ -47,6 +50,7 @@ impl CudaExecute for DecimalBytePartsExecutor { } let msp = array.msp().clone(); + let len = msp.len(); let PrimitiveDataParts { buffer, ptype, @@ -58,10 +62,32 @@ impl CudaExecute for DecimalBytePartsExecutor { .into_primitive() .into_data_parts(); - // SAFETY: The primitive array's buffer is already validated with correct type. - // The decimal dtype matches the array's dtype, and validity is preserved. + // NB: Older writers can store an MSP wider than the canonical width, because they narrowed + // a decimal by the min and max of every slot, null slots included. The device cannot narrow + // it, because `widen_decimal_values` only widens. The CUDA Arrow export also rejects such + // storage, so these columns failed on the GPU before this check as well. + let msp_type = DecimalType::try_from(ptype)?; + let canonical_type = decimal_dtype.canonical_values_type(); + vortex_ensure!( + msp_type <= canonical_type, + "DecimalBytePartsArray on GPU requires an MSP no wider than the canonical width \ + {canonical_type}, got {msp_type}" + ); + + // An MSP at the canonical width already holds the decimal values, so it needs no copy. + // The CPU decoder returns a `Widen` over a narrow MSP, but CUDA has no `Widen` kernel, so + // the device widens to the canonical width. + let values = if msp_type == canonical_type { + buffer + } else { + widen_decimal_values(buffer, msp_type, canonical_type, len, ctx).await? + }; + + // SAFETY: The values hold one value per row of the MSP at the canonical width of + // `decimal_dtype`, so the validity of the MSP fits them. A widening copy keeps every + // value, and reads trust the values of the MSP to be within the precision. Ok(Canonical::Decimal(unsafe { - DecimalArray::new_unchecked_handle(buffer, ptype.try_into()?, decimal_dtype, validity) + DecimalArray::new_unchecked_handle(values, canonical_type, decimal_dtype, validity) })) } } @@ -78,42 +104,72 @@ mod tests { use vortex::buffer::Buffer; use vortex::dtype::DecimalDType; use vortex::encodings::decimal_byte_parts::DecimalByteParts; - use vortex::error::VortexExpect; + use vortex::error::VortexResult; use vortex::session::VortexSession; use vortex_array::VortexSessionExecute; use super::*; + use crate::CanonicalCudaExt; use crate::session::CudaSession; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); #[rstest] + #[case::i8_p3_s0(Buffer::from(vec![99i8, 1, 0, -1, -99]), 3, 0)] #[case::i8_p5_s2(Buffer::from(vec![100i8, 101, 102, -1, -100]), 5, 2)] #[case::i16_p10_s2(Buffer::from(vec![100i16, 200, 300, 400, 500]), 10, 2)] + #[case::i32_p9_s2(Buffer::from(vec![100i32, 200, 300, -400, -500]), 9, 2)] #[case::i32_p18_s4(Buffer::from(vec![100i32, 200, 300, 400, 500]), 18, 4)] #[case::i64_p38_s6(Buffer::from(vec![100i64, 200, 300, 400, 500]), 38, 6)] + #[case::i64_p40_s2(Buffer::from(vec![i64::MAX, 1, 0, -1, i64::MIN]), 40, 2)] #[crate::test] async fn test_decimal_byte_parts_gpu_decode( #[case] encoded: Buffer, #[case] precision: u8, #[case] scale: i8, - ) { + ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session()) - .vortex_expect("create execution context"); + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; let decimal_dtype = DecimalDType::new(precision, scale); let dbp_array = DecimalByteParts::try_new( PrimitiveArray::new(encoded, Validity::NonNullable).into_array(), decimal_dtype, - ) - .vortex_expect("create DecimalBytePartsArray"); + )?; let gpu_result = DecimalBytePartsExecutor .execute(dbp_array.clone().into_array(), &mut cuda_ctx) - .await - .vortex_expect("GPU decode"); + .await? + .into_host(cuda_ctx.execution_ctx()) + .await? + .into_decimal(); + + assert_eq!( + gpu_result.values_type(), + decimal_dtype.canonical_values_type() + ); + assert_arrays_eq!(gpu_result, dbp_array, &mut ctx); + Ok(()) + } - assert_arrays_eq!(dbp_array, gpu_result.into_array(), &mut ctx); + #[crate::test] + async fn test_decimal_byte_parts_gpu_decode_rejects_msp_wider_than_canonical() + -> VortexResult<()> { + let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + + let dbp_array = DecimalByteParts::try_new( + PrimitiveArray::new(Buffer::from(vec![99i64, -99]), Validity::NonNullable).into_array(), + DecimalDType::new(4, 0), + )?; + + let err = DecimalBytePartsExecutor + .execute(dbp_array.into_array(), &mut cuda_ctx) + .await + .expect_err("CUDA decoding rejects an MSP wider than the canonical decimal width"); + assert!( + err.to_string() + .contains("requires an MSP no wider than the canonical width") + ); + Ok(()) } } diff --git a/vortex-cuda/src/kernel/mod.rs b/vortex-cuda/src/kernel/mod.rs index f1eb24b1fad..94d15ffa6d0 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -19,6 +19,10 @@ use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::sys::CUevent_flags; mod arrays; + +mod decimal_cast; +pub(crate) use decimal_cast::widen_decimal_values; + // Fatbins let the CUDA driver select among the configured real and virtual architectures. mod embedded_kernels { include!(concat!(env!("OUT_DIR"), "/embedded_kernels.rs")); diff --git a/vortex/benches/common_encoding_tree_throughput.rs b/vortex/benches/common_encoding_tree_throughput.rs index de446630fd7..5fc06308bab 100644 --- a/vortex/benches/common_encoding_tree_throughput.rs +++ b/vortex/benches/common_encoding_tree_throughput.rs @@ -15,7 +15,6 @@ use mimalloc::MiMalloc; use rand::RngExt; use rand::SeedableRng; use vortex::array::ArrayRef; -use vortex::array::Canonical; use vortex::array::IntoArray; use vortex::array::VortexSessionExecute; use vortex::array::aggregate_fn::session::AggregateFnSession; @@ -24,6 +23,7 @@ use vortex::array::arrays::PrimitiveArray; use vortex::array::arrays::TemporalArray; use vortex::array::arrays::VarBinArray; use vortex::array::arrays::VarBinViewArray; +use vortex::array::arrays::Widen; use vortex::array::arrays::varbin::VarBinArraySlotsExt; use vortex::array::builtins::ArrayBuiltins; use vortex::array::memory::MemorySession; @@ -491,7 +491,11 @@ macro_rules! setup_fn { }; } -/// Benchmark decompression of various encoding trees +/// Benchmark decompression of various encoding trees. +/// +/// Decompression stops at a [`Widen`], the form that decoding a narrow DecimalByteParts MSP +/// produces. The other trees never yield a `Widen`, so they do the same work as a full +/// canonicalization. #[divan::bench( args = [ setup_fn!(setup::for_bp_u64), @@ -512,5 +516,5 @@ fn decompress(bencher: Bencher, setup_fn: SetupFn) { with_byte_counter(bencher, nbytes) .with_inputs(|| (&compressed, SESSION.create_execution_ctx())) - .bench_refs(|(a, ctx)| (**a).clone().execute::(ctx)); + .bench_refs(|(a, ctx)| (**a).clone().execute_until::(ctx)); }