Share decimal value assembly and validate lower-part dtypes

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs

@@ -252,7 +252,13 @@ pub fn assemble_decimal(
     let len = msp.len();
     let lower: Vec<&[u64]> = lower_parts
         .iter()
-        .map(|part| {
+        .enumerate()
+        .map(|(idx, part)| {
+            vortex_ensure!(
+                part.dtype() == &LOWER_PART_DTYPE,
+                "lower part {idx} must have dtype {LOWER_PART_DTYPE}, got {}",
+                part.dtype()
+            );
             let part = part.as_slice::<u64>();
             vortex_ensure!(
                 part.len() == len,
@@ -265,17 +271,17 @@ pub fn assemble_decimal(
 
     Ok(match lower.as_slice() {
         [first] => DecimalArray::new(
-            assemble_wide::<i128, 1>(msp, [first]),
+            assemble_wide_decimal::<i128, 1>(msp, [first]),
             decimal_dtype,
             validity,
         ),
         [first, second] => DecimalArray::new(
-            assemble_wide::<i256, 2>(msp, [first, second]),
+            assemble_wide_decimal::<i256, 2>(msp, [first, second]),
             decimal_dtype,
             validity,
         ),
         [first, second, third] => DecimalArray::new(
-            assemble_wide::<i256, 3>(msp, [first, second, third]),
+            assemble_wide_decimal::<i256, 3>(msp, [first, second, third]),
             decimal_dtype,
             validity,
         ),
@@ -286,39 +292,10 @@ pub fn assemble_decimal(
     })
 }
 
-/// Combine a single row's parts into an `i128`.
-#[inline]
-pub(crate) fn combine_i128(msp: i64, lower: impl IntoIterator<Item = u64>) -> i128 {
-    lower.into_iter().fold(i128::from(msp), |acc, part| {
-        (acc << LOWER_PART_BITS) | i128::from(part)
-    })
-}
-
-/// Combine a signed MSP and two or three lower parts into an `i256`.
-#[inline]
-pub(crate) fn combine_i256(msp: i64, lower: impl ExactSizeIterator<Item = u64>) -> i256 {
-    let count = lower.len();
-    let mut high = i128::from(msp);
-    let mut low = 0u128;
-    for (index, part) in lower.enumerate() {
-        if count == 3 && index == 0 {
-            high = (high << LOWER_PART_BITS) | i128::from(part);
-        } else {
-            low = (low << LOWER_PART_BITS) | u128::from(part);
-        }
-    }
-    i256::from_parts(low, high)
-}
-
-/// Reassemble a signed MSP and `K` unsigned lower parts into wide integers.
+/// Assemble a column of wide decimal values from the MSP and `K` lower-part columns.
 ///
-/// Each row starts with the MSP sign-extended to `T`. Appending a lower word shifts the
-/// accumulated value left by 64 bits and fills the low bits with that word. Lower parts
-/// are appended most significant first.
-///
-/// The callers select `i128` for one lower part and `i256` for two or three. Since `K`
-/// is constant, the compiler can unroll the loop that appends the lower words.
-fn assemble_wide<T, const K: usize>(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer<T>
+/// A fixed part count lets the compiler unroll each call to [`assemble_wide_decimal_value`].
+fn assemble_wide_decimal<T, const K: usize>(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer<T>
 where
     T: NativeDecimalType + Shl<usize, Output = T> + BitOr<Output = T>,
 {
@@ -329,16 +306,30 @@ where
                 clippy::useless_conversion,
                 reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match"
             )]
-            let mut value = T::from(i64::from(*value)).vortex_expect("MSP fits in the output type");
-            for part in lower {
-                value = (value << LOWER_PART_BITS)
-                    | T::from(part[row]).vortex_expect("lower word fits in the output type");
-            }
-            value
+            let msp = i64::from(*value);
+            assemble_wide_decimal_value(msp, lower.map(|part| part[row]))
         }));
     });
     out.freeze()
 }
 
+/// Reassemble a decimal's unscaled integer from its signed MSP and `K` lower words.
+///
+/// Sign-extend the MSP to `T`, then append each lower word by shifting left 64 bits and
+/// filling the low bits. Lower words are ordered most significant first. Callers select
+/// `i128` for one lower word and `i256` for two or three.
+#[inline]
+pub(crate) fn assemble_wide_decimal_value<T, const K: usize>(msp: i64, lower: [u64; K]) -> T
+where
+    T: NativeDecimalType + Shl<usize, Output = T> + BitOr<Output = T>,
+{
+    let mut value = T::from(msp).vortex_expect("MSP fits in the output type");
+    for part in lower {
+        value = (value << LOWER_PART_BITS)
+            | T::from(part).vortex_expect("lower word fits in the output type");
+    }
+    value
+}
+
 #[cfg(test)]
 mod tests;

encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs

@@ -181,6 +181,23 @@ fn test_split_i256_part_order(
     Ok(())
 }
 
+#[rstest]
+#[case::signed(PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable))]
+#[case::narrow_unsigned(PrimitiveArray::new(buffer![0u32; 2], Validity::NonNullable))]
+#[case::nullable_all_valid(PrimitiveArray::new(buffer![0u64; 2], Validity::AllValid))]
+#[case::nullable_all_null(PrimitiveArray::new(buffer![0u64; 2], Validity::AllInvalid))]
+#[case::nullable_mixed(PrimitiveArray::new(buffer![0u64; 2], Validity::from_iter([true, false])))]
+fn test_assemble_rejects_invalid_lower_dtype(
+    #[case] invalid_lower: PrimitiveArray,
+    #[values(1, 2, 3)] lower_count: usize,
+) {
+    let msp = PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable);
+    let mut lower = vec![PrimitiveArray::new(buffer![0u64; 2], Validity::NonNullable); lower_count];
+    lower[lower_count - 1] = invalid_lower;
+    let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0);
+    assert!(assemble_decimal(&msp, &lower, dtype).is_err());
+}
+
 #[rstest]
 fn test_assemble_rejects_mismatched_lower_lengths(
     #[values(1, 2, 3)] lower_count: usize,

encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs

@@ -57,8 +57,7 @@ use vortex_session::registry::CachedId;
 
 use crate::decimal_byte_parts::limbs::LOWER_PART_DTYPE;
 use crate::decimal_byte_parts::limbs::assemble_decimal;
-use crate::decimal_byte_parts::limbs::combine_i128;
-use crate::decimal_byte_parts::limbs::combine_i256;
+use crate::decimal_byte_parts::limbs::assemble_wide_decimal_value;
 use crate::decimal_byte_parts::rules::PARENT_RULES;
 
 /// A [`DecimalByteParts`]-encoded Vortex array.
@@ -440,10 +439,19 @@ impl OperationsVTable<DecimalByteParts> for DecimalByteParts {
             })
             .collect::<VortexResult<Vec<_>>>()?;
 
-        let value = match lower_parts.len() {
-            0 => DecimalValue::I64(msp),
-            1 => DecimalValue::I128(combine_i128(msp, lower_parts)),
-            _ => DecimalValue::I256(combine_i256(msp, lower_parts.into_iter())),
+        let value = match lower_parts.as_slice() {
+            [] => DecimalValue::I64(msp),
+            [first] => DecimalValue::I128(assemble_wide_decimal_value(msp, [*first])),
+            [first, second] => {
+                DecimalValue::I256(assemble_wide_decimal_value(msp, [*first, *second]))
+            }
+            [first, second, third] => {
+                DecimalValue::I256(assemble_wide_decimal_value(msp, [*first, *second, *third]))
+            }
+            _ => vortex_bail!(
+                "at most {MAX_LOWER_PARTS} lower parts are supported, got {}",
+                lower_parts.len()
+            ),
         };
 
         Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value)))