Build DecimalByteParts from the canonical decimal array

encodings/decimal-byte-parts/benches/common/mod.rs

@@ -88,4 +88,15 @@ pub(super) mod arrays {
             _ => vortex_panic!("unsupported benchmark storage type: {values_type}"),
         }
     }
+
+    /// The number of lower parts that holds every value of `values_type`, so the split and
+    /// assembly benchmarks keep the shape of each storage width.
+    pub(crate) fn full_width_lower_part_count(values_type: DecimalType) -> usize {
+        match values_type {
+            DecimalType::I64 => 0,
+            DecimalType::I128 => 1,
+            DecimalType::I256 => 3,
+            _ => vortex_panic!("unsupported benchmark storage type: {values_type}"),
+        }
+    }
 }

encodings/decimal-byte-parts/benches/dbp_assemble.rs

@@ -126,6 +126,7 @@ mod arrays {
 
     use crate::common::arrays::cases;
     use crate::common::arrays::decimal_array;
+    use crate::common::arrays::full_width_lower_part_count;
 
     pub(super) static SESSION: LazyLock<VortexSession> =
         LazyLock::new(|| VortexSession::empty().with::<KernelSession>());
@@ -134,7 +135,8 @@ mod arrays {
     fn dbp_assemble(bencher: Bencher, (values_type, len): (DecimalType, usize)) {
         let decimal = decimal_array(values_type, len, Validity::NonNullable);
         let mut ctx = SESSION.create_execution_ctx();
-        let parts = split_decimal(&decimal, &mut ctx).vortex_expect("split benchmark input");
+        let parts = split_decimal(&decimal, full_width_lower_part_count(values_type), &mut ctx)
+            .vortex_expect("split benchmark input");
         bench_parts(
             bencher,
             parts.msp,

encodings/decimal-byte-parts/benches/dbp_split.rs

@@ -94,6 +94,7 @@ mod arrays {
 
     use crate::common::arrays::cases;
     use crate::common::arrays::decimal_array;
+    use crate::common::arrays::full_width_lower_part_count;
 
     pub(super) static SESSION: LazyLock<VortexSession> =
         LazyLock::new(|| VortexSession::empty().with::<KernelSession>());
@@ -118,11 +119,13 @@ mod arrays {
     fn bench_split(bencher: Bencher, values_type: DecimalType, len: usize, validity: Validity) {
         let decimal = decimal_array(values_type, len, validity);
         mark_optimized(decimal.as_array(), &SESSION).vortex_expect("mark the decimal as optimized");
+        let lower_part_count = full_width_lower_part_count(values_type);
 
         bencher
             .with_inputs(|| SESSION.create_execution_ctx())
             .bench_refs(|ctx| {
-                split_decimal(black_box(&decimal), ctx).vortex_expect("split decimal array")
+                split_decimal(black_box(&decimal), lower_part_count, ctx)
+                    .vortex_expect("split decimal array")
             });
     }
 }

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

@@ -41,6 +41,7 @@ use vortex_session::VortexSession;
 
 use super::assemble_decimal;
 use crate::decimal_byte_parts::split_decimal;
+use crate::decimal_byte_parts::testing::full_width_lower_part_count;
 
 static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
 
@@ -129,7 +130,8 @@ fn assert_form(decoded: &ArrayRef, decimal_dtype: DecimalDType, child_precision:
 
 fn round_trip(decimal: DecimalArray) -> VortexResult<DecimalArray> {
     let mut ctx = SESSION.create_execution_ctx();
-    let parts = split_decimal(&decimal, &mut ctx)?;
+    let lower_part_count = full_width_lower_part_count(decimal.values_type());
+    let parts = split_decimal(&decimal, lower_part_count, &mut ctx)?;
     assemble_decimal(
         &parts.msp,
         &parts.lower_parts,
@@ -166,7 +168,7 @@ fn test_split_without_valid_rows(
             validity,
         )
     });
-    let parts = split_decimal(&decimal, &mut ctx)?;
+    let parts = split_decimal(&decimal, full_width_lower_part_count(values_type), &mut ctx)?;
     assert!(parts.msp.is::<Constant>());
     assert!(parts.lower_parts.iter().all(|part| part.is::<Constant>()));
     assert_eq!(parts.msp.len(), len);
@@ -213,8 +215,9 @@ fn test_split_zeroes_null_words(
             .collect::<Buffer<_>>(),
         Validity::NonNullable,
     );
-    let parts = split_decimal(&decimal, &mut ctx)?;
-    assert_eq!(parts.lower_parts.len(), if wide_256 { 3 } else { 1 });
+    let lower_part_count = if wide_256 { 3 } else { 1 };
+    let parts = split_decimal(&decimal, lower_part_count, &mut ctx)?;
+    assert_eq!(parts.lower_parts.len(), lower_part_count);
     assert_eq!(
         parts.msp.dtype(),
         &DType::Primitive(PType::I64, decimal.dtype().nullability())
@@ -265,7 +268,7 @@ fn test_split_narrow_decimal_reuses_values(
     let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(9, 0), validity);
     let values_ptr = decimal.buffer::<i32>().as_ptr();
 
-    let parts = split_decimal(&decimal, &mut ctx)?;
+    let parts = split_decimal(&decimal, 0, &mut ctx)?;
     assert!(parts.lower_parts.is_empty());
     assert_eq!(parts.msp.dtype().as_ptype(), PType::I32);
     let msp = parts.msp.clone().execute::<PrimitiveArray>(&mut ctx)?;

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

@@ -21,23 +21,38 @@ use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
 
+/// The maximum number of 64-bit lower parts an encoded `i128` decimal can carry.
+const MAX_I128_LOWER_PARTS: usize = 1;
+
+/// The maximum number of 64-bit lower parts an encoded `i256` decimal can carry.
+const MAX_I256_LOWER_PARTS: usize = 3;
+
+/// The maximum number of 64-bit lower parts an encoded decimal can carry.
+const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS;
+
+/// Number of bits stored in each lower part.
+const LOWER_PART_BITS: usize = 64;
+
+/// The dtype of every lower part.
+const LOWER_PART_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable);
+
 mod array;
+pub use array::*;
+
 mod assemble;
 pub(crate) mod compute;
-mod plugin;
-#[cfg(test)]
-mod prop_tests;
-mod rules;
-mod split;
-#[cfg(test)]
-mod testing;
 
-pub use array::*;
+mod plugin;
 pub use plugin::DecimalBytePartsPlugin;
 pub use plugin::DecimalBytePartsV2Metadata;
 pub use plugin::decimal_byte_parts_v1_id;
 pub use plugin::decimal_byte_parts_v2_id;
+
+mod rules;
+
+mod split;
 pub use split::DecimalParts;
+pub use split::lower_part_count;
 pub use split::split_decimal;
 
 #[doc(hidden)]
@@ -49,17 +64,7 @@ pub mod _benchmarking {
     pub use super::split::split_wide;
 }
 
-/// The maximum number of 64-bit lower parts an encoded `i128` decimal can carry.
-const MAX_I128_LOWER_PARTS: usize = 1;
-
-/// The maximum number of 64-bit lower parts an encoded `i256` decimal can carry.
-const MAX_I256_LOWER_PARTS: usize = 3;
-
-/// The maximum number of 64-bit lower parts an encoded decimal can carry.
-const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS;
-
-/// Number of bits stored in each lower part.
-const LOWER_PART_BITS: usize = 64;
-
-/// The dtype of every lower part.
-const LOWER_PART_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable);
+#[cfg(test)]
+mod prop_tests;
+#[cfg(test)]
+mod testing;

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

@@ -18,6 +18,7 @@ use vortex_array::dtype::DType;
 use vortex_array::dtype::DecimalDType;
 use vortex_array::dtype::Nullability;
 use vortex_array::dtype::PType;
+use vortex_array::dtype::i256;
 use vortex_array::serde::SerializeOptions;
 use vortex_array::serde::SerializedArray;
 use vortex_array::session::ArraySessionExt;
@@ -46,6 +47,24 @@ static SESSION: LazyLock<VortexSession> = 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::wider_i256_storage_all_null(DecimalByteParts::encode(
+    &DecimalArray::new(
+        buffer![i256::from_parts(1 << 127, 0); 3],
+        DecimalDType::new(2, 0), Validity::AllInvalid,
+    ),
+    &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()],

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

@@ -149,8 +149,9 @@ fn decoded_survives_encode_then_decode(tc: TestCase) {
 ///
 /// Starting from the encoded side reaches part counts `split_decimal` never produces, so this
 /// covers layouts the property above cannot generate. It compares decoded values rather than
-/// the arrays themselves because re-encoding normalizes the part count: splitting an `i256`
-/// always yields three lower parts, whatever the original array carried.
+/// the arrays themselves because re-encoding normalizes the part count:
+/// [`DecimalByteParts::encode`] uses the fewest lower parts that hold the valid values, whatever
+/// the original array carried.
 #[hegel::test]
 fn encoded_survives_decode_then_encode(tc: TestCase) {
     let (array, _len) = draw_encoded(&tc);