From 93d56ad276934efec6c8710e47e77828ddaadbee Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 17:22:39 -0400 Subject: [PATCH 1/4] Require u64 lower parts in DecimalByteParts Every lower part of a DecimalByteParts array now has the dtype `u64`. Only the MSP may use a narrower signed integer dtype. The constructor, `assemble_decimal` and the v2 reader reject `u8`, `u16` and `u32` lower parts, and the error names `u64`. The constructor also bounds the lower-part count, so the parts never assemble wider than the canonical width of the decimal dtype. Precision up to 18 allows no lower parts, 19 to 38 at most one, and 39 to 76 up to three. `assemble_decimal` runs the same validation as the constructor. `DecimalByteParts::encode` now uses the checked constructor, because a decimal stored wider than its canonical width splits into more lower parts than its dtype allows. The unchecked constructor had no other caller, so it is removed. Assembly no longer casts lower parts to `u64`. The docs, tests and the `dbp_assemble_narrowed` bench now use `u64` lower parts. Signed-off-by: Connor Tsui --- .../benches/dbp_assemble.rs | 13 +- .../src/decimal_byte_parts/array.rs | 149 ++++++++++-------- .../src/decimal_byte_parts/assemble.rs | 47 ++---- .../src/decimal_byte_parts/mod.rs | 7 +- .../src/decimal_byte_parts/plugin/tests.rs | 54 ++++--- .../src/decimal_byte_parts/plugin/v2.rs | 6 +- .../src/decimal_byte_parts/split.rs | 10 +- encodings/decimal-byte-parts/src/lib.rs | 14 +- 8 files changed, 151 insertions(+), 149 deletions(-) diff --git a/encodings/decimal-byte-parts/benches/dbp_assemble.rs b/encodings/decimal-byte-parts/benches/dbp_assemble.rs index 33ab5efbf2a..c7c6557f047 100644 --- a/encodings/decimal-byte-parts/benches/dbp_assemble.rs +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -2,7 +2,7 @@ // 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; @@ -145,15 +145,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/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs index f0b03720e00..9ccf2e7371d 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -23,6 +23,7 @@ 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; @@ -40,6 +41,9 @@ 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; @@ -51,11 +55,14 @@ 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. +/// 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 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. +/// 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. /// /// e.g. for a decimal i128 \[ 127..64 | 63..0 \] msp = 127..64 and lower_part\[0\] = 63..0 /// @@ -85,8 +92,9 @@ impl DecimalBytePartsData { /// # Errors /// /// Returns an error if the MSP is not a signed integer array of length `len`, if `dtype` - /// does not match the MSP's nullability, if there are more than `MAX_LOWER_PARTS` - /// lower parts, or if any lower part is not a non-nullable unsigned integer array of length `len`. + /// 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, @@ -103,15 +111,16 @@ impl DecimalBytePartsData { 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, - "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" + 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().is_unsigned_int() && !part.dtype().is_nullable(), - "lower part {idx} must have a non-nullable unsigned integer dtype, got {}", + 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"); @@ -120,6 +129,18 @@ impl DecimalBytePartsData { } } +/// The most lower parts that assemble to at most the canonical width of `decimal_dtype`. +/// +/// An MSP assembles to `i64` with no lower parts, to `i128` with one, and to `i256` with two or +/// three. 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, + } +} + #[derive(Clone, Debug)] pub struct DecimalByteParts; @@ -139,8 +160,8 @@ impl DecimalByteParts { /// 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 + /// 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 @@ -163,34 +184,6 @@ impl DecimalByteParts { 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 { @@ -369,7 +362,6 @@ impl OperationsVTable for DecimalByteParts { 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() @@ -577,12 +569,7 @@ mod tests { } #[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], + fn test_narrow_msp_with_u64_lower_parts( #[values(PType::I8, PType::I16, PType::I32, PType::I64)] msp_ptype: PType, #[values(1, 2, 3)] lower_count: usize, ) -> VortexResult<()> { @@ -591,18 +578,14 @@ mod tests { 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 + // Set the highest bit of each word to catch accidental sign extension. The low bits + // differ, so a swapped word shows up too. + let words = [(1u64 << 63) | 1, (1 << 63) | 2, (1 << 63) | 3]; + let lower = words .into_iter() - .zip(words) .take(lower_count) - .map(|(ptype, word)| { - buffer![0u64, word, word] - .into_array() - .cast(DType::Primitive(ptype, Nullability::NonNullable)) - }) - .collect::>>()?; + .map(|word| buffer![0u64, word, word].into_array()) + .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 { @@ -659,22 +642,62 @@ mod tests { } #[rstest] - #[case::signed_lower_part(vec![buffer![1i64, 2, 3].into_array()], DecimalDType::new(38, 2))] + #[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) + DecimalDType::new(38, 2), + "must have dtype u64, got u64?" + )] + #[case::u8_lower_part( + vec![buffer![1u8, 2, 3].into_array()], + DecimalDType::new(38, 2), + "must have dtype u64, got u8" + )] + #[case::u16_lower_part( + vec![buffer![1u16, 2, 3].into_array()], + DecimalDType::new(38, 2), + "must have dtype u64, got u16" + )] + #[case::u32_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::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) + 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!( - DecimalByteParts::try_new_with_lower_parts(msp(), lower_parts, decimal_dtype).is_err() + result + .as_ref() + .is_err_and(|err| err.to_string().contains(expected_error)), + "expected {expected_error}, got {result:?}" ); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs index 1cdf27c15bd..5053f3537d1 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs @@ -11,7 +11,7 @@ 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::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::NativeDecimalType; use vortex_array::dtype::i256; @@ -22,25 +22,20 @@ 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::DecimalBytePartsData; 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. +/// The parts must satisfy [`DecimalBytePartsData::validate`]: the MSP has a signed integer dtype, +/// every lower part is a non-nullable `u64` array, all parts have the same length, and the lower +/// parts do not assemble wider than the canonical width of `decimal_dtype`. /// /// 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. @@ -50,10 +45,8 @@ pub fn assemble_decimal( decimal_dtype: DecimalDType, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - vortex_ensure!( - msp.dtype().is_signed_int(), - "MSP must have a signed integer dtype" - ); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + DecimalBytePartsData::validate(msp, lower_parts.iter(), decimal_dtype, &dtype, msp.len())?; let validity = msp.validity()?; @@ -61,21 +54,6 @@ pub fn assemble_decimal( 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) } @@ -92,7 +70,7 @@ fn assemble_narrow_decimal( })) } -/// Execute the MSP at its signed integer width and cast lower parts to `u64` before assembly. +/// Execute the MSP at its signed integer width and the `u64` lower parts, then assemble them. /// 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( @@ -104,15 +82,9 @@ fn assemble_wide_decimal_from_arrays( ) -> 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) - }) + .map(|part| part.clone().execute::(exec_ctx)) .collect::>>()?; Ok(match_each_signed_integer_ptype!(msp.ptype(), |Msp| { @@ -389,6 +361,7 @@ mod tests { #[rstest] #[case::signed(PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable))] + #[case::narrow(PrimitiveArray::new(buffer![0u32; 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))] 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..0790a750786 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -13,7 +13,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. +//! The MSP is a signed integer of at most 64 bits, and it may use a narrower dtype when its values +//! fit. Every lower part has the dtype `u64`. An array has at most as many lower parts as the +//! canonical width of its decimal dtype holds, so the assembled values are never wider than that +//! width. use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; @@ -59,5 +62,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..2fb7fbeffe6 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 @@ -47,39 +47,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), ))] @@ -135,6 +124,25 @@ fn serde_round_trip(#[case] array: VortexResult) -> Vorte 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 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:?}" + ); +} + #[test] fn v1_metadata_is_unchanged() -> VortexResult<()> { let session = session(); @@ -183,7 +191,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()) } } diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 2557555eac8..bb279aa74f9 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -4,14 +4,12 @@ mod decimal_byte_parts; 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. +/// This encoding allows compression of decimals using integer compression schemes. +/// A decimal is split into a signed most significant part (MSP) and up to three `u64` lower +/// parts, and each part can be compressed independently. +/// A decimal whose values fit an `i64` needs no lower parts. +/// Only the MSP can be narrowed: it can use the smallest signed integer dtype that holds its +/// values, while every lower part stays `u64`. pub use decimal_byte_parts::*; use vortex_array::ArrayVTable; use vortex_array::aggregate_fn::AggregateFnVTable; From 64cfbd3f6409bd2f8ad626ab71e492d653efde59 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 17:25:31 -0400 Subject: [PATCH 2/4] Decode DecimalByteParts at the canonical width `assemble_decimal` now returns a decimal array at the canonical width of its dtype, whatever the width of the MSP. With no lower parts, an MSP at the canonical width is reused with no copy and no scan. A narrower MSP is widened into a copy, and a TODO marks where a later change can return a `Widen` instead. A wider MSP is narrowed with `copy_to_canonical_width`, which checks every valid value against the precision and writes zeros in the null slots. One lower part assembles into `i128`, or straight into `i256` at a precision above 38. Two or three lower parts assemble into `i256`. The constructor bounds the lower parts by the canonical width, so assembly never narrows. Every result is built with the unchecked constructor, because its width is the canonical width and reads trust the precision. Signed-off-by: Connor Tsui --- .../src/decimal_byte_parts/assemble.rs | 337 ++++++++++++++---- .../src/decimal_byte_parts/mod.rs | 2 +- .../src/decimal_byte_parts/plugin/tests.rs | 5 +- 3 files changed, 275 insertions(+), 69 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs index 5053f3537d1..9076e1f5e60 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs @@ -11,8 +11,11 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::decimal::copy_to_canonical_width; +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::match_each_signed_integer_ptype; @@ -21,24 +24,33 @@ 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_panic; use super::DecimalBytePartsData; use super::LOWER_PART_BITS; -use super::MAX_LOWER_PARTS; -/// Reassemble decimal parts into a decimal array. +/// Reassemble decimal parts into the canonical decimal array of `decimal_dtype`. /// /// The parts must satisfy [`DecimalBytePartsData::validate`]: the MSP has a signed integer dtype, /// every lower part is a non-nullable `u64` array, all parts have the same length, and the lower /// parts do not assemble wider than the canonical width of `decimal_dtype`. /// -/// 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`. +/// The result always has the canonical width of `decimal_dtype`: +/// +/// - With no lower parts, an MSP at the canonical width is reused as the decimal values. A +/// narrower MSP is widened into a copy. A wider MSP is narrowed into a copy, which checks every +/// valid value against the precision and writes zeros in the null slots (see +/// [`copy_to_canonical_width`]). +/// - One lower part assembles into `i128`, or into `i256` at a precision above 38. Two or three +/// lower parts assemble into `i256`. +/// +/// Only the narrowing copy checks the precision. Every other path trusts that the parts hold +/// values within the precision, so it does not scan them. /// /// # Errors /// -/// Returns an error for invalid part dtypes, lengths, or counts, or if executing a part fails. +/// Returns an error for invalid part dtypes, lengths, or counts, if executing a part fails, or if +/// a narrowing copy finds a valid value outside the precision of `decimal_dtype`. pub fn assemble_decimal( msp: &ArrayRef, lower_parts: &[ArrayRef], @@ -49,37 +61,75 @@ pub fn assemble_decimal( DecimalBytePartsData::validate(msp, lower_parts.iter(), decimal_dtype, &dtype, msp.len())?; let validity = msp.validity()?; + let values = if lower_parts.is_empty() { + msp_to_canonical_width(msp, &validity, decimal_dtype, exec_ctx)? + } else { + assemble_wide_decimal_from_arrays(msp, lower_parts, decimal_dtype, exec_ctx)? + }; - if lower_parts.is_empty() { - return assemble_narrow_decimal(msp, validity, decimal_dtype, exec_ctx); - } + // SAFETY: Both helpers return one value per row of the MSP at the canonical width of + // `decimal_dtype`, so the validity of the MSP fits the values. Every valid value is within the + // precision: a narrowing copy of the MSP checks it, a widening copy cannot leave it, and the + // other paths keep the values of the parts, which reads trust to be within the precision. + let decimal = unsafe { + DecimalArray::new_unchecked_handle( + values, + decimal_dtype.canonical_values_type(), + decimal_dtype, + validity, + ) + }; - assemble_wide_decimal_from_arrays(msp, lower_parts, validity, decimal_dtype, exec_ctx) + Ok(decimal.into_array()) } -fn assemble_narrow_decimal( +/// Execute an MSP that has no lower parts, and return its values at the canonical width of +/// `decimal_dtype`. +fn msp_to_canonical_width( msp: &ArrayRef, - validity: Validity, + validity: &Validity, decimal_dtype: DecimalDType, exec_ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> 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() - })) + let msp_type = DecimalType::try_from(msp.ptype())?; + let canonical_type = decimal_dtype.canonical_values_type(); + + // An MSP at the canonical width already holds the decimal values, so it needs no copy. + if msp_type == canonical_type { + return Ok(BufferHandle::new_host(msp.buffer_handle().to_host_sync())); + } + + let mask = validity.execute_mask(msp.len(), exec_ctx)?; + + // TODO(connor): Find the max of a narrow MSP and return a `Widen` over a child that reuses + // its buffer, once benchmarks show that this widening copy costs too much. + if msp_type < canonical_type { + return match_each_signed_integer_ptype!(msp.ptype(), |P| { + copy_to_canonical_width(msp.as_slice::

(), &mask, decimal_dtype) + }); + } + + // The MSP is wider than the canonical width, so this copy narrows it. It fails on a valid + // value outside the precision and writes zeros in the null slots. + match_each_signed_integer_ptype!(msp.ptype(), |P| { + copy_to_canonical_width(msp.as_slice::

(), &mask, decimal_dtype) + }) } -/// Execute the MSP at its signed integer width and the `u64` lower parts, then assemble them. -/// The number of lower parts determines the decimal storage type: one produces `i128`, while -/// two or three produce `i256`. +/// Execute the MSP at its signed integer width and the `u64` lower parts, then assemble them at +/// the canonical width of `decimal_dtype`. +/// +/// One lower part assembles into `i128`, or into `i256` when that is the canonical width. Two or +/// three lower parts assemble into `i256`. [`DecimalBytePartsData::validate`] rejects lower parts +/// that assemble wider than the canonical width, so this never narrows. fn assemble_wide_decimal_from_arrays( msp: &ArrayRef, lower_parts: &[ArrayRef], - validity: Validity, decimal_dtype: DecimalDType, exec_ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult { // TODO(mk): Broadcast constant parts directly instead of materializing their buffers. let msp = msp.clone().execute::(exec_ctx)?; let lower = lower_parts @@ -87,50 +137,48 @@ fn assemble_wide_decimal_from_arrays( .map(|part| part.clone().execute::(exec_ctx)) .collect::>>()?; - Ok(match_each_signed_integer_ptype!(msp.ptype(), |Msp| { + let canonical_type = decimal_dtype.canonical_values_type(); + let values = 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, + match (lower.as_slice(), canonical_type) { + ([first], DecimalType::I128) => assemble_wide_decimal::( + msp, + first.as_slice::().iter().map(|&word| [word]), ) - .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_byte_buffer(), + ([first], DecimalType::I256) => assemble_wide_decimal::( + msp, + first.as_slice::().iter().map(|&word| [word]), ) - .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_byte_buffer(), + ([first, second], DecimalType::I256) => assemble_wide_decimal::( + msp, + first + .as_slice::() + .iter() + .zip(second.as_slice::()) + .map(|(&a, &b)| [a, b]), ) - .into_array(), - _ => { - vortex_bail!("expected between one and {MAX_LOWER_PARTS} lower parts") - } + .into_byte_buffer(), + ([first, second, third], DecimalType::I256) => 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(), + (lower, canonical_type) => vortex_panic!( + "validate must bound the lower parts by the canonical width, got {} lower parts \ + at {canonical_type}", + lower.len() + ), } - })) + }); + + Ok(BufferHandle::new_host(values)) } /// Assemble a signed MSP slice and one array of `K` lower words per row. @@ -176,6 +224,7 @@ mod tests { use std::sync::LazyLock; use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; @@ -183,6 +232,8 @@ mod tests { 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::BigCast; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::DecimalType; @@ -193,6 +244,7 @@ mod tests { 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; @@ -346,16 +398,171 @@ mod tests { validity: Validity, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); + 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.execute::(&mut ctx)?; + 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(()) + } + + /// Decode an MSP with no lower parts. + fn assemble_msp(msp: ArrayRef, decimal_dtype: DecimalDType) -> VortexResult { + let mut ctx = SESSION.create_execution_ctx(); + assemble_decimal(&msp, &[], decimal_dtype, &mut ctx)?.execute::(&mut ctx) + } + + /// The decimal array of `values` at the canonical width of `decimal_dtype`. + fn canonical_decimal( + values: &[i64], + decimal_dtype: DecimalDType, + validity: Validity, + ) -> DecimalArray { + match_each_decimal_value_type!(decimal_dtype.canonical_values_type(), |T| { + DecimalArray::new( + values + .iter() + .map(|&value| { + ::from(value).vortex_expect("value fits the canonical width") + }) + .collect::>(), + decimal_dtype, + validity, + ) + }) + } + + #[rstest] + #[case::reuse_i8_at_p2(PType::I8, 2)] + #[case::widen_i8_to_p4(PType::I8, 4)] + #[case::widen_i32_to_p18(PType::I32, 18)] + #[case::widen_i64_to_p38(PType::I64, 38)] + #[case::widen_i16_to_p76(PType::I16, 76)] + #[case::narrow_i64_to_p9(PType::I64, 9)] + #[case::narrow_i32_to_p2(PType::I32, 2)] + fn test_assemble_returns_canonical_width( + #[case] msp_ptype: PType, + #[case] precision: u8, + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] + validity: Validity, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(precision, 0); + let msp = PrimitiveArray::new(buffer![-5i64, 0, 7], validity.clone()) + .into_array() + .cast(DType::Primitive(msp_ptype, validity.nullability()))?; + + let decoded = assemble_msp(msp, decimal_dtype)?; + + assert_eq!(decoded.values_type(), decimal_dtype.canonical_values_type()); + assert_arrays_eq!( + decoded, + canonical_decimal(&[-5, 0, 7], decimal_dtype, validity), + &mut ctx + ); + Ok(()) + } + + #[test] + fn test_assemble_reuses_msp_buffer() -> VortexResult<()> { + let msp = PrimitiveArray::new(buffer![1i32, -2, 3], Validity::NonNullable); + let msp_ptr = msp.as_slice::().as_ptr(); + + let decoded = assemble_msp(msp.into_array(), DecimalDType::new(9, 0))?; + + assert_eq!(decoded.buffer::().as_ptr(), msp_ptr); + Ok(()) + } + + #[test] + fn test_assemble_narrowing_copy_zeroes_null_slots() -> VortexResult<()> { + let msp = PrimitiveArray::new( + buffer![1i64, 1000, -1], + Validity::from_iter([true, false, true]), + ); + + let decoded = assemble_msp(msp.into_array(), DecimalDType::new(2, 0))?; + + assert_eq!(decoded.values_type(), DecimalType::I8); + assert_eq!(decoded.buffer::().as_slice(), &[1, 0, -1]); + 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 = assemble_msp(msp.into_array(), DecimalDType::new(2, 0)); + + assert!( + result + .as_ref() + .is_err_and(|err| err.to_string().contains("must be within the precision")), + "expected the narrowing copy to reject 100, got {result:?}" + ); + } + + #[test] + fn test_assemble_widening_sign_extends() -> VortexResult<()> { + let msp = PrimitiveArray::new(buffer![-1i8, 1], Validity::NonNullable); + + let decoded = assemble_msp(msp.into_array(), DecimalDType::new(40, 0))?; + + assert_eq!(decoded.values_type(), DecimalType::I256); assert_eq!( - msp.as_slice::().as_ptr(), - decimal.buffer::().as_ptr() + decoded.buffer::().as_slice(), + &[i256::from_i128(-1), i256::ONE] ); - assert_arrays_eq!(round_trip(decimal.clone())?, decimal, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::one_lower_part_at_p38(DecimalArray::new( + buffer![(3i128 << 64) | 2, (-3i128 << 64) | 2], + DecimalDType::new(38, 0), + Validity::NonNullable, + ), 1)] + #[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)] + #[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)] + fn test_assemble_lower_parts_canonical_width( + #[case] expected: DecimalArray, + #[case] lower_count: usize, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = expected.decimal_dtype(); + let msp = 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)? + .execute::(&mut ctx)?; + + assert_eq!(decoded.values_type(), decimal_dtype.canonical_values_type()); + assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 0790a750786..54ef5d35ec0 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -16,7 +16,7 @@ //! The MSP is a signed integer of at most 64 bits, and it may use a narrower dtype when its values //! fit. Every lower part has the dtype `u64`. An array has at most as many lower parts as the //! canonical width of its decimal dtype holds, so the assembled values are never wider than that -//! width. +//! width. Decoding always produces a decimal array at the canonical width. use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; 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 2fb7fbeffe6..42f7a301b53 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 @@ -43,9 +43,8 @@ static SESSION: LazyLock = LazyLock::new(VortexSession::empty); #[case::one_lower_part(DecimalByteParts::try_new_with_lower_parts( msp(), vec![lower_part()], DecimalDType::new(38, 2), ))] -#[case::wider_i64_storage(DecimalByteParts::encode( - &DecimalArray::new(buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable), - &mut SESSION.create_execution_ctx(), +#[case::wider_i64_storage(DecimalByteParts::try_new( + buffer![-99i64, 0, 99].into_array(), DecimalDType::new(2, 0), ))] #[case::redundant_lower_parts(DecimalByteParts::try_new_with_lower_parts( buffer![0i64; 3].into_array(), From 07eee640e5eae928ba23c26cc6661547f9d203f8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 27 Oct 2026 17:33:59 -0400 Subject: [PATCH 3/4] Decode DecimalByteParts on the GPU at the canonical width The CUDA DecimalByteParts decoder now returns a decimal array at the canonical width of its dtype, as the CPU decoder does. An MSP at the canonical width is reused with no copy. A narrower MSP is widened on the device with the decimal_cast kernels. An MSP wider than the canonical width returns an error: narrowing it must check the precision, and that check on the device needs an overflow flag that synchronizes with the host. The btrblocks scheme never writes such an MSP. Lower parts still return an error on the GPU. The device widening moves out of the Arrow export into a crate-private kernel module, so the export and the decoder share one entry point. That entry point rejects a narrowing cast. The decimal_cast kernels gain an i16 output, because an i8 MSP at precision 3 or 4 widens to i16. The GPU tests now assert the canonical width, and cover an i8 MSP at precision 3, an i64 MSP at precision 40, an i32 MSP at precision 9 that is reused, and an i64 MSP at precision 4 that is rejected. Signed-off-by: Connor Tsui --- vortex-cuda/kernels/src/decimal_cast.cu | 32 +++-- vortex-cuda/src/arrow/canonical.rs | 71 +---------- vortex-cuda/src/kernel/decimal_cast.rs | 118 ++++++++++++++++++ .../kernel/encodings/decimal_byte_parts.rs | 77 ++++++++++-- vortex-cuda/src/kernel/mod.rs | 2 + 5 files changed, 207 insertions(+), 93 deletions(-) create mode 100644 vortex-cuda/src/kernel/decimal_cast.rs diff --git a/vortex-cuda/kernels/src/decimal_cast.cu b/vortex-cuda/kernels/src/decimal_cast.cu index 45c8c75e2a0..86cd092e24f 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..60586d4a710 --- /dev/null +++ b/vortex-cuda/src/kernel/decimal_cast.rs @@ -0,0 +1,118 @@ +// 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!("i8 is the narrowest decimal type, so no cast widens to it") + } + 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..8263b796116 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -9,6 +9,7 @@ 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; @@ -18,6 +19,7 @@ use vortex::error::vortex_bail; use crate::CudaExecutionCtx; use crate::executor::CudaArrayExt; use crate::executor::CudaExecute; +use crate::kernel::widen_decimal_values; // See `DecimalBytePartsArray` #[derive(Debug)] @@ -47,6 +49,7 @@ impl CudaExecute for DecimalBytePartsExecutor { } let msp = array.msp().clone(); + let len = msp.len(); let PrimitiveDataParts { buffer, ptype, @@ -58,10 +61,30 @@ 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. + // Narrowing the MSP must check every valid value against the precision, and that check on + // the device needs an overflow flag that synchronizes with the host. The btrblocks scheme + // never writes an MSP wider than the canonical width, so the GPU does not support one. + let msp_type = DecimalType::try_from(ptype)?; + let canonical_type = decimal_dtype.canonical_values_type(); + if msp_type > canonical_type { + vortex_bail!( + "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. + 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 +101,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(()) + } + + #[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())?; - assert_arrays_eq!(dbp_array, gpu_result.into_array(), &mut ctx); + 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 + .unwrap_err(); + 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 f81a5c05119..c919bc54463 100644 --- a/vortex-cuda/src/kernel/mod.rs +++ b/vortex-cuda/src/kernel/mod.rs @@ -19,6 +19,7 @@ use vortex_cuda_driver::LaunchConfig; use vortex_cuda_driver::sys::CUevent_flags; mod arrays; +mod decimal_cast; // Fatbins let the CUDA driver select among the configured real and virtual architectures. mod embedded_kernels { include!(concat!(env!("OUT_DIR"), "/embedded_kernels.rs")); @@ -33,6 +34,7 @@ pub(crate) use arrays::DictExecutor; pub(crate) use arrays::ListExecutor; pub(crate) use arrays::MaskedExecutor; pub(crate) use arrays::SharedExecutor; +pub(crate) use decimal_cast::widen_decimal_values; pub use encodings::ZstdKernelPrep; pub use encodings::zstd_kernel_prepare; pub(crate) use encodings::*; From ed4f73008ffe700c13ecb7ede90fddc1fbff4435 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 13:58:29 -0400 Subject: [PATCH 4/4] Share the DecimalByteParts MSP copy path Use one copy_to_canonical_width call for both widening and narrowing an MSP after the zero-copy reuse path. Keep the widening TODO and the explanation of precision checks and null-slot zeros. Signed-off-by: Connor Tsui --- .../src/decimal_byte_parts/assemble.rs | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs index 9076e1f5e60..2a2c8297e10 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs @@ -105,14 +105,9 @@ fn msp_to_canonical_width( // TODO(connor): Find the max of a narrow MSP and return a `Widen` over a child that reuses // its buffer, once benchmarks show that this widening copy costs too much. - if msp_type < canonical_type { - return match_each_signed_integer_ptype!(msp.ptype(), |P| { - copy_to_canonical_width(msp.as_slice::

(), &mask, decimal_dtype) - }); - } - // The MSP is wider than the canonical width, so this copy narrows it. It fails on a valid - // value outside the precision and writes zeros in the null slots. + // If the MSP is wider than the canonical width, the copy checks every valid value against + // the precision and writes zeros in the null slots. match_each_signed_integer_ptype!(msp.ptype(), |P| { copy_to_canonical_width(msp.as_slice::

(), &mask, decimal_dtype) })