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..2a2c8297e10 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::builtins::ArrayBuiltins; +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,144 +24,156 @@ 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 vortex_error::vortex_panic; +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. +/// Reassemble decimal parts into the canonical decimal array of `decimal_dtype`. /// -/// 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`. +/// The result always has the canonical width of `decimal_dtype`: /// -/// 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. +/// - 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], 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()?; + 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, + ) + }; - 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) + 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 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) + }) } -/// 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`. +/// 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)?; - // 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| { + 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. @@ -204,6 +219,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; @@ -211,6 +227,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; @@ -221,6 +239,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; @@ -374,21 +393,177 @@ 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(()) } #[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..54ef5d35ec0 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. Decoding always produces a decimal array at the canonical 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..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,43 +43,31 @@ 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_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::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(), 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 +123,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 +190,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; 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::*;