Add wide decimal byte-part splitting and assembly logic (#9808)

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

Author: mhk197

Cargo.lock

@@ -26,5 +26,15 @@ vortex-mask = { workspace = true }
 vortex-session = { workspace = true }
 
 [dev-dependencies]
+divan = { workspace = true }
+rand = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { path = "../../vortex-array", features = ["_test-harness"] }
+
+[[bench]]
+name = "dbp_assemble"
+harness = false
+
+[[bench]]
+name = "dbp_split"
+harness = false

encodings/decimal-byte-parts/Cargo.toml

@@ -0,0 +1,51 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Shared decimal inputs for splitting and assembly benchmarks.
+
+use rand::RngExt;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use vortex_array::arrays::DecimalArray;
+use vortex_array::dtype::DecimalDType;
+use vortex_array::dtype::DecimalType;
+use vortex_array::dtype::i256;
+use vortex_array::validity::Validity;
+use vortex_buffer::Buffer;
+use vortex_error::vortex_panic;
+
+pub(super) fn cases() -> Vec<(DecimalType, usize)> {
+    [DecimalType::I64, DecimalType::I128, DecimalType::I256]
+        .into_iter()
+        .flat_map(|values_type| [1_024, 8_192].map(|len| (values_type, len)))
+        .collect()
+}
+
+pub(super) fn decimal_array(
+    values_type: DecimalType,
+    len: usize,
+    validity: Validity,
+) -> DecimalArray {
+    let mut rng = StdRng::seed_from_u64(42);
+
+    macro_rules! decimal {
+        ($T:ty, $precision:literal) => {{
+            let max = <$T>::pow(10, $precision) - 1;
+            let values: Buffer<$T> = (0..len).map(|_| rng.random_range(-max..=max)).collect();
+            DecimalArray::new(values, DecimalDType::new($precision, 2), validity)
+        }};
+    }
+
+    match values_type {
+        DecimalType::I64 => decimal!(i64, 18),
+        DecimalType::I128 => decimal!(i128, 38),
+        DecimalType::I256 => {
+            // Keep the magnitude below 10^76 while exercising all four signed/unsigned words.
+            let values: Buffer<i256> = (0..len)
+                .map(|_| i256::from_parts(rng.random(), rng.random::<i128>() >> 4))
+                .collect();
+            DecimalArray::new(values, DecimalDType::new(76, 2), validity)
+        }
+        _ => vortex_panic!("unsupported benchmark storage type: {values_type}"),
+    }
+}

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

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Reassembling primitive decimal parts across storage widths and lengths.
+
+mod common;
+
+use divan::Bencher;
+use divan::black_box;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DecimalType;
+use vortex_array::validity::Validity;
+use vortex_decimal_byte_parts::_benchmarking::assemble_decimal;
+use vortex_decimal_byte_parts::split_decimal;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+
+use crate::common::cases;
+use crate::common::decimal_array;
+
+fn main() {
+    divan::main();
+}
+
+#[divan::bench(args = cases())]
+fn dbp_assemble(bencher: Bencher, (values_type, len): (DecimalType, usize)) {
+    let decimal = decimal_array(values_type, len, Validity::NonNullable);
+    let mut ctx = array_session().create_execution_ctx();
+    let parts = split_decimal(&decimal, &mut ctx).vortex_expect("split benchmark input");
+    let msp = parts
+        .msp
+        .execute::<PrimitiveArray>(&mut ctx)
+        .vortex_expect("execute benchmark MSP");
+    let lower_parts = parts
+        .lower_parts
+        .into_iter()
+        .map(|part| part.execute::<PrimitiveArray>(&mut ctx))
+        .collect::<VortexResult<Vec<_>>>()
+        .vortex_expect("execute benchmark lower parts");
+    let decimal_dtype = decimal.decimal_dtype();
+
+    bencher.bench(|| {
+        assemble_decimal(black_box(&msp), black_box(&lower_parts), decimal_dtype)
+            .vortex_expect("assemble decimal byte parts")
+    });
+}

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

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Splitting decimal arrays across storage widths, lengths, and validity paths.
+
+mod common;
+
+use divan::Bencher;
+use divan::black_box;
+use rand::RngExt;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::dtype::DecimalType;
+use vortex_array::validity::Validity;
+use vortex_decimal_byte_parts::split_decimal;
+use vortex_error::VortexExpect;
+
+use crate::common::cases;
+use crate::common::decimal_array;
+
+fn main() {
+    divan::main();
+}
+
+#[divan::bench(args = cases())]
+fn dbp_split_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) {
+    bench_split(bencher, values_type, len, Validity::AllValid);
+}
+
+#[divan::bench(args = cases())]
+fn dbp_split_mixed_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) {
+    let mut rng = StdRng::seed_from_u64(42);
+    let validity = Validity::from_iter((0..len).map(|_| rng.random_bool(0.5)));
+    bench_split(bencher, values_type, len, validity);
+}
+
+fn bench_split(bencher: Bencher, values_type: DecimalType, len: usize, validity: Validity) {
+    let decimal = decimal_array(values_type, len, validity);
+    let session = array_session();
+    bencher
+        .with_inputs(|| session.create_execution_ctx())
+        .bench_refs(|ctx| {
+            split_decimal(black_box(&decimal), ctx).vortex_expect("split decimal array")
+        });
+}