From 56ea0c091c32aa9973b548ec2bdcd9b4af1b9217 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 16:27:02 -0400 Subject: [PATCH 1/5] Support multi-part decimal arrays and kernels Represent wide decimals with a signed high part and up to three unsigned low parts. Add validation, execution, kernel support, property tests, and assembly benchmarks while keeping serialization on the frozen format. Signed-off-by: Matt Katz --- .gitignore | 2 + Cargo.lock | 84 ++- Cargo.toml | 1 + encodings/decimal-byte-parts/Cargo.toml | 1 + .../src/decimal_byte_parts/compute/cast.rs | 17 +- .../src/decimal_byte_parts/compute/compare.rs | 47 ++ .../src/decimal_byte_parts/compute/filter.rs | 44 +- .../decimal_byte_parts/compute/is_constant.rs | 27 +- .../src/decimal_byte_parts/compute/kernel.rs | 8 - .../src/decimal_byte_parts/compute/mask.rs | 16 +- .../src/decimal_byte_parts/compute/mod.rs | 33 + .../src/decimal_byte_parts/compute/take.rs | 107 +++- .../src/decimal_byte_parts/limbs/mod.rs | 28 +- .../src/decimal_byte_parts/mod.rs | 606 +++++++++++++++--- .../src/decimal_byte_parts/rules.rs | 42 +- .../src/decimal_byte_parts/slice.rs | 14 +- .../src/decimal_byte_parts/testing.rs | 53 ++ encodings/decimal-byte-parts/tests/props.rs | 198 ++++++ vortex-btrblocks/src/trace_tests.rs | 5 +- .../kernel/encodings/decimal_byte_parts.rs | 7 + 20 files changed, 1139 insertions(+), 201 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs create mode 100644 encodings/decimal-byte-parts/tests/props.rs diff --git a/.gitignore b/.gitignore index f9613807332..6db14ce5f6a 100644 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,8 @@ coverage.xml *.cover *.py,cover .hypothesis/ +# hegeltest's example database, the Rust equivalent of .hypothesis/ +.hegel/ .pytest_cache/ cover/ diff --git a/Cargo.lock b/Cargo.lock index 04baecf6504..639e06bf078 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1120,7 +1120,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f21ff1fc630079352bae9b024f85519bf1f641cf7f326623f4c0b59f7ea834fd" dependencies = [ "compact_str", - "miniz_oxide", + "miniz_oxide 0.9.1", "thiserror 2.0.20", ] @@ -2257,6 +2257,25 @@ dependencies = [ "parking_lot_core", ] +[[package]] +name = "dashu-base" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "993b95dc1b248e3f5747dcb017a41d6e75853a2e5ee4504f7d537c5b8dffdae4" + +[[package]] +name = "dashu-int" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49c05a0d5cb0b39fcc87c46432fdac24b90dce239857c7f6b798be4ffc3c42c6" +dependencies = [ + "cfg-if", + "dashu-base", + "num-modular", + "rustversion", + "static_assertions", +] + [[package]] name = "datafusion" version = "54.1.0" @@ -4098,7 +4117,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.9.1", "zlib-rs", ] @@ -4699,6 +4718,51 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hegeltest" +version = "0.28.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "100bcd6ef825f5b6a60e2f55c05bb626ebf254dd8a09d16e006c4bb7883e7f1c" +dependencies = [ + "crc32fast", + "dashu-int", + "hegeltest-c", + "hegeltest-macros", + "miniz_oxide 0.8.9", + "parking_lot", + "paste", + "rand 0.10.2", + "rustc-hash", + "tempfile", +] + +[[package]] +name = "hegeltest-c" +version = "0.30.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a672fd53360ca4122c1a145a85e8fef835508d7b40eb9de43498978e796c54b" +dependencies = [ + "dashu-int", + "hashbrown 0.17.1", + "libm", + "miniz_oxide 0.8.9", + "parking_lot", + "rand 0.10.2", + "rustc-hash", + "tempfile", +] + +[[package]] +name = "hegeltest-macros" +version = "0.28.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba792d78fa3740a7c1627085c34618b998b8aa0f63625721235234f525aad1aa" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "hermit-abi" version = "0.5.3" @@ -6555,6 +6619,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", +] + [[package]] name = "miniz_oxide" version = "0.9.1" @@ -6862,6 +6935,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-modular" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc41a1374056e9672221567958a66c16be12d0e2c1b408761e14d901c237d5e0" + [[package]] name = "num-rational" version = "0.4.2" @@ -10922,6 +11001,7 @@ name = "vortex-decimal-byte-parts" version = "0.1.0" dependencies = [ "codspeed-divan-compat", + "hegeltest", "num-traits", "prost 0.14.4", "rand 0.10.2", diff --git a/Cargo.toml b/Cargo.toml index 6a9afca5e25..1bf8177ec0e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -177,6 +177,7 @@ glob = "0.3.2" goldenfile = "1" half = { version = "2.7.1", features = ["std", "num-traits"] } hashbrown = "0.17.1" +hegeltest = "0.28.7" http = "1.5.0" humansize = "2.1.3" indicatif = "0.18.0" diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 9f2e387a4da..e9ea8569af1 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -27,6 +27,7 @@ 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"] } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index 5ae1bf0101e..7b949fcd695 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/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> { @@ -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(), diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs index 3044bd6e605..fe4d69801a3 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/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 = 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); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index a47a6ed846b..e4fb03a5ca0 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs @@ -5,22 +5,15 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::filter::FilterReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::map_parts; + impl FilterReduce for DecimalByteParts { fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { - DecimalByteParts::try_new( - array.msp().filter(mask.clone())?, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - ) - .map(|d| Some(d.into_array())) + map_parts(array, |part| part.filter(mask.clone())).map(|d| Some(d.into_array())) } } @@ -32,9 +25,13 @@ mod test { use vortex_array::arrays::PrimitiveArray; use vortex_array::compute::conformance::filter::test_filter_conformance; use vortex_array::dtype::DecimalDType; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use crate::DecimalByteParts; + 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_filter_decimal_byte_parts() { @@ -59,4 +56,31 @@ mod test { &mut array_session().create_execution_ctx(), ); } + + #[test] + fn test_filter_decimal_byte_parts_with_lower_parts() { + let array = i128_parts( + vec![1i128 << 70, -(1i128 << 70), 5, (1i128 << 64) - 1, 0], + Validity::NonNullable, + ); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + + let array = i256_parts( + vec![ + i256_of(1, 0), + i256_of(-1, 5), + i256_of(0, u128::MAX), + i256_of(1 << 64, 7), + i256_of(0, 0), + ], + Validity::from_iter([true, false, true, true, false]), + ); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs index 065bc5e0051..3fe59111f6e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::ArrayRef; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; @@ -15,7 +16,9 @@ use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; /// DecimalByteParts-specific is_constant kernel. /// -/// Delegates to checking if the MSP (most significant part) is constant. +/// Delegates to checking that every part is constant: the MSP (most significant part) plus +/// each lower part. An all-null array is constant regardless of the bits its lower parts +/// hold in null slots. #[derive(Debug)] pub(crate) struct DecimalBytePartsIsConstantKernel; @@ -34,7 +37,27 @@ impl DynAggregateKernel for DecimalBytePartsIsConstantKernel { return Ok(None); }; - let result = is_constant(array.msp(), ctx)?; + let result = is_constant_parts(array, ctx)?; Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } + +fn is_constant_parts( + array: ArrayView<'_, DecimalByteParts>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !is_constant(array.msp(), ctx)? { + return Ok(false); + } + // Null slots hold undefined bits in the lower parts, so they cannot make a constant + // (all-null) array non-constant. + if array.array().all_invalid(ctx)? { + return Ok(true); + } + for part in array.lower_parts().iter() { + if !is_constant(part, ctx)? { + return Ok(false); + } + } + Ok(true) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 5e8d28e3526..cb71ba7880c 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -1,9 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::ArrayVTable; -use vortex_array::arrays::Dict; -use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; @@ -19,9 +16,4 @@ pub(crate) fn initialize(session: &VortexSession) { DecimalByteParts, CompareExecuteAdaptor(DecimalByteParts), ); - kernels.register_execute_parent_kernel( - Dict.id(), - DecimalByteParts, - TakeExecuteAdaptor(DecimalByteParts), - ); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index e7dc95af84f..2eea785794b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs @@ -6,24 +6,18 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; use vortex_array::scalar_fn::fns::mask::MaskReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::decimal_dtype; +use crate::decimal_byte_parts::with_msp; impl MaskReduce for DecimalByteParts { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { + // Validity lives in the MSP, so only that part needs masking: the lower parts hold + // undefined bits in null slots, which is exactly what a masked-out row is. let masked_msp = MaskExpr::try_new(array.msp().clone(), mask.clone())?.into_array(); - Ok(Some( - DecimalByteParts::try_new( - masked_msp, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - )? - .into_array(), - )) + with_msp(array, masked_msp, decimal_dtype(array)).map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 6c2d0dabb31..844468545cf 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -19,10 +19,36 @@ mod tests { use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::i256; + 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; + + /// Values needing more than 64 bits, so the encoding carries lower parts. + fn wide_i128() -> Vec { + vec![ + 1 << 70, + -(1 << 70), + (1 << 64) - 1, + 0, + 99_999_999_999_999_999_999_999_999_999_999_999_999, + ] + } + + fn wide_i256() -> Vec { + vec![ + i256_of(1, 0), + i256_of(-1, 0), + i256_of(0, u128::MAX), + i256_of(1 << 64, 7), + i256_of(0, 0), + ] + } #[rstest] // Basic decimal byte parts arrays @@ -70,6 +96,11 @@ mod tests { PrimitiveArray::from_iter((0..2000i64).map(|i| i * 1000000)).into_array(), DecimalDType::new(19, 6) ).unwrap())] + // Wide decimals carrying lower parts + #[case::decimal_i128_one_lower_part(i128_parts(wide_i128(), Validity::NonNullable))] + #[case::decimal_i128_nullable(i128_parts(wide_i128(), Validity::from_iter([true, false, true, true, false])))] + #[case::decimal_i256_three_lower_parts(i256_parts(wide_i256(), Validity::NonNullable))] + #[case::decimal_i256_nullable(i256_parts(wide_i256(), Validity::from_iter([false, true, true, false, true])))] fn test_decimal_byte_parts_consistency(#[case] array: DecimalBytePartsArray) { let ctx = &mut array_session().create_execution_ctx(); @@ -89,6 +120,8 @@ mod tests { buffer![-100i32, -200, 300, -400, 500].into_array(), DecimalDType::new(10, 2) ).unwrap())] + #[case::decimal_i128_one_lower_part(i128_parts(wide_i128(), Validity::NonNullable))] + #[case::decimal_i256_three_lower_parts(i256_parts(wide_i256(), Validity::NonNullable))] fn test_decimal_byte_parts_binary_numeric(#[case] array: DecimalBytePartsArray) { let ctx = &mut array_session().create_execution_ctx(); test_binary_numeric_array(&array.into_array(), ctx); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index 7a18f7bf91b..578834635b8 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs @@ -3,28 +3,101 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; -use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::arrays::dict::TakeExecute; -use vortex_error::VortexExpect; +use vortex_array::arrays::dict::TakeReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::map_parts; -impl TakeExecute for DecimalByteParts { - fn take( - array: ArrayView<'_, Self>, - indices: &ArrayRef, - _ctx: &mut ExecutionCtx, - ) -> VortexResult> { - DecimalByteParts::try_new( - array.msp().take(indices.clone())?, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - ) - .map(|a| Some(a.into_array())) +impl TakeReduce for DecimalByteParts { + /// Taking wraps each part in a `Dict` without reading any buffer, so it reduces rather + /// than executes. + fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + // Taking with nullable indices makes every taken part nullable, but lower parts must + // stay non-nullable `u64` — validity belongs to the MSP alone. Fall back to the + // canonical path rather than rebuilding parts we would have to strip nullability from. + if indices.dtype().is_nullable() && !array.lower_parts().is_empty() { + return Ok(None); + } + + map_parts(array, |part| part.take(indices.clone())).map(|a| Some(a.into_array())) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + 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::assert_arrays_eq; + use vortex_array::dtype::DecimalDType; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::encode; + use crate::decimal_byte_parts::testing::i256_of; + + /// Taking pushes down into the parts during optimization, with no execution context in + /// play: `ArrayRef::take` wraps the array in a `Dict` and optimizes, and the reduce rule + /// must rewrite that into a `DecimalByteParts` of taken parts. + #[test] + fn take_pushes_down_without_executing() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + + let decimal = DecimalArray::new( + Buffer::from(vec![1i128 << 70, 2, 3]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let indices = buffer![0u64, 2].into_array(); + let taken = encode(&decimal)?.into_array().take(indices)?; + + assert!( + taken.is::(), + "expected the take to reduce into the encoding, got {}", + taken.encoding_id() + ); + Ok(()) + } + + /// Taking with nullable indices must still round-trip the wide values, including the + /// null row, on arrays that carry lower parts. + #[rstest] + #[case::one_lower_part(DecimalArray::new( + Buffer::from(vec![1i128 << 70, 2, 3]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ))] + #[case::three_lower_parts(DecimalArray::new( + Buffer::from(vec![i256_of(1, 1 << 70), i256_of(0, 2), i256_of(0, 3)]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ))] + fn take_with_nullable_indices(#[case] decimal: DecimalArray) -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let indices = PrimitiveArray::from_option_iter([Some(0u64), None, Some(2u64)]).into_array(); + let expected = decimal + .clone() + .into_array() + .take(indices.clone())? + .execute::(&mut ctx)?; + + let taken = encode(&decimal)?.into_array().take(indices)?; + let actual = taken.execute::(&mut ctx)?; + + assert_arrays_eq!(expected, actual, &mut ctx); + Ok(()) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs index 1e561b149fe..119e0a63854 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -253,10 +253,6 @@ pub fn assemble_decimal( let lower: Vec<&[u64]> = lower_parts .iter() .map(|part| { - vortex_ensure!( - part.dtype() == &LOWER_PART_DTYPE, - "lower part must be non-nullable u64" - ); let part = part.as_slice::(); vortex_ensure!( part.len() == len, @@ -290,6 +286,30 @@ pub fn assemble_decimal( }) } +/// Combine a single row's parts into an `i128`. +#[inline] +pub(crate) fn combine_i128(msp: i64, lower: impl IntoIterator) -> 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) -> 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. /// /// Each row starts with the MSP sign-extended to `T`. Appending a lower word shifts the 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 a7f63bfa082..4a104e82c2b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -10,12 +10,14 @@ use vortex_array::ArrayParts; use vortex_array::ArrayView; pub(crate) mod compute; mod limbs; -pub use limbs::DecimalParts; -pub use limbs::MAX_LOWER_PARTS; -pub use limbs::split_decimal; mod rules; mod slice; +#[cfg(test)] +pub(crate) mod testing; +pub use limbs::DecimalParts; +pub use limbs::MAX_LOWER_PARTS; +pub use limbs::split_decimal; #[doc(hidden)] pub mod _benchmarking { pub use super::limbs::assemble_decimal; @@ -26,23 +28,22 @@ use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayRef; +use vortex_array::ArraySlots; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; use vortex_array::IntoArray; use vortex_array::array_slots; -use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; use vortex_array::dtype::PType; -use vortex_array::match_each_signed_integer_ptype; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::serde::ArrayChildren; -use vortex_array::smallvec::smallvec; use vortex_array::vtable::OperationsVTable; use vortex_array::vtable::VTable; use vortex_array::vtable::ValidityChild; @@ -51,10 +52,15 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; 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::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. @@ -78,6 +84,70 @@ pub struct DecimalBytesPartsMetadata { lower_part_count: u32, } +impl DecimalBytesPartsMetadata { + fn from_array(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { + Ok(Self { + zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, + lower_part_count: u32::try_from(array.lower_parts().len()) + .map_err(|_| vortex_err!("lower part count exceeds u32"))?, + }) + } + + fn into_array_parts( + self, + dtype: &DType, + len: usize, + children: &dyn ArrayChildren, + ) -> VortexResult> { + vortex_ensure!( + dtype.as_decimal_opt().is_some(), + "decoding decimal but given non decimal dtype {dtype}" + ); + + let encoded_dtype = DType::Primitive(self.zeroth_child_ptype(), dtype.nullability()); + + let lower_part_count = self.lower_part_count()?; + vortex_ensure!( + children.len() == DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, + "expected {} children, got {}", + DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, + children.len() + ); + + let msp = children.get(DecimalBytePartsSlots::MSP, &encoded_dtype, len)?; + + let mut slots = ArraySlots::with_capacity(children.len()); + slots.push(Some(msp)); + for idx in 0..lower_part_count { + slots.push(Some(children.get( + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx, + &LOWER_PART_DTYPE, + len, + )?)); + } + + Ok( + ArrayParts::new(DecimalByteParts, dtype.clone(), len, DecimalBytePartsData) + .with_slots(slots), + ) + } + + /// The number of lower parts encoded in this array. + /// + /// # Errors + /// + /// Returns an error if the count exceeds [`MAX_LOWER_PARTS`]. + fn lower_part_count(&self) -> VortexResult { + let count = usize::try_from(self.lower_part_count) + .map_err(|_| vortex_err!("lower part count {} out of range", self.lower_part_count))?; + vortex_ensure!( + count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {count}" + ); + Ok(count) + } +} + impl VTable for DecimalByteParts { type TypedArrayData = DecimalBytePartsData; @@ -99,8 +169,14 @@ impl VTable for DecimalByteParts { let Some(decimal_dtype) = dtype.as_decimal_opt() else { vortex_bail!("expected decimal dtype, got {}", dtype) }; - let msp = DecimalBytePartsSlotsView::from_slots(slots).msp; - DecimalBytePartsData::validate(msp, *decimal_dtype, dtype, len) + let slots = DecimalBytePartsSlotsView::from_slots(slots); + DecimalBytePartsData::validate( + slots.msp, + slots.lower_parts.iter(), + *decimal_dtype, + dtype, + len, + ) } fn nbuffers(_array: ArrayView<'_, Self>) -> usize { @@ -127,12 +203,12 @@ impl VTable for DecimalByteParts { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { + vortex_ensure!( + array.lower_parts().is_empty(), + "serializing DecimalByteParts with lower parts is not supported" + ); Ok(Some( - DecimalBytesPartsMetadata { - zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, - lower_part_count: 0, - } - .encode_to_vec(), + DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), )) } @@ -146,26 +222,15 @@ impl VTable for DecimalByteParts { _session: &VortexSession, ) -> VortexResult> { let metadata = DecimalBytesPartsMetadata::decode(metadata)?; - let Some(decimal_dtype) = dtype.as_decimal_opt() else { - vortex_bail!("decoding decimal but given non decimal dtype {}", dtype) - }; - - let encoded_dtype = DType::Primitive(metadata.zeroth_child_ptype(), dtype.nullability()); - - let msp = children.get(0, &encoded_dtype, len)?; - - assert_eq!( - metadata.lower_part_count, 0, - "lower_part_count > 0 not currently supported" + vortex_ensure!( + metadata.lower_part_count()? == 0, + "vortex.decimal_byte_parts must not carry lower parts" ); - - let slots = smallvec![Some(msp.clone())]; - let data = DecimalBytePartsData::try_new(msp.dtype(), msp.len(), *decimal_dtype)?; - Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) + metadata.into_array_parts(dtype, len, children) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - DecimalBytePartsSlots::NAMES[idx].to_string() + DecimalBytePartsSlots::slot_name(idx) } fn reduce_parent( @@ -186,20 +251,21 @@ pub struct DecimalBytePartsSlots { /// The most significant parts of the decimal values. #[slot(0)] pub msp: ArrayRef, + /// The remaining 64-bit windows of the decimal values, most significant first. + #[slot(1..)] + pub lower_parts: Vec, } /// This array encodes decimals as between 1-4 columns of primitive typed children. -/// The most significant part (msp) sorting the most significant decimal bits. +/// The most significant part (msp) storing the most significant decimal bits. /// This array must be signed and is nullable iff the decimal is nullable. +/// Every lower part is a non-nullable `u64` holding a raw 64-bit window of the value. +/// +/// e.g. for a decimal i128 \[ 127..64 | 63..0 \] msp = 127..64 and lower_part\[0\] = 63..0 /// -/// e.g. for a decimal i128 \[ 127..64 | 64..0 \] msp = 127..64 and lower_part\[0\] = 64..0 +/// All parts live in slots, so the array carries no additional data. #[derive(Clone, Debug)] -pub struct DecimalBytePartsData { - // NOTE: the lower_parts is currently unused, we reserve this field so that it is properly - // read/written during serde, but provide no constructor to initialize this to anything - // other than the empty Vec. - _lower_parts: Vec, -} +pub struct DecimalBytePartsData; impl Display for DecimalBytePartsData { fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { @@ -207,13 +273,17 @@ impl Display for DecimalBytePartsData { } } -pub struct DecimalBytePartsDataParts { - pub msp: ArrayRef, -} - impl DecimalBytePartsData { - pub fn validate( + /// Validate the parts of a [`DecimalBytePartsArray`]. + /// + /// # 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 `u64` array of length `len`. + pub fn validate<'a>( msp: &ArrayRef, + lower_parts: impl ExactSizeIterator, decimal_dtype: DecimalDType, dtype: &DType, len: usize, @@ -228,24 +298,26 @@ impl DecimalBytePartsData { "expected dtype {expected_dtype}, got {dtype}" ); vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); - Ok(()) - } - pub(crate) fn try_new( - msp_dtype: &DType, - msp_len: usize, - decimal_dtype: DecimalDType, - ) -> VortexResult { - let expected_dtype = DType::Decimal(decimal_dtype, msp_dtype.nullability()); + let lower_part_count = lower_parts.len(); + // Physical storage may be wider than the declared precision, as for DecimalArray. vortex_ensure!( - msp_dtype.is_signed_int(), - "decimal bytes parts, first part must be a signed array" + lower_part_count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" ); - let _ = msp_len; - drop(expected_dtype); - Ok(Self { - _lower_parts: Vec::new(), - }) + for (idx, part) in lower_parts.enumerate() { + vortex_ensure!( + part.dtype() == &LOWER_PART_DTYPE, + "lower part {idx} must have dtype {LOWER_PART_DTYPE}, got {}", + part.dtype() + ); + vortex_ensure!( + part.len() == len, + "lower part {idx} has len {}, expected {len}", + part.len() + ); + } + Ok(()) } } @@ -254,47 +326,108 @@ pub struct DecimalByteParts; impl DecimalByteParts { /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array. pub fn try_new( msp: ArrayRef, decimal_dtype: DecimalDType, ) -> VortexResult { + Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + } + + /// 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 `u64` + /// array of the same length as the MSP. See [`split_decimal`] for producing them from a + /// canonical decimal array. + /// + /// # Errors + /// + /// Returns an error if the parts do not describe a valid decimal, see + /// [`DecimalBytePartsData::validate`]. + pub fn try_new_with_lower_parts( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> VortexResult { + // Lower parts are supported in memory; the frozen serializer still rejects them. let len = msp.len(); let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); - let slots = smallvec![Some(msp.clone())]; - let data = DecimalBytePartsData::try_new(msp.dtype(), msp.len(), decimal_dtype)?; - Ok(unsafe { - Array::from_parts_unchecked( - ArrayParts::new(DecimalByteParts, dtype, len, data).with_slots(slots), - ) - }) + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + Array::try_from_parts( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + ) + } +} + +/// The decimal storage type this array canonicalizes to. +fn values_type(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { + match array.lower_parts().len() { + 0 => DecimalType::try_from(array.msp().dtype().as_ptype()), + 1 => Ok(DecimalType::I128), + 2 | 3 => Ok(DecimalType::I256), + count => vortex_bail!("at most {MAX_LOWER_PARTS} lower parts are supported, got {count}"), } } +/// The decimal dtype this array carries. +/// +/// Guaranteed to be a decimal by construction: [`DecimalBytePartsData::validate`] rejects +/// every other dtype. +pub(crate) fn decimal_dtype(array: ArrayView<'_, DecimalByteParts>) -> DecimalDType { + *array + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype") +} + +/// Rebuild the array by applying `f` to the MSP and to every lower part, in slot order. +/// +/// Part-wise operations must touch every part. Going through this rather than calling +/// [`DecimalByteParts::try_new_with_lower_parts`] directly makes dropping a lower part — +/// which silently corrupts wide values — unrepresentable. +pub(crate) fn map_parts( + array: ArrayView<'_, DecimalByteParts>, + mut f: impl FnMut(&ArrayRef) -> VortexResult, +) -> VortexResult { + let msp = f(array.msp())?; + let lower_parts = array + .lower_parts() + .iter() + .map(&mut f) + .collect::>>()?; + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal_dtype(array)) +} + +/// Rebuild the array with a replacement MSP, keeping its lower parts untouched. +/// +/// Only valid for operations that cannot change a row's magnitude bits — a nullability cast +/// or a mask — since the lower parts keep whatever bits they held. That is sound because +/// validity lives in the MSP alone, so lower-part bits in a null row are already undefined. +pub(crate) fn with_msp( + array: ArrayView<'_, DecimalByteParts>, + msp: ArrayRef, + decimal_dtype: DecimalDType, +) -> VortexResult { + DecimalByteParts::try_new_with_lower_parts(msp, array.lower_parts().to_vec(), decimal_dtype) +} + /// Converts a DecimalBytePartsArray to its canonical DecimalArray representation. fn to_canonical_decimal( array: &DecimalBytePartsArray, ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): support parts len != 1 - let prim = array.msp().clone().execute::(ctx)?; - // Depending on the decimal type and the min/max of the primitive array we can choose - // the correct buffer size - - Ok(match_each_signed_integer_ptype!(prim.ptype(), |P| { - // SAFETY: The primitive array's buffer is already validated with correct type. - // The decimal dtype matches the array's dtype, and validity is preserved. - unsafe { - DecimalArray::new_unchecked( - prim.to_buffer::

(), - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - prim.validity()?, - ) - } - .into_array() - })) + let msp = array.msp().clone().execute::(ctx)?; + let lower_parts = array + .lower_parts() + .iter() + .map(|part| part.clone().execute::(ctx)) + .collect::>>()?; + + Ok(assemble_decimal(&msp, &lower_parts, decimal_dtype(array.as_view()))?.into_array()) } impl OperationsVTable for DecimalByteParts { @@ -303,17 +436,34 @@ impl OperationsVTable for DecimalByteParts { index: usize, ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): support parts len != 1 let scalar = array.msp().execute_scalar(index, ctx)?; // Note. values in msp, can only be signed integers upto size i64. let primitive_scalar = scalar.as_primitive(); - // TODO(joe): extend this to support multiple parts. - let value = primitive_scalar.as_::().vortex_expect("non-null"); - Scalar::try_new( - array.dtype().clone(), - Some(ScalarValue::Decimal(DecimalValue::I64(value))), - ) + let msp = primitive_scalar.as_::().vortex_expect("non-null"); + + let lower_parts = array + .lower_parts() + .iter() + .map(|part| { + Ok(part + .execute_scalar(index, ctx)? + .as_primitive() + .as_::() + .vortex_expect("lower parts are non-nullable")) + }) + .collect::>>()?; + + let value = if lower_parts.is_empty() { + DecimalValue::I64(msp) + } else { + match values_type(array)? { + DecimalType::I256 => DecimalValue::I256(combine_i256(msp, lower_parts.into_iter())), + _ => DecimalValue::I128(combine_i128(msp, lower_parts)), + } + }; + + Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value))) } } @@ -326,21 +476,32 @@ impl ValidityChild for DecimalByteParts { #[cfg(test)] mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; 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::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; + use super::*; use crate::DecimalByteParts; + 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_scalar_at_decimal_parts() { @@ -380,4 +541,267 @@ mod tests { .unwrap() ); } + + /// The largest unscaled value a `Decimal(38, _)` can hold: `10^38 - 1`. + const MAX_PRECISION_38: i128 = 99_999_999_999_999_999_999_999_999_999_999_999_999; + + /// The largest unscaled value a `Decimal(76, _)` can hold: `10^76 - 1`. + fn max_precision_76() -> i256 { + i256::from_i128(10).wrapping_pow(76) - i256::ONE + } + + /// Values that exercise every 64-bit window of an `i128`, both signs, and the boundaries + /// where a lower part carries into the MSP. + fn wide_i128_values() -> Vec { + vec![ + 0, + 1, + -1, + (1 << 64) - 1, + 1 << 64, + -(1 << 64), + -((1 << 64) + 1), + MAX_PRECISION_38, + -MAX_PRECISION_38, + 1 << 100, + ] + } + + /// Values that exercise every 64-bit window of an `i256`. + fn wide_i256_values() -> Vec { + vec![ + i256::ZERO, + i256::ONE, + i256::ZERO - i256::ONE, + i256_of(0, u128::MAX), + i256_of(1, 0), + i256_of(-1, 0), + i256_of(-1, u128::MAX - 1), + i256_of(1 << 64, 12345), + max_precision_76(), + i256::ZERO - max_precision_76(), + ] + } + + #[rstest] + #[case::i128_non_nullable(i128_parts(wide_i128_values(), Validity::NonNullable))] + #[case::i256_non_nullable(i256_parts(wide_i256_values(), Validity::NonNullable))] + fn test_canonical_decimal_round_trips( + #[case] array: DecimalBytePartsArray, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(array, canonical, &mut ctx); + Ok(()) + } + + #[test] + fn test_lower_part_layout_i128() -> VortexResult<()> { + let array = i128_parts(vec![(3i128 << 64) | 7], Validity::NonNullable); + assert_eq!(array.lower_parts().len(), 1); + assert_eq!(array.msp().dtype().as_ptype(), PType::I64); + assert_eq!(array.lower_parts()[0].dtype(), &LOWER_PART_DTYPE); + + let mut ctx = array_session().create_execution_ctx(); + let msp = array.msp().clone().execute::(&mut ctx)?; + let lower = array.lower_parts()[0] + .clone() + .execute::(&mut ctx)?; + assert_eq!(msp.as_slice::(), &[3]); + assert_eq!(lower.as_slice::(), &[7]); + Ok(()) + } + + #[test] + fn test_lower_part_layout_i256() -> VortexResult<()> { + let array = i256_parts( + vec![i256_of((5i128 << 64) | 6, (7u128 << 64) | 8)], + Validity::NonNullable, + ); + assert_eq!(array.lower_parts().len(), MAX_LOWER_PARTS); + + let mut ctx = array_session().create_execution_ctx(); + let msp = array.msp().clone().execute::(&mut ctx)?; + assert_eq!(msp.as_slice::(), &[5]); + for (part, expected) in array.lower_parts().iter().zip([6u64, 7, 8]) { + let part = part.clone().execute::(&mut ctx)?; + assert_eq!(part.as_slice::(), &[expected]); + } + Ok(()) + } + + #[rstest] + #[case::i128(i128_parts(wide_i128_values(), Validity::AllValid))] + #[case::i256(i256_parts(wide_i256_values(), Validity::AllValid))] + fn test_scalar_at_matches_canonical(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)? + .into_array(); + let array = array.into_array(); + for idx in 0..array.len() { + assert_eq!( + array.execute_scalar(idx, &mut ctx)?, + canonical.execute_scalar(idx, &mut ctx)?, + "scalar mismatch at index {idx}" + ); + } + Ok(()) + } + + #[rstest] + fn test_scalar_at_matches_canonical_for_each_part_count( + #[values(false, true)] narrow_msp: bool, + #[values(0, 1, 2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let validity = Validity::from_iter([false, true, true]); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![0i8, 3, -3], validity) + } else { + PrimitiveArray::new(buffer![0i64, 3, -3], validity) + }; + let lower = [4u64, 1, 2] + .into_iter() + .take(lower_count) + .map(|word| PrimitiveArray::new(buffer![word; 3], Validity::NonNullable).into_array()) + .collect(); + let dtype = DecimalDType::new(if lower_count <= 1 { 38 } else { 76 }, 0); + let array = DecimalByteParts::try_new_with_lower_parts(msp.into_array(), lower, dtype)?; + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)?; + for row in 0..array.len() { + assert_eq!( + array.execute_scalar(row, &mut ctx)?, + canonical.execute_scalar(row, &mut ctx)? + ); + } + Ok(()) + } + + #[test] + fn test_scalar_at_null_with_lower_parts() -> VortexResult<()> { + let array = i128_parts( + vec![1i128 << 100, 2, 3], + Validity::Array(BoolArray::from_iter([false, true, true]).into_array()), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::decimal( + DecimalValue::I128(2), + DecimalDType::new(38, 2), + Nullability::Nullable + ) + ); + Ok(()) + } + + fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() + } + + fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() + } + + #[rstest] + #[case::signed_lower_part(vec![buffer![1i64, 2, 3].into_array()], DecimalDType::new(38, 2))] + #[case::nullable_lower_part( + vec![PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid).into_array()], + DecimalDType::new(38, 2) + )] + #[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) + )] + fn test_rejects_invalid_parts( + #[case] lower_parts: Vec, + #[case] decimal_dtype: DecimalDType, + ) { + assert!( + DecimalByteParts::try_new_with_lower_parts(msp(), lower_parts, decimal_dtype).is_err() + ); + } + + #[test] + fn test_wide_decimal_buffer_types() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let i128_array = i128_parts(vec![1i128 << 100], Validity::NonNullable); + let canonical = i128_array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + + let i256_array = i256_parts(vec![i256_of(1 << 100, 0)], Validity::NonNullable); + let canonical = i256_array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + + // A narrow MSP with a single lower part still fits 128 bits. + let array = DecimalByteParts::try_new_with_lower_parts( + buffer![1i8, -1, 0].into_array(), + vec![buffer![7u64, 7, 7].into_array()], + DecimalDType::new(38, 2), + )?; + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + assert_eq!( + canonical.buffer::().as_slice(), + &[(1i128 << 64) | 7, (-1i128 << 64) | 7, 7] + ); + + // Two lower parts under a narrow MSP overflow 128 bits, so the value widens. + let array = DecimalByteParts::try_new_with_lower_parts( + buffer![1i8].into_array(), + vec![buffer![0u64].into_array(), buffer![9u64].into_array()], + DecimalDType::new(76, 2), + )?; + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + assert_eq!(canonical.buffer::().as_slice(), &[i256_of(1, 9)]); + Ok(()) + } + + #[test] + fn test_unused_buffer_of_values_is_ignored_for_null_rows() -> VortexResult<()> { + // Null rows may hold arbitrary bits in the lower parts; they must stay null. + let array = DecimalByteParts::try_new_with_lower_parts( + PrimitiveArray::new( + buffer![0i64, 0, 0], + Validity::Array(BoolArray::from_iter([false, false, true]).into_array()), + ) + .into_array(), + vec![buffer![7u64, 9, 11].into_array()], + DecimalDType::new(38, 2), + )? + .into_array(); + + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + let canonical = array.clone().execute::(&mut ctx)?; + assert_arrays_eq!(array, canonical.into_array(), &mut ctx); + Ok(()) + } + #[test] + fn test_frozen_serializer_rejects_lower_parts() -> VortexResult<()> { + let session = array_session(); + let array = i128_parts(vec![1i128 << 70], Validity::NonNullable); + assert!(VTable::serialize(array.as_view(), &session).is_err()); + Ok(()) + } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs index d4052a4bed8..28503d5d8af 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/rules.rs @@ -1,57 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::ArrayRef; -use vortex_array::ArrayView; -use vortex_array::IntoArray; -use vortex_array::arrays::Filter; +use vortex_array::arrays::dict::TakeReduceAdaptor; use vortex_array::arrays::filter::FilterReduceAdaptor; use vortex_array::arrays::slice::SliceReduceAdaptor; -use vortex_array::optimizer::rules::ArrayParentReduceRule; use vortex_array::optimizer::rules::ParentRuleSet; use vortex_array::scalar_fn::fns::cast::CastReduceAdaptor; use vortex_array::scalar_fn::fns::mask::MaskReduceAdaptor; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; pub(super) const PARENT_RULES: ParentRuleSet = ParentRuleSet::new(&[ - ParentRuleSet::lift(&DecimalBytePartsFilterPushDownRule), ParentRuleSet::lift(&CastReduceAdaptor(DecimalByteParts)), ParentRuleSet::lift(&FilterReduceAdaptor(DecimalByteParts)), ParentRuleSet::lift(&MaskReduceAdaptor(DecimalByteParts)), ParentRuleSet::lift(&SliceReduceAdaptor(DecimalByteParts)), + ParentRuleSet::lift(&TakeReduceAdaptor(DecimalByteParts)), ]); - -#[derive(Debug)] -struct DecimalBytePartsFilterPushDownRule; - -impl ArrayParentReduceRule for DecimalBytePartsFilterPushDownRule { - type Parent = Filter; - - fn reduce_parent( - &self, - child: ArrayView<'_, DecimalByteParts>, - parent: ArrayView<'_, Filter>, - _child_idx: usize, - ) -> VortexResult> { - // TODO(ngates): we should benchmark whether to push-down filters with "lower parts". - // For now, we only push down if there are no lower parts. - if !child._lower_parts.is_empty() { - return Ok(None); - } - - let new_msp = child.msp().filter(parent.filter_mask().clone())?; - let new_child = DecimalByteParts::try_new( - new_msp, - *child - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - )? - .into_array(); - Ok(Some(new_child)) - } -} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs index 14807421c73..e31f717d389 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs @@ -7,23 +7,13 @@ use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::arrays::slice::SliceReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::map_parts; impl SliceReduce for DecimalByteParts { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { - Ok(Some( - DecimalByteParts::try_new( - array.msp().slice(range)?, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - )? - .into_array(), - )) + map_parts(array, |part| part.slice(range.clone())).map(|d| Some(d.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs new file mode 100644 index 00000000000..d2ce68f3700 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Test-only helpers for building byte-parts arrays. + +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::DecimalByteParts; +use crate::DecimalBytePartsArray; +use crate::decimal_byte_parts::limbs::split_decimal; + +/// Encode a canonical decimal array as byte parts, splitting wide values into lower parts. +pub(crate) fn encode(decimal: &DecimalArray) -> VortexResult { + let parts = split_decimal(decimal, &mut array_session().create_execution_ctx())?; + DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ) +} + +/// An `i128`-backed decimal array, encoded as byte parts with one lower part. +pub(crate) fn i128_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { + encode(&DecimalArray::new( + Buffer::from(values), + DecimalDType::new(38, 2), + validity, + )) + .vortex_expect("valid decimal byte parts") +} + +/// An `i256`-backed decimal array, encoded as byte parts with three lower parts. +pub(crate) fn i256_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { + encode(&DecimalArray::new( + Buffer::from(values), + DecimalDType::new(76, 2), + validity, + )) + .vortex_expect("valid decimal byte parts") +} + +/// Build an `i256` from a signed high `i128` and unsigned low `u128`. +pub(crate) fn i256_of(high: i128, low: u128) -> i256 { + i256::from_parts(low, high) +} diff --git a/encodings/decimal-byte-parts/tests/props.rs b/encodings/decimal-byte-parts/tests/props.rs new file mode 100644 index 00000000000..e33606b2880 --- /dev/null +++ b/encodings/decimal-byte-parts/tests/props.rs @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Property tests for splitting decimals into byte parts and putting them back together. +//! +//! Every property here is the same shape: whatever the encoding does must be indistinguishable +//! from doing it to the canonical `DecimalArray`. Round tripping covers the split/assemble +//! pair directly; the compute properties cover it indirectly, since each one canonicalizes an +//! encoded array at the end. +//! +//! The generators deliberately reach the cases hand-written tests tend to miss: values that +//! straddle a 64-bit word boundary, negative values whose sign extension fills the words above +//! the most significant part, and null rows whose lower parts hold arbitrary bits. + +#![expect(clippy::tests_outside_test_module)] + +use hegel::TestCase; +use hegel::generators as gs; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +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::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArray; +use vortex_decimal_byte_parts::split_decimal; +use vortex_error::VortexExpect; + +/// Largest magnitude a `Decimal(38, _)` can hold: 38 nines. +const MAX_I128: i128 = 10i128.pow(38) - 1; + +/// Bound on the high `i128` half of an `i256` draw. `10^37 * 2^128` is about `3.4e75`, so any +/// value built from it stays inside the 76 digits a `Decimal(76, _)` can hold. +const MAX_I256_HIGH: i128 = 10i128.pow(37); + +/// Rows per generated array. Small enough to shrink usefully, large enough that a chunked or +/// vectorized path is not trivially degenerate. +const MAX_LEN: usize = 48; + +fn ctx() -> ExecutionCtx { + let session = array_session(); + vortex_decimal_byte_parts::initialize(&session); + session.create_execution_ctx() +} + +/// Encode a canonical decimal as byte parts, splitting wide values into lower parts. +fn encode(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> DecimalBytePartsArray { + let parts = split_decimal(decimal, ctx).vortex_expect("split"); + DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ) + .vortex_expect("valid byte parts") +} + +/// A validity mask of exactly `len` entries, so null rows exercise lower parts holding bits +/// that must never be read. +fn draw_validity(tc: &TestCase, len: usize) -> Validity { + let valid: Vec = tc.draw(gs::vecs(gs::booleans()).min_size(len).max_size(len)); + Validity::from_iter(valid) +} + +/// An `i128`-backed decimal. The bounds keep values inside `Decimal(38, 2)` while still +/// reaching both sides of the 64-bit word boundary the encoding splits on. +fn draw_i128_decimal(tc: &TestCase) -> DecimalArray { + let values: Vec = tc.draw( + gs::vecs( + gs::integers::() + .min_value(-MAX_I128) + .max_value(MAX_I128), + ) + .min_size(1) + .max_size(MAX_LEN), + ); + let validity = draw_validity(tc, values.len()); + DecimalArray::new(Buffer::from(values), DecimalDType::new(38, 2), validity) +} + +/// An `i256`-backed decimal, built from a signed high half and an unsigned low half so the +/// draw covers sign extension above the most significant part. +fn draw_i256_decimal(tc: &TestCase) -> DecimalArray { + let halves: Vec<(i128, u128)> = tc.draw( + gs::vecs(gs::tuples2( + gs::integers::() + .min_value(-MAX_I256_HIGH) + .max_value(MAX_I256_HIGH), + gs::integers::(), + )) + .min_size(1) + .max_size(MAX_LEN), + ); + let values: Vec = halves + .into_iter() + .map(|(high, low)| i256::from_parts(low, high)) + .collect(); + let validity = draw_validity(tc, values.len()); + DecimalArray::new(Buffer::from(values), DecimalDType::new(76, 2), validity) +} + +fn draw_decimal(tc: &TestCase) -> DecimalArray { + if tc.draw(gs::booleans()) { + draw_i128_decimal(tc) + } else { + draw_i256_decimal(tc) + } +} + +/// Canonicalize an encoded array back to a `DecimalArray`. +fn canonicalize(array: ArrayRef, ctx: &mut ExecutionCtx) -> DecimalArray { + array.execute::(ctx).vortex_expect("execute") +} + +/// A byte-parts array built directly from drawn parts, rather than by splitting a decimal. +/// +/// `split_decimal` only ever emits 0, 1 or 3 lower parts under an `i64` most significant +/// part, so drawing the part count here is the only way to reach the two-part shape and the +/// sign extension that sits above a most significant part below the top word. +fn draw_encoded(tc: &TestCase) -> (DecimalBytePartsArray, usize) { + let lower_part_count = tc.draw(gs::integers::().min_value(0).max_value(3)); + let msp: Vec = tc.draw( + gs::vecs(gs::integers::()) + .min_size(1) + .max_size(MAX_LEN), + ); + let len = msp.len(); + + let lower: Vec = (0..lower_part_count) + .map(|_| { + let part: Vec = + tc.draw(gs::vecs(gs::integers::()).min_size(len).max_size(len)); + PrimitiveArray::new(Buffer::from(part), Validity::NonNullable).into_array() + }) + .collect(); + + // The declared precision must be wide enough for what the parts assemble into. + let precision = match lower_part_count { + 0 => 18, + 1 => 38, + _ => 76, + }; + let msp = PrimitiveArray::new(Buffer::from(msp), draw_validity(tc, len)).into_array(); + let array = + DecimalByteParts::try_new_with_lower_parts(msp, lower, DecimalDType::new(precision, 2)) + .vortex_expect("valid byte parts"); + (array, len) +} + +/// Encoding a decimal and decoding it again must reproduce it exactly, including null rows +/// and the storage width. +#[hegel::test] +fn decoded_survives_encode_then_decode(tc: TestCase) { + let decimal = draw_decimal(&tc); + let mut ctx = ctx(); + + let round_tripped = canonicalize(encode(&decimal, &mut ctx).into_array(), &mut ctx); + + assert_eq!(round_tripped.values_type(), decimal.values_type()); + assert_arrays_eq!(decimal, round_tripped, &mut ctx); +} + +/// Decoding an encoded array and encoding it again must not change the values it decodes to. +/// +/// 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. +#[hegel::test] +fn encoded_survives_decode_then_encode(tc: TestCase) { + let (array, _len) = draw_encoded(&tc); + let mut ctx = ctx(); + + let decoded = canonicalize(array.into_array(), &mut ctx); + let re_decoded = canonicalize(encode(&decoded, &mut ctx).into_array(), &mut ctx); + + assert_arrays_eq!(decoded, re_decoded, &mut ctx); +} + +// TODO(joe): restore the coverage removed alongside these two round trips. Each of the +// following was a property here and caught mutations that the round trips do not: +// +// - `scalar_at` against bulk canonicalization. `combine_i128`/`combine_i256` are a second +// implementation of the assembly loops and can drift from them silently. +// - filter, slice and take against the same operation on the canonical array. These caught +// part-order and word-placement mutations, though the round trips catch those too. +// - a serialize/decode round trip, which is the only property that exercised the metadata +// carrying the lower part count. +// - sign extension above a most significant part below the top word, checked against an +// expectation computed independently of the assembly loop. This is the one real gap: a +// round trip compares decode against decode, so a decode-side sign-extension bug is +// invisible to it. Dropping the sign extension is caught by neither property here. diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 07069f6309a..f173440f26d 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -418,7 +418,7 @@ fn trace_scan_filter_on_compressed_table() -> VortexResult<()> { optimize root=vortex.filter(i16, len=43) session=false reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(i16, len=43) child=vortex.dict(i16, len=4096) -> vortex.dict(i16, len=43) done output=vortex.dict(i16, len=43) - reduce_parent static:DecimalBytePartsFilterPushDownRule slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) + reduce_parent static:FilterReduceAdaptor(DecimalByteParts) slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) done output=vortex.decimal_byte_parts(decimal(15,2), len=43) optimize root=vortex.filter(vortex.date[days](i32), len=43) session=false optimize root=vortex.filter(i32, len=43) session=false @@ -454,6 +454,9 @@ fn trace_scan_take_on_compressed_table() -> VortexResult<()> { insta::assert_snapshot!(optimized.trace.to_string(), @" optimize root=vortex.dict({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) session=false + optimize root=vortex.dict(decimal(15,2), len=64) session=false + reduce_parent static:TakeReduceAdaptor(DecimalByteParts) slot=1 parent=vortex.dict(decimal(15,2), len=64) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=64) + done output=vortex.decimal_byte_parts(decimal(15,2), len=64) optimize root=vortex.dict(vortex.date[days](i32), len=64) session=false reduce_parent static:TakeReduceAdaptor(Extension) slot=1 parent=vortex.dict(vortex.date[days](i32), len=64) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.ext(vortex.date[days](i32), len=64) done output=vortex.ext(vortex.date[days](i32), len=64) diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 3475f26a175..a54df06fb4c 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -39,6 +39,13 @@ impl CudaExecute for DecimalBytePartsExecutor { .dtype() .as_decimal_opt() .vortex_expect("DecimalBytePartsArray dtype must be decimal"); + + // Reassembling lower parts into wide decimals is not implemented on the GPU; the MSP + // alone is not the value. + if !array.lower_parts().is_empty() { + vortex_bail!("DecimalBytePartsArray with lower parts is not supported on GPU") + } + let msp = array.msp().clone(); let PrimitiveDataParts { buffer, From 13f8cd18064fa40fb03b08ca3b76d37f437c29dd Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 23:38:41 -0400 Subject: [PATCH 2/5] Simplify decimal byte-parts array helpers Move array helpers onto a crate-private extension trait, preserve decimal precision and scale when replacing the MSP, and group slicing with the other compute operations. Inline canonical execution and select scalar storage directly from the lower-part count. Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/compute/cast.rs | 6 +- .../src/decimal_byte_parts/compute/filter.rs | 6 +- .../src/decimal_byte_parts/compute/mask.rs | 5 +- .../src/decimal_byte_parts/compute/mod.rs | 1 + .../decimal_byte_parts/{ => compute}/slice.rs | 6 +- .../src/decimal_byte_parts/compute/take.rs | 6 +- .../src/decimal_byte_parts/mod.rs | 212 ++++++++---------- 7 files changed, 116 insertions(+), 126 deletions(-) rename encodings/decimal-byte-parts/src/decimal_byte_parts/{ => compute}/slice.rs (73%) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index 7b949fcd695..0594f15a7c6 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -10,8 +10,8 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; 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> { @@ -20,7 +20,7 @@ impl CastReduce for DecimalByteParts { return Ok(None); } // DecimalBytePartsArray can only have Decimal dtype, so we only handle decimal-to-decimal casts - let DType::Decimal(target_decimal, target_nullability) = dtype else { + let DType::Decimal(_, target_nullability) = dtype else { // Cannot cast decimal to non-decimal types - delegate to canonical form return Ok(None); }; @@ -30,7 +30,7 @@ impl CastReduce for DecimalByteParts { .msp() .cast(array.msp().dtype().with_nullability(*target_nullability))?; - with_msp(array, new_msp, *target_decimal).map(|a| Some(a.into_array())) + array.with_msp(new_msp).map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index e4fb03a5ca0..49c4021dd18 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs @@ -9,11 +9,13 @@ use vortex_error::VortexResult; use vortex_mask::Mask; use crate::DecimalByteParts; -use crate::decimal_byte_parts::map_parts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; impl FilterReduce for DecimalByteParts { fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult> { - map_parts(array, |part| part.filter(mask.clone())).map(|d| Some(d.into_array())) + array + .map_parts(|part| part.filter(mask.clone())) + .map(|d| Some(d.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index 2eea785794b..9a022ef34ce 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs @@ -9,15 +9,14 @@ use vortex_array::scalar_fn::fns::mask::MaskReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; -use crate::decimal_byte_parts::decimal_dtype; -use crate::decimal_byte_parts::with_msp; impl MaskReduce for DecimalByteParts { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { // Validity lives in the MSP, so only that part needs masking: the lower parts hold // undefined bits in null slots, which is exactly what a masked-out row is. let masked_msp = MaskExpr::try_new(array.msp().clone(), mask.clone())?.into_array(); - with_msp(array, masked_msp, decimal_dtype(array)).map(|a| Some(a.into_array())) + array.with_msp(masked_msp).map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 844468545cf..f9848e1b2e7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -7,6 +7,7 @@ mod filter; pub(crate) mod is_constant; pub(crate) mod kernel; mod mask; +mod slice; mod take; #[cfg(test)] diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs similarity index 73% rename from encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs rename to encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs index e31f717d389..1a2efc9034e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs @@ -10,10 +10,12 @@ use vortex_array::arrays::slice::SliceReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; -use crate::decimal_byte_parts::map_parts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; impl SliceReduce for DecimalByteParts { fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { - map_parts(array, |part| part.slice(range.clone())).map(|d| Some(d.into_array())) + array + .map_parts(|part| part.slice(range.clone())) + .map(|d| Some(d.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index 578834635b8..bdf7dda4f74 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs @@ -8,8 +8,8 @@ use vortex_array::arrays::dict::TakeReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; -use crate::decimal_byte_parts::map_parts; impl TakeReduce for DecimalByteParts { /// Taking wraps each part in a `Dict` without reading any buffer, so it reduces rather @@ -22,7 +22,9 @@ impl TakeReduce for DecimalByteParts { return Ok(None); } - map_parts(array, |part| part.take(indices.clone())).map(|a| Some(a.into_array())) + array + .map_parts(|part| part.take(indices.clone())) + .map(|a| Some(a.into_array())) } } 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 4a104e82c2b..ae02c23902a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -11,7 +11,6 @@ use vortex_array::ArrayView; pub(crate) mod compute; mod limbs; mod rules; -mod slice; #[cfg(test)] pub(crate) mod testing; @@ -32,13 +31,12 @@ use vortex_array::ArraySlots; use vortex_array::EqMode; use vortex_array::ExecutionCtx; use vortex_array::ExecutionResult; -use vortex_array::IntoArray; +use vortex_array::TypedArrayRef; use vortex_array::array_slots; use vortex_array::arrays::PrimitiveArray; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::DecimalType; use vortex_array::dtype::PType; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; @@ -148,6 +146,48 @@ impl DecimalBytesPartsMetadata { } } +#[derive(Clone, Debug)] +pub struct DecimalByteParts; + +impl DecimalByteParts { + /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array. + pub fn try_new( + msp: ArrayRef, + decimal_dtype: DecimalDType, + ) -> VortexResult { + Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + } + + /// 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 `u64` + /// array of the same length as the MSP. See [`split_decimal`] for producing them from a + /// canonical decimal array. + /// + /// # Errors + /// + /// Returns an error if the parts do not describe a valid decimal, see + /// [`DecimalBytePartsData::validate`]. + pub fn try_new_with_lower_parts( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> VortexResult { + // Lower parts are supported in memory; the frozen serializer still rejects them. + let len = msp.len(); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + Array::try_from_parts( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + ) + } +} + impl VTable for DecimalByteParts { type TypedArrayData = DecimalBytePartsData; @@ -242,7 +282,17 @@ impl VTable for DecimalByteParts { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - to_canonical_decimal(&array, ctx).map(ExecutionResult::done) + // Reassemble DecimalArray from split parts + let msp = array.msp().clone().execute::(ctx)?; + let lower_parts = array + .lower_parts() + .iter() + .map(|part| part.clone().execute::(ctx)) + .collect::>>()?; + + let assembled = assemble_decimal(&msp, &lower_parts, array.decimal_dtype())?; + + Ok(ExecutionResult::done(assembled)) } } @@ -256,9 +306,11 @@ pub struct DecimalBytePartsSlots { pub lower_parts: Vec, } -/// This array encodes decimals as between 1-4 columns of primitive typed children. -/// The most significant part (msp) storing the most significant decimal bits. -/// This array must be signed and is nullable iff the decimal is nullable. +/// 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. +/// /// Every lower part is a non-nullable `u64` holding a raw 64-bit window of the value. /// /// e.g. for a decimal i128 \[ 127..64 | 63..0 \] msp = 127..64 and lower_part\[0\] = 63..0 @@ -300,7 +352,7 @@ impl DecimalBytePartsData { vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); let lower_part_count = lower_parts.len(); - // Physical storage may be wider than the declared precision, as for DecimalArray. + vortex_ensure!( lower_part_count <= MAX_LOWER_PARTS, "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" @@ -321,114 +373,47 @@ impl DecimalBytePartsData { } } -#[derive(Clone, Debug)] -pub struct DecimalByteParts; +pub(crate) trait DecimalBytePartsArrayExt: DecimalBytePartsArraySlotsExt { + /// The decimal precision and scale, validated when the array was constructed. + fn decimal_dtype(&self) -> DecimalDType { + *self + .as_ref() + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype") + } -impl DecimalByteParts { - /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. - /// - /// # Errors + /// Rebuild the array by applying `f` to the MSP and every lower part, in slot order. /// - /// Returns an error if the MSP is not a signed integer array. - pub fn try_new( - msp: ArrayRef, - decimal_dtype: DecimalDType, + /// This applies row operations such as slicing and filtering to all parts together, + /// preserving the decimal precision and scale. + fn map_parts( + &self, + mut f: impl FnMut(&ArrayRef) -> VortexResult, ) -> VortexResult { - Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + let msp = f(self.msp())?; + let lower_parts = self + .lower_parts() + .iter() + .map(&mut f) + .collect::>>()?; + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, self.decimal_dtype()) } - /// Construct a new [`DecimalBytePartsArray`] from an MSP array, its lower parts, and a - /// decimal dtype. + /// Rebuild the array with a replacement MSP, preserving its lower parts, precision and scale. /// - /// Lower parts are ordered most significant first and must each be a non-nullable `u64` - /// array of the same length as the MSP. See [`split_decimal`] for producing them from a - /// canonical decimal array. - /// - /// # Errors - /// - /// Returns an error if the parts do not describe a valid decimal, see - /// [`DecimalBytePartsData::validate`]. - pub fn try_new_with_lower_parts( - msp: ArrayRef, - lower_parts: Vec, - decimal_dtype: DecimalDType, - ) -> VortexResult { - // Lower parts are supported in memory; the frozen serializer still rejects them. - let len = msp.len(); - let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); - let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); - Array::try_from_parts( - ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + /// Use this for operations such as masking and nullability casts that only affect the MSP. + /// The replacement MSP determines the result's nullability. + fn with_msp(&self, msp: ArrayRef) -> VortexResult { + DecimalByteParts::try_new_with_lower_parts( + msp, + self.lower_parts().to_vec(), + self.decimal_dtype(), ) } } -/// The decimal storage type this array canonicalizes to. -fn values_type(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { - match array.lower_parts().len() { - 0 => DecimalType::try_from(array.msp().dtype().as_ptype()), - 1 => Ok(DecimalType::I128), - 2 | 3 => Ok(DecimalType::I256), - count => vortex_bail!("at most {MAX_LOWER_PARTS} lower parts are supported, got {count}"), - } -} - -/// The decimal dtype this array carries. -/// -/// Guaranteed to be a decimal by construction: [`DecimalBytePartsData::validate`] rejects -/// every other dtype. -pub(crate) fn decimal_dtype(array: ArrayView<'_, DecimalByteParts>) -> DecimalDType { - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype") -} - -/// Rebuild the array by applying `f` to the MSP and to every lower part, in slot order. -/// -/// Part-wise operations must touch every part. Going through this rather than calling -/// [`DecimalByteParts::try_new_with_lower_parts`] directly makes dropping a lower part — -/// which silently corrupts wide values — unrepresentable. -pub(crate) fn map_parts( - array: ArrayView<'_, DecimalByteParts>, - mut f: impl FnMut(&ArrayRef) -> VortexResult, -) -> VortexResult { - let msp = f(array.msp())?; - let lower_parts = array - .lower_parts() - .iter() - .map(&mut f) - .collect::>>()?; - DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal_dtype(array)) -} - -/// Rebuild the array with a replacement MSP, keeping its lower parts untouched. -/// -/// Only valid for operations that cannot change a row's magnitude bits — a nullability cast -/// or a mask — since the lower parts keep whatever bits they held. That is sound because -/// validity lives in the MSP alone, so lower-part bits in a null row are already undefined. -pub(crate) fn with_msp( - array: ArrayView<'_, DecimalByteParts>, - msp: ArrayRef, - decimal_dtype: DecimalDType, -) -> VortexResult { - DecimalByteParts::try_new_with_lower_parts(msp, array.lower_parts().to_vec(), decimal_dtype) -} - -/// Converts a DecimalBytePartsArray to its canonical DecimalArray representation. -fn to_canonical_decimal( - array: &DecimalBytePartsArray, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let msp = array.msp().clone().execute::(ctx)?; - let lower_parts = array - .lower_parts() - .iter() - .map(|part| part.clone().execute::(ctx)) - .collect::>>()?; - - Ok(assemble_decimal(&msp, &lower_parts, decimal_dtype(array.as_view()))?.into_array()) -} +impl> DecimalBytePartsArrayExt for T {} impl OperationsVTable for DecimalByteParts { fn scalar_at( @@ -438,7 +423,8 @@ impl OperationsVTable for DecimalByteParts { ) -> VortexResult { let scalar = array.msp().execute_scalar(index, ctx)?; - // Note. values in msp, can only be signed integers upto size i64. + // Widen the MSP's signed value (i8/i16/i32/i64) to i64 for scalar reconstruction. + // The array retains its original MSP storage type. let primitive_scalar = scalar.as_primitive(); let msp = primitive_scalar.as_::().vortex_expect("non-null"); @@ -454,13 +440,10 @@ impl OperationsVTable for DecimalByteParts { }) .collect::>>()?; - let value = if lower_parts.is_empty() { - DecimalValue::I64(msp) - } else { - match values_type(array)? { - DecimalType::I256 => DecimalValue::I256(combine_i256(msp, lower_parts.into_iter())), - _ => DecimalValue::I128(combine_i128(msp, lower_parts)), - } + 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())), }; Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value))) @@ -487,6 +470,7 @@ mod tests { use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::DecimalType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; From 76a0db42e4450e936a845e87050e9fcccb0bdf17 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 16:06:25 -0400 Subject: [PATCH 3/5] Share decimal value assembly and validate lower-part dtypes Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/limbs/mod.rs | 75 ++++++++----------- .../src/decimal_byte_parts/limbs/tests.rs | 17 +++++ .../src/decimal_byte_parts/mod.rs | 20 +++-- 3 files changed, 64 insertions(+), 48 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs index 119e0a63854..ff722c05908 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ b/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::(); vortex_ensure!( part.len() == len, @@ -265,17 +271,17 @@ pub fn assemble_decimal( Ok(match lower.as_slice() { [first] => DecimalArray::new( - assemble_wide::(msp, [first]), + assemble_wide_decimal::(msp, [first]), decimal_dtype, validity, ), [first, second] => DecimalArray::new( - assemble_wide::(msp, [first, second]), + assemble_wide_decimal::(msp, [first, second]), decimal_dtype, validity, ), [first, second, third] => DecimalArray::new( - assemble_wide::(msp, [first, second, third]), + assemble_wide_decimal::(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) -> 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) -> 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(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer +/// A fixed part count lets the compiler unroll each call to [`assemble_wide_decimal_value`]. +fn assemble_wide_decimal(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer where T: NativeDecimalType + Shl + BitOr, { @@ -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(msp: i64, lower: [u64; K]) -> T +where + T: NativeDecimalType + Shl + BitOr, +{ + 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; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs index 3e3de06c44e..dd90dd7e197 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs +++ b/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, 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 ae02c23902a..f064a5210d3 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/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 for DecimalByteParts { }) .collect::>>()?; - 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))) From 215e0d56ec2d32729a285fbfb14dd70a1809f237 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 16:18:47 -0400 Subject: [PATCH 4/5] Validate required decimal byte-parts slots before typed access Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/mod.rs | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) 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 f064a5210d3..b9babd52aaa 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -208,6 +208,18 @@ impl VTable for DecimalByteParts { let Some(decimal_dtype) = dtype.as_decimal_opt() else { vortex_bail!("expected decimal dtype, got {}", dtype) }; + + let min_slots = DecimalBytePartsSlots::FIXED_COUNT; + let max_slots = min_slots + MAX_LOWER_PARTS; + vortex_ensure!( + (min_slots..=max_slots).contains(&slots.len()), + "expected {min_slots}..={max_slots} slots, got {}", + slots.len() + ); + for (idx, slot) in slots.iter().enumerate() { + vortex_ensure!(slot.is_some(), "missing required slot {idx}"); + } + let slots = DecimalBytePartsSlotsView::from_slots(slots); DecimalBytePartsData::validate( slots.msp, @@ -729,6 +741,22 @@ mod tests { ); } + #[rstest] + #[case::no_slots(vec![])] + #[case::missing_msp(vec![None])] + #[case::missing_lower(vec![Some(msp()), None])] + #[case::gap_in_lower(vec![Some(msp()), None, Some(lower_part())])] + fn test_rejects_missing_slots(#[case] slots: Vec>) { + let parts = ArrayParts::new( + DecimalByteParts, + DType::Decimal(DecimalDType::new(76, 2), Nullability::NonNullable), + 3, + DecimalBytePartsData, + ) + .with_slots(slots.into_iter().collect()); + assert!(Array::try_from_parts(parts).is_err()); + } + #[test] fn test_wide_decimal_buffer_types() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); From 7283d4f616f12f6cd061a682ea78c810717a2690 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 16:27:45 -0400 Subject: [PATCH 5/5] fix comment Signed-off-by: Matt Katz --- .../decimal-byte-parts/src/decimal_byte_parts/compute/take.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index bdf7dda4f74..9e504433a2b 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs @@ -12,8 +12,6 @@ use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl TakeReduce for DecimalByteParts { - /// Taking wraps each part in a `Dict` without reading any buffer, so it reduces rather - /// than executes. fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { // Taking with nullable indices makes every taken part nullable, but lower parts must // stay non-nullable `u64` — validity belongs to the MSP alone. Fall back to the