Support wide decimals in DecimalByteParts with 64-bit lower parts

.gitignore

@@ -52,6 +52,8 @@ coverage.xml
 *.cover
 *.py,cover
 .hypothesis/
+# hegeltest's example database, the Rust equivalent of .hypothesis/
+.hegel/
 .pytest_cache/
 cover/
 

Cargo.toml

@@ -214,6 +214,7 @@ pyo3-bytes = "0.7"
 pyo3-log = "0.13.0"
 pyo3-object_store = "0.11.0"
 quote = "1.0.44"
+hegeltest = "0.28.7"
 rand = "0.10.1"
 rand_distr = "0.6"
 ratatui = { version = "0.30", default-features = false }

encodings/decimal-byte-parts/Cargo.toml

@@ -16,6 +16,11 @@ version = { workspace = true }
 [lints]
 workspace = true
 
+[features]
+# Serializing more than one child, which readers that predate lower parts cannot open.
+# Building and reading them is always allowed; only writing is gated.
+unstable_encodings = []
+
 [dependencies]
 num-traits = { workspace = true }
 prost = { workspace = true }
@@ -26,5 +31,12 @@ vortex-mask = { workspace = true }
 vortex-session = { workspace = true }
 
 [dev-dependencies]
+divan = { workspace = true }
+hegeltest = { workspace = true }
+rand = { workspace = true }
 rstest = { workspace = true }
 vortex-array = { path = "../../vortex-array", features = ["_test-harness"] }
+
+[[bench]]
+name = "decimal_assemble"
+harness = false

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

@@ -0,0 +1,106 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Canonicalizing `DecimalByteParts` into `i128`/`i256` values.
+//!
+//! Reassembly walks a most significant part plus one (`i128`) or three (`i256`) unsigned
+//! 64-bit lower parts and produces one wide value per row. This benchmark measures the
+//! shipped path through the array API, so it tracks whatever shape the crate currently uses
+//! and cannot drift away from it.
+//!
+//! Alternative shapes were compared while choosing that path and then removed, since keeping
+//! hand-written copies of the assembly loop here means maintaining the same loop twice. At
+//! 65,536 rows, `fastest` of three runs:
+//!
+//! - **`i256` is dominated by the part count being visible to the compiler.** Specializing it
+//!   to a constant is 1.85x (190 µs against 351 µs). How the output is written barely matters
+//!   at 32 bytes per row.
+//! - **`i128` is dominated by the write.** Specializing the part count is worth only ~1.04x,
+//!   while storing into a pre-sized buffer instead of pushing into a reserved one is 1.6x
+//!   (83 µs against 138 µs) — the bounds-checked `push` is the whole cost at 16 bytes per row.
+//! - **Columnar always loses.** For `i256` each lane store is strided by 32 bytes, 2.3x slower
+//!   than the row loop (438 µs); cache blocking the passes recovered part of that and was
+//!   still 1.6x slower; expressing them as whole-value `i256` shifts was 11x slower. For
+//!   `i128` the two-pass column shape (103 µs) beats the *pushing* row loop but still loses to
+//!   the single-pass write, so the second pass buys nothing once the push is gone.
+//! - **Hand-written 64-bit words do not beat the `u128` packing.** `i256::from_parts` takes a
+//!   `u128` and an `i128`, so each row ends in `u128::from(w0) | (u128::from(w1) << 64)`.
+//!   Writing four `u64` lanes by hand instead ties it. Disassembly says why: neither emits a
+//!   single `shld`/`shrd`, and both compile to four plain 64-bit stores per row at offsets
+//!   0x0/0x8/0x10/0x18. The `i128` loop is the same — `(i128::from(msp) << 64) | i128::from(p)`
+//!   becomes two 64-bit stores. A shift by a constant multiple of 64 followed by an or is pure
+//!   data movement and LLVM recognizes it; the 128-bit codegen worth avoiding is division and
+//!   remainder, which call into compiler-rt, and shifts by a runtime amount. Neither is here.
+//!
+//! The removed variants are recoverable from git history if a future change needs to re-run
+//! the comparison rather than trust these numbers.
+
+#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)]
+
+use divan::Bencher;
+use divan::black_box;
+use rand::RngExt;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DecimalDType;
+use vortex_array::validity::Validity;
+use vortex_buffer::Buffer;
+use vortex_decimal_byte_parts::DecimalByteParts;
+
+fn main() {
+    divan::main();
+}
+
+/// Rows per benchmark: a typical scan chunk, and large enough that the output does not fit
+/// in L2.
+const LEN: usize = 65_536;
+
+/// Deterministic pseudo-random words, so no part is constant or a sequence.
+fn words(seed: u64, len: usize) -> Buffer<u64> {
+    let mut rng = StdRng::seed_from_u64(seed);
+    (0..len).map(|_| rng.random()).collect()
+}
+
+fn msp(seed: u64, len: usize) -> Buffer<i64> {
+    words(seed, len)
+        .iter()
+        .map(|w| (w >> 40).cast_signed())
+        .collect()
+}
+
+/// Canonicalizing through the public array API, so the child execution and validity handling
+/// around the assembly loop are included.
+#[divan::bench(args = [1, 3])]
+fn canonicalize_byte_parts(bencher: Bencher, lower_parts: usize) {
+    let msp = PrimitiveArray::new(msp(1, LEN), Validity::NonNullable);
+    let lower = (0..lower_parts)
+        .map(|i| PrimitiveArray::new(words(7 + i as u64, LEN), Validity::NonNullable))
+        .collect::<Vec<_>>();
+
+    let dtype = if lower_parts == 1 {
+        DecimalDType::new(38, 2)
+    } else {
+        DecimalDType::new(76, 2)
+    };
+    let array = DecimalByteParts::try_new_with_lower_parts(
+        msp.into_array(),
+        lower.into_iter().map(IntoArray::into_array).collect(),
+        dtype,
+    )
+    .unwrap()
+    .into_array();
+
+    let session = array_session();
+    bencher
+        .with_inputs(|| session.create_execution_ctx())
+        .bench_refs(|ctx| {
+            black_box(array.clone())
+                .execute::<DecimalArray>(ctx)
+                .unwrap()
+        });
+}

encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs

@@ -11,6 +11,7 @@ use vortex_error::VortexResult;
 
 use crate::DecimalByteParts;
 use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt;
+use crate::decimal_byte_parts::with_msp;
 
 impl CastReduce for DecimalByteParts {
     fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
@@ -29,9 +30,7 @@ impl CastReduce for DecimalByteParts {
             .msp()
             .cast(array.msp().dtype().with_nullability(*target_nullability))?;
 
-        Ok(Some(
-            DecimalByteParts::try_new(new_msp, *target_decimal)?.into_array(),
-        ))
+        with_msp(array, new_msp, *target_decimal).map(|a| Some(a.into_array()))
     }
 }
 
@@ -49,10 +48,14 @@ mod tests {
     use vortex_array::dtype::DType;
     use vortex_array::dtype::DecimalDType;
     use vortex_array::dtype::Nullability;
+    use vortex_array::validity::Validity;
     use vortex_buffer::buffer;
 
     use crate::DecimalByteParts;
     use crate::DecimalBytePartsArray;
+    use crate::decimal_byte_parts::testing::i128_parts;
+    use crate::decimal_byte_parts::testing::i256_of;
+    use crate::decimal_byte_parts::testing::i256_parts;
 
     #[test]
     fn test_cast_decimal_byte_parts_nullability() {
@@ -117,6 +120,14 @@ mod tests {
         buffer![-100i32, -200, 300, -400, 500].into_array(),
         DecimalDType::new(10, 2),
     ).unwrap())]
+    #[case::one_lower_part(i128_parts(
+        vec![1i128 << 70, -(1i128 << 70), 5, (1i128 << 64) - 1, 0],
+        Validity::NonNullable,
+    ))]
+    #[case::three_lower_parts(i256_parts(
+        vec![i256_of(1, 0), i256_of(-1, 5), i256_of(0, u128::MAX)],
+        Validity::NonNullable,
+    ))]
     fn test_cast_decimal_byte_parts_conformance(#[case] array: DecimalBytePartsArray) {
         test_cast_conformance(
             &array.into_array(),

encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs

@@ -39,6 +39,12 @@ impl CompareKernel for DecimalByteParts {
             return Ok(None);
         };
 
+        // The MSP alone only determines the ordering when it holds the whole value. With
+        // lower parts present, fall back to comparing the canonical decimal.
+        if !lhs.lower_parts().is_empty() {
+            return Ok(None);
+        }
+
         let nullability = lhs.dtype().nullability() | rhs.dtype().nullability();
         let scalar_type = lhs.msp().dtype().with_nullability(nullability);
 
@@ -158,10 +164,12 @@ mod tests {
     use vortex_array::scalar_fn::fns::operators::Operator;
     use vortex_array::validity::Validity;
     use vortex_buffer::buffer;
+    use vortex_error::VortexExpect;
     use vortex_error::VortexResult;
     use vortex_session::VortexSession;
 
     use crate::DecimalByteParts;
+    use crate::decimal_byte_parts::testing::i128_parts;
 
     static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
         let session = vortex_array::array_session();
@@ -220,6 +228,45 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn compare_decimal_const_with_lower_parts() -> VortexResult<()> {
+        // The MSP-only pushdown is invalid once lower parts carry part of the value, so this
+        // must fall back to the canonical comparison rather than compare MSPs.
+        let values = vec![1i128 << 70, (1i128 << 70) + 1, 5, -(1i128 << 70)];
+        let lhs = i128_parts(values.clone(), Validity::NonNullable).into_array();
+        let decimal_dtype = *lhs
+            .dtype()
+            .as_decimal_opt()
+            .vortex_expect("decimal byte parts array");
+
+        let pivot = (1i128 << 70) + 1;
+        let rhs = ConstantArray::new(
+            Scalar::decimal(
+                DecimalValue::I128(pivot),
+                decimal_dtype,
+                Nullability::NonNullable,
+            ),
+            lhs.len(),
+        )
+        .into_array();
+
+        let mut ctx = SESSION.create_execution_ctx();
+        for (operator, predicate) in [
+            (Operator::Eq, (|v, p| v == p) as fn(i128, i128) -> bool),
+            (Operator::NotEq, |v, p| v != p),
+            (Operator::Lt, |v, p| v < p),
+            (Operator::Lte, |v, p| v <= p),
+            (Operator::Gt, |v, p| v > p),
+            (Operator::Gte, |v, p| v >= p),
+        ] {
+            let res = lhs.clone().binary(rhs.clone(), operator)?;
+            let expected =
+                BoolArray::from_iter(values.iter().map(|v| predicate(*v, pivot))).into_array();
+            assert_arrays_eq!(res, expected, &mut ctx);
+        }
+        Ok(())
+    }
+
     #[test]
     fn compare_decimal_const_unconvertible_comparison() {
         let decimal_dtype = DecimalDType::new(40, 2);