From 38a1dee6cd8fc548b4a00881292d20bf93c3c69d Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 16:27:02 -0400 Subject: [PATCH 1/9] 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 c4c7877156937417aa40e190e8b11bbcec9b4e52 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 23:38:41 -0400 Subject: [PATCH 2/9] 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 e30bd995bd9893fbb5b6db5a4b0345cb6fe75762 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 16:06:25 -0400 Subject: [PATCH 3/9] 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 6a7b033aa19c1a9a6da6a28d7899c890f293c070 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 16:18:47 -0400 Subject: [PATCH 4/9] 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 a20094e4d1f10279909b0988994dc311b438c835 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 16:27:45 -0400 Subject: [PATCH 5/9] 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 From 38b5bc461350492ea1a722c188e5d67755f9cccf Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 10 Sep 2026 16:26:39 -0400 Subject: [PATCH 6/9] QOL and perf changes for splitting Signed-off-by: Matt Katz --- Cargo.lock | 1 + encodings/decimal-byte-parts/Cargo.toml | 1 + .../benches/dbp_assemble.rs | 1 + .../decimal-byte-parts/benches/dbp_split.rs | 8 ++ .../src/decimal_byte_parts/limbs/mod.rs | 84 +++++++++++++------ .../src/decimal_byte_parts/limbs/tests.rs | 60 +++++++++++-- 6 files changed, 121 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 639e06bf078..41715bdec35 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11007,6 +11007,7 @@ dependencies = [ "rand 0.10.2", "rstest", "vortex-array", + "vortex-bench-support", "vortex-buffer", "vortex-error", "vortex-mask", diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index e9ea8569af1..e5d15c2c4f7 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -31,6 +31,7 @@ hegeltest = { workspace = true } rand = { workspace = true } rstest = { workspace = true } vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } +vortex-bench-support = { workspace = true } [[bench]] name = "dbp_assemble" diff --git a/encodings/decimal-byte-parts/benches/dbp_assemble.rs b/encodings/decimal-byte-parts/benches/dbp_assemble.rs index 327899fa640..cab7de81baa 100644 --- a/encodings/decimal-byte-parts/benches/dbp_assemble.rs +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -24,6 +24,7 @@ fn main() { divan::main(); } +#[vortex_bench_support::cpu_features] #[divan::bench(args = cases())] fn dbp_assemble(bencher: Bencher, (values_type, len): (DecimalType, usize)) { let decimal = decimal_array(values_type, len, Validity::NonNullable); diff --git a/encodings/decimal-byte-parts/benches/dbp_split.rs b/encodings/decimal-byte-parts/benches/dbp_split.rs index ba716d1ba5b..7bf335cfcce 100644 --- a/encodings/decimal-byte-parts/benches/dbp_split.rs +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -24,11 +24,19 @@ fn main() { divan::main(); } +#[vortex_bench_support::cpu_features] #[divan::bench(args = cases())] fn dbp_split_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { bench_split(bencher, values_type, len, Validity::AllValid); } +#[vortex_bench_support::cpu_features] +#[divan::bench(args = cases())] +fn dbp_split_all_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Validity::AllInvalid); +} + +#[vortex_bench_support::cpu_features] #[divan::bench(args = cases())] fn dbp_split_mixed_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { let mut rng = StdRng::seed_from_u64(42); 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 ff722c05908..b88486f1556 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 @@ -20,6 +20,7 @@ use std::ops::Shl; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; @@ -31,6 +32,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::i256; use vortex_array::match_each_signed_integer_ptype; +use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -40,11 +42,17 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_mask::Mask; -/// The maximum number of lower parts an encoded decimal can carry. Each is 64 bits. +/// The maximum number of 64-bit lower parts an encoded `i128` decimal can carry. +pub const MAX_I128_LOWER_PARTS: usize = 1; + +/// The maximum number of 64-bit lower parts an encoded `i256` decimal can carry. +pub const MAX_I256_LOWER_PARTS: usize = 3; + +/// The maximum number of 64-bit lower parts an encoded decimal can carry. /// /// Since the MSP is at most 64 bits wide, three additional 64-bit parts saturates /// the 256-bit maximum width of a Vortex decimal. -pub const MAX_LOWER_PARTS: usize = 3; +pub const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS; /// Number of bits stored in each lower part. const LOWER_PART_BITS: usize = 64; @@ -62,15 +70,18 @@ pub struct DecimalParts { } impl DecimalParts { - /// Construct decimal parts from an MSP with no lower parts. - fn from_msp(values: Buffer, validity: Validity) -> Self { + /// Construct decimal parts from the MSP buffer constituting a narrow decimal (`i64` or narrower). + /// Narrow decimals have an MSP at most as wide as `i64` and no lower parts. + fn from_narrow(values: Buffer, validity: Validity) -> Self { Self { msp: PrimitiveArray::new(values, validity).into_array(), lower_parts: Vec::new(), } } - fn new( + /// Construct decimal parts arrays from the buffers constituting a wide decimal (`i128` or `i256`). + /// Wide decimals have an `i64` MSP and up to [`MAX_LOWER_PARTS`] `u64` lower parts. + fn from_wide( msp: Buffer, lower_parts: impl IntoIterator>, validity: Validity, @@ -94,35 +105,63 @@ impl DecimalParts { /// /// The MSP retains the decimal's validity while lower parts are non-nullable. Lower parts /// are constructed with zeroes at null positions instead of invalid bytes. +/// Empty and all-null arrays use constant parts, preserving the part types and MSP's nullability. /// /// # Errors /// -/// Returns an error if the array's validity cannot be derived or executed. +/// * If the array's validity cannot be derived or executed. pub fn split_decimal(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { let validity = decimal.validity()?; + let len = decimal.len(); + let mask = validity.execute_mask(len, ctx)?; + + if mask.all_false() || decimal.is_empty() { + return Ok(split_no_valid_row(decimal, &validity)); + } + Ok(match decimal.values_type() { - DecimalType::I8 => DecimalParts::from_msp(decimal.buffer::(), validity), - DecimalType::I16 => DecimalParts::from_msp(decimal.buffer::(), validity), - DecimalType::I32 => DecimalParts::from_msp(decimal.buffer::(), validity), - DecimalType::I64 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I8 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I16 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I32 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I64 => DecimalParts::from_narrow(decimal.buffer::(), validity), DecimalType::I128 => { - let mask = validity.execute_mask(decimal.len(), ctx)?; let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i128_to_parts); - DecimalParts::new(msp, lower, validity) + DecimalParts::from_wide(msp, lower, validity) } DecimalType::I256 => { - let mask = validity.execute_mask(decimal.len(), ctx)?; let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i256_to_parts); - DecimalParts::new(msp, lower, validity) + DecimalParts::from_wide(msp, lower, validity) } }) } +/// Splits decimals with no valid rows (all null or empty) into constant decimal parts with the +/// corresponding nullability. +fn split_no_valid_row(decimal: &DecimalArray, validity: &Validity) -> DecimalParts { + let (msp_ptype, lower_part_count) = match decimal.values_type() { + DecimalType::I8 => (PType::I8, 0), + DecimalType::I16 => (PType::I16, 0), + DecimalType::I32 => (PType::I32, 0), + DecimalType::I64 => (PType::I64, 0), + DecimalType::I128 => (PType::I64, MAX_I128_LOWER_PARTS), + DecimalType::I256 => (PType::I64, MAX_I256_LOWER_PARTS), + }; + // Empty masks are also all-false. The default scalar is null for nullable inputs + // and zero for non-nullable empty inputs, preserving the MSP's nullability. + let msp = Scalar::default_value(&DType::Primitive(msp_ptype, validity.nullability())); + let len = decimal.len(); + DecimalParts { + msp: ConstantArray::new(msp, len).into_array(), + lower_parts: vec![ConstantArray::new(0u64, len).into_array(); lower_part_count], + } +} + /// Split wide integers into a signed MSP and `N` unsigned lower parts. /// /// `to_parts` returns the MSP and lower words in most-significant-first order. /// It is specialized for each input type: `i128` has one lower word and `i256` -/// has three. Null rows get zeros in every output buffer. +/// has three. Null rows get zeros in every output buffer. The caller handles empty +/// and all-null arrays before calling this function. fn split_wide( values: &Buffer, validity: &Mask, @@ -132,15 +171,6 @@ fn split_wide( let mut msp = BufferMut::::with_capacity(len); let mut lower = std::array::from_fn::<_, N, _>(|_| BufferMut::::with_capacity(len)); - // Zero out all parts if all null - if validity.all_false() { - msp.push_n(0, len); - for part in &mut lower { - part.push_n(0, len); - } - return (msp.freeze(), lower.map(BufferMut::freeze)); - } - // Allocate without zeroing, then initialize every part of each row together. let msp_out = &mut msp.spare_capacity_mut()[..len]; let mut lower_out = lower @@ -180,7 +210,7 @@ fn split_wide( } } } - Mask::AllFalse(_) => unreachable!("AllFalse case addressed above"), + Mask::AllFalse(_) => unreachable!("all-null arrays are handled by split_decimal"), } // SAFETY: the input and all output slices have len elements. Both branches @@ -197,7 +227,7 @@ fn split_wide( /// Extract the high signed word and low unsigned word of an `i128`. #[inline] -const fn i128_to_parts(value: i128) -> (i64, [u64; 1]) { +const fn i128_to_parts(value: i128) -> (i64, [u64; MAX_I128_LOWER_PARTS]) { #[expect( clippy::cast_possible_truncation, clippy::cast_sign_loss, @@ -208,7 +238,7 @@ const fn i128_to_parts(value: i128) -> (i64, [u64; 1]) { /// Extract the signed MSP and three unsigned lower words of an `i256`. #[inline] -const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_LOWER_PARTS]) { +const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_I256_LOWER_PARTS]) { let (low, high) = value.to_parts(); #[expect( clippy::cast_possible_truncation, 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 dd90dd7e197..fa04aceaf4d 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 @@ -4,10 +4,13 @@ use rstest::rstest; use vortex_array::VortexSessionExecute; use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Constant; use vortex_array::arrays::DecimalArray; use vortex_array::assert_arrays_eq; use vortex_array::dtype::DecimalDType; use vortex_array::dtype::i256; +use vortex_array::match_each_decimal_value_type; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::buffer; @@ -15,6 +18,47 @@ use vortex_error::VortexResult; use super::*; +#[rstest] +#[case::empty_non_nullable(0, Validity::NonNullable)] +#[case::empty_nullable(0, Validity::AllValid)] +#[case::empty_all_null(0, Validity::AllInvalid)] +#[case::all_null(3, Validity::AllInvalid)] +#[case::all_null_array(3, Validity::Array(BoolArray::from_iter([false; 3]).into_array()))] +fn test_split_without_valid_rows( + #[case] len: usize, + #[case] validity: Validity, + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + values_type: DecimalType, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = match_each_decimal_value_type!(values_type, |T| { + DecimalArray::new( + Buffer::::zeroed(len), + DecimalDType::new(T::MAX_PRECISION, 0), + validity, + ) + }); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.msp.is::()); + assert!(parts.lower_parts.iter().all(|part| part.is::())); + assert_eq!(parts.msp.len(), len); + assert_eq!( + parts.msp.dtype().nullability(), + decimal.dtype().nullability() + ); + let round_tripped = round_trip(decimal.clone())?; + assert_eq!(round_tripped.values_type(), values_type); + assert_arrays_eq!(decimal, round_tripped, &mut ctx); + Ok(()) +} + #[rstest] #[case::non_nullable(Validity::NonNullable)] #[case::all_valid(Validity::AllValid)] @@ -40,16 +84,19 @@ fn test_split_zeroes_null_words( let decimal = decimal .slice(3..len + 3)? .execute::(&mut ctx)?; + let mask = decimal.validity()?.execute_mask(len, &mut ctx)?; let expected = PrimitiveArray::new( - decimal - .validity()? - .execute_mask(len, &mut ctx)? - .iter() + mask.iter() .map(|valid| if valid { u64::MAX } else { 0 }) .collect::>(), Validity::NonNullable, ); let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), if wide_256 { 3 } else { 1 }); + assert_eq!( + parts.msp.dtype(), + &DType::Primitive(PType::I64, decimal.dtype().nullability()) + ); for lower in parts.lower_parts { assert_arrays_eq!(expected.clone(), lower, &mut ctx); } @@ -111,9 +158,8 @@ fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { } #[rstest] -fn test_split_narrow_decimal_has_no_lower_parts( - #[values(Validity::NonNullable, Validity::AllInvalid, Validity::from_iter([true, false, true]))] - validity: Validity, +fn test_split_narrow_decimal_reuses_values( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] validity: Validity, ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); From df7cbfce9194cf1f644eb4eec5efb900a153cd81 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 10 Sep 2026 16:27:16 -0400 Subject: [PATCH 7/9] comment fix Signed-off-by: Matt Katz --- .../decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 b88486f1556..12c5c7064e1 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 @@ -109,7 +109,7 @@ impl DecimalParts { /// /// # Errors /// -/// * If the array's validity cannot be derived or executed. +/// Returns an error if the array's validity cannot be derived or executed. pub fn split_decimal(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { let validity = decimal.validity()?; let len = decimal.len(); From b2e49fb63d81cd762a14da14737133d963b8c4fd Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 10 Sep 2026 18:36:47 -0400 Subject: [PATCH 8/9] reorg Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/array.rs | 848 +++++++++++++++++ .../src/decimal_byte_parts/assemble.rs | 383 ++++++++ .../src/decimal_byte_parts/limbs/tests.rs | 289 ------ .../src/decimal_byte_parts/mod.rs | 850 +----------------- .../decimal_byte_parts/prop_tests.rs} | 43 +- .../{limbs/mod.rs => split.rs} | 235 ++--- .../src/decimal_byte_parts/testing.rs | 16 +- 7 files changed, 1379 insertions(+), 1285 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs delete mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs rename encodings/decimal-byte-parts/{tests/props.rs => src/decimal_byte_parts/prop_tests.rs} (81%) rename encodings/decimal-byte-parts/src/decimal_byte_parts/{limbs/mod.rs => split.rs} (59%) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs new file mode 100644 index 00000000000..b76f39a8388 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -0,0 +1,848 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal byte-parts array types, validation, and VTable implementations. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +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::PType; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::serde::ArrayChildren; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityChild; +use vortex_array::vtable::ValidityVTableFromChild; +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 super::LOWER_PART_DTYPE; +use super::MAX_LOWER_PARTS; +use super::assemble::assemble_decimal; +use super::assemble::assemble_wide_decimal_value; +use super::rules::PARENT_RULES; + +/// A [`DecimalByteParts`]-encoded Vortex array. +pub type DecimalBytePartsArray = Array; + +#[derive(Clone, prost::Message)] +pub struct DecimalBytesPartsMetadata { + #[prost(enumeration = "PType", tag = "1")] + zeroth_child_ptype: i32, + #[prost(uint32, tag = "2")] + 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) + } +} + +/// 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 +/// +/// All parts live in slots, so the array carries no additional data. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsData; + +impl Display for DecimalBytePartsData { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl ArrayHash for DecimalBytePartsData { + fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} +} + +impl ArrayEq for DecimalBytePartsData { + fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { + true + } +} + +impl DecimalBytePartsData { + /// 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, + ) -> VortexResult<()> { + if !msp.dtype().is_signed_int() { + vortex_bail!("msp must be a signed integer array") + } + + let expected_dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + vortex_ensure!( + dtype == &expected_dtype, + "expected dtype {expected_dtype}, got {dtype}" + ); + vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); + + let lower_part_count = lower_parts.len(); + + vortex_ensure!( + lower_part_count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" + ); + 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(()) + } +} + +#[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 [`super::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 { + 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), + ) + } + + /// Construct a [`DecimalBytePartsArray`] from parts whose invariants are already established. + /// + /// # Safety + /// + /// The MSP must have a signed integer dtype (`i8`, `i16`, `i32`, or `i64`). There must be + /// at most [`MAX_LOWER_PARTS`] lower parts, each a non-nullable `u64` array with the same + /// length as the MSP. Lower parts are ordered most significant first. + pub(super) unsafe fn new_unchecked( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> DecimalBytePartsArray { + let len = msp.len(); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + // SAFETY: the caller guarantees the part types, lengths, and count. The slot builder + // fills every required slot, and the length and nullability come from the MSP. + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData) + .with_slots(slots), + ) + } + } +} + +impl VTable for DecimalByteParts { + type TypedArrayData = DecimalBytePartsData; + + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + 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, + slots.lower_parts.iter(), + *decimal_dtype, + dtype, + len, + ) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("DecimalBytePartsArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + vortex_ensure!( + array.lower_parts().is_empty(), + "serializing DecimalByteParts with lower parts is not supported" + ); + Ok(Some( + DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + _buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = DecimalBytesPartsMetadata::decode(metadata)?; + vortex_ensure!( + metadata.lower_part_count()? == 0, + "vortex.decimal_byte_parts must not carry lower parts" + ); + metadata.into_array_parts(dtype, len, children) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + DecimalBytePartsSlots::slot_name(idx) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + // Reassemble DecimalArray from split parts + // TODO(mk): handle constant array parts directly instead canonicalizing. + 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)) + } +} + +#[array_slots(DecimalByteParts)] +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, +} + +pub(crate) trait DecimalBytePartsArrayExt: DecimalBytePartsArraySlotsExt { + /// The decimal dtype of this array. + fn decimal_dtype(&self) -> DecimalDType { + *self + .as_ref() + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype") + } + + /// Rebuild the array by applying `f` to the MSP and every lower part, in slot order. + /// + /// 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 { + 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()) + } + + /// Rebuild the array with a replacement MSP, preserving its lower parts, precision and scale. + /// + /// 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(), + ) + } +} + +impl> DecimalBytePartsArrayExt for T {} + +impl OperationsVTable for DecimalByteParts { + fn scalar_at( + array: ArrayView<'_, DecimalByteParts>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let scalar = array.msp().execute_scalar(index, ctx)?; + + // 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"); + + 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 = 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))) + } +} + +impl ValidityChild for DecimalByteParts { + fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { + // validity stored in 0th child + array.msp().clone() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::Array; + use vortex_array::ArrayParts; + 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::DecimalType; + 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_array::vtable::VTable; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::DecimalByteParts; + use super::DecimalBytePartsArray; + use super::DecimalBytePartsArraySlotsExt; + use super::DecimalBytePartsData; + use crate::decimal_byte_parts::LOWER_PART_DTYPE; + use crate::decimal_byte_parts::MAX_LOWER_PARTS; + 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() { + let decimal_dtype = DecimalDType::new(8, 2); + let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); + let array = DecimalByteParts::try_new( + PrimitiveArray::new( + buffer![100i32, 200i32, 400i32], + Validity::Array(BoolArray::from_iter(vec![false, true, true]).into_array()), + ) + .into_array(), + decimal_dtype, + ) + .unwrap() + .into_array(); + + assert_eq!( + Scalar::null(dtype.clone()), + array + .execute_scalar(0, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Decimal(DecimalValue::I64(200))) + ) + .unwrap(), + array + .execute_scalar(1, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), + array + .execute_scalar(2, &mut array_session().create_execution_ctx()) + .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() + ); + } + + #[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(); + + 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/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs new file mode 100644 index 00000000000..5077c98e17c --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs @@ -0,0 +1,383 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reassembling decimal arrays and values from their parts. + +use std::ops::BitOr; +use std::ops::Shl; + +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::i256; +use vortex_array::match_each_signed_integer_ptype; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use super::LOWER_PART_BITS; +use super::LOWER_PART_DTYPE; +use super::MAX_LOWER_PARTS; + +/// Reassemble primitive arrays that constitute decimal byte parts into a canonical decimal array. +/// +/// The MSP must be signed. There must be between zero and three (inclusive) `u64` lower parts, ordered +/// most significant first. The lower parts must be non-nullable. Every input array must have the same length. +/// +/// With no lower parts, the MSP buffer is reused as the decimal values. One lower part +/// assembles into `i128`. Two or three lower parts assemble into `i256`. +/// +/// # Errors +/// +/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity +/// cannot be derived. +pub fn assemble_decimal( + msp: &PrimitiveArray, + lower_parts: &[PrimitiveArray], + decimal_dtype: DecimalDType, +) -> VortexResult { + let validity = msp.validity()?; + vortex_ensure!(msp.dtype().as_ptype().is_signed_int()); + + if lower_parts.is_empty() { + return Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { + // SAFETY: the buffer is typed by the array's own ptype, the decimal dtype is the + // array's, and the validity is taken from the same array. + unsafe { DecimalArray::new_unchecked(msp.to_buffer::

(), decimal_dtype, validity) } + })); + } + + let len = msp.len(); + let lower: Vec<&[u64]> = lower_parts + .iter() + .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, + "lower part has len {}, expected {len}", + part.len() + ); + Ok(part) + }) + .collect::>()?; + + Ok(match lower.as_slice() { + [first] => DecimalArray::new( + assemble_wide_decimal::(msp, [first]), + decimal_dtype, + validity, + ), + [first, second] => DecimalArray::new( + assemble_wide_decimal::(msp, [first, second]), + decimal_dtype, + validity, + ), + [first, second, third] => DecimalArray::new( + assemble_wide_decimal::(msp, [first, second, third]), + decimal_dtype, + validity, + ), + _ => vortex_bail!( + "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", + lower.len() + ), + }) +} + +/// Assemble a column of wide decimal values from the MSP and `K` lower-part columns. +/// +/// 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, +{ + let mut out = BufferMut::::with_capacity(msp.len()); + match_each_signed_integer_ptype!(msp.ptype(), |P| { + out.extend_trusted(msp.as_slice::

().iter().enumerate().map(|(row, value)| { + #[allow( + clippy::useless_conversion, + reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" + )] + 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 { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Constant; + 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::DecimalType; + use vortex_array::dtype::NativeDecimalType; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::match_each_decimal_value_type; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::assemble_decimal; + use crate::decimal_byte_parts::split_decimal; + + #[rstest] + #[case::empty_non_nullable(0, Validity::NonNullable)] + #[case::empty_nullable(0, Validity::AllValid)] + #[case::empty_all_null(0, Validity::AllInvalid)] + #[case::all_null(3, Validity::AllInvalid)] + #[case::all_null_array(3, Validity::Array(BoolArray::from_iter([false; 3]).into_array()))] + fn test_split_without_valid_rows( + #[case] len: usize, + #[case] validity: Validity, + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + values_type: DecimalType, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = match_each_decimal_value_type!(values_type, |T| { + DecimalArray::new( + Buffer::::zeroed(len), + DecimalDType::new(T::MAX_PRECISION, 0), + validity, + ) + }); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.msp.is::()); + assert!(parts.lower_parts.iter().all(|part| part.is::())); + assert_eq!(parts.msp.len(), len); + assert_eq!( + parts.msp.dtype().nullability(), + decimal.dtype().nullability() + ); + let round_tripped = round_trip(decimal.clone())?; + assert_eq!(round_tripped.values_type(), values_type); + assert_arrays_eq!(decimal, round_tripped, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::non_nullable(Validity::NonNullable)] + #[case::all_valid(Validity::AllValid)] + #[case::all_null(Validity::AllInvalid)] + #[case::mixed(Validity::from_iter((0..263).map(|i| i % 3 != 1)))] + #[case::sparse(Validity::from_iter((0..263).map(|i| i % 16 == 0)))] + #[case::null_prefix_and_suffix(Validity::from_iter((0..263).map(|i| (67..196).contains(&i))))] + fn test_split_zeroes_null_words( + #[case] validity: Validity, + #[values(false, true)] wide_256: bool, + #[values(0, 1, 63, 64, 65, 257)] len: usize, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = if wide_256 { + DecimalArray::new( + buffer![i256::from_i128(-1); 263], + DecimalDType::new(76, 2), + validity, + ) + } else { + DecimalArray::new(buffer![-1i128; 263], DecimalDType::new(38, 2), validity) + }; + let decimal = decimal + .slice(3..len + 3)? + .execute::(&mut ctx)?; + let mask = decimal.validity()?.execute_mask(len, &mut ctx)?; + let expected = PrimitiveArray::new( + mask.iter() + .map(|valid| if valid { u64::MAX } else { 0 }) + .collect::>(), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), if wide_256 { 3 } else { 1 }); + assert_eq!( + parts.msp.dtype(), + &DType::Primitive(PType::I64, decimal.dtype().nullability()) + ); + for lower in parts.lower_parts { + assert_arrays_eq!(expected.clone(), lower, &mut ctx); + } + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + Ok(()) + } + + fn round_trip(decimal: DecimalArray) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal, &mut ctx)?; + let msp = parts.msp.execute::(&mut ctx)?; + let lower = parts + .lower_parts + .into_iter() + .map(|part| part.execute::(&mut ctx)) + .collect::>>()?; + assemble_decimal(&msp, &lower, decimal.decimal_dtype()) + } + + #[rstest] + #[case::zero(0)] + #[case::one(1)] + #[case::minus_one(-1)] + #[case::limb_boundary(1i128 << 64)] + #[case::just_below_limb_boundary((1i128 << 64) - 1)] + #[case::negative_limb_boundary(-(1i128 << 64))] + #[case::max(i128::MAX)] + #[case::min(i128::MIN)] + fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) + } + + #[rstest] + #[case::zero(i256::ZERO)] + #[case::one(i256::ONE)] + #[case::minus_one(i256::ZERO - i256::ONE)] + #[case::max(i256::MAX)] + #[case::min(i256::MIN)] + #[case::word_1(i256::from_parts(1u128 << 64, 0))] + #[case::word_2(i256::from_parts(0, 1))] + #[case::word_3(i256::from_parts(0, 1i128 << 64))] + #[case::mixed(i256::from_parts(u128::MAX, -3))] + fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) + } + + #[rstest] + fn test_split_narrow_decimal_reuses_values( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] + validity: Validity, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.lower_parts.is_empty()); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); + let msp = parts.msp.execute::(&mut ctx)?; + assert_eq!( + msp.as_slice::().as_ptr(), + decimal.buffer::().as_ptr() + ); + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + 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, + #[values(0, 1, 3)] lower_len: 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] = + PrimitiveArray::new(buffer![0u64; lower_len], Validity::NonNullable); + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + assert!(assemble_decimal(&msp, &lower, dtype).is_err()); + } + + #[rstest] + fn test_assemble_i256_part_order_and_sign_extension( + #[values(false, true)] narrow_msp: bool, + #[values(2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![3i8, -3], Validity::NonNullable) + } else { + PrimitiveArray::new(buffer![3i64, -3], Validity::NonNullable) + }; + let lower = + [4u64, 1, 2].map(|word| PrimitiveArray::new(buffer![word; 2], Validity::NonNullable)); + let dtype = DecimalDType::new(76, 0); + let actual = assemble_decimal(&msp, &lower[3 - lower_count..], dtype)?; + let low = (1u128 << 64) | 2; + let expected = if lower_count == 2 { + buffer![i256::from_parts(low, 3), i256::from_parts(low, -3)] + } else { + buffer![ + i256::from_parts(low, (3i128 << 64) | 4), + i256::from_parts(low, (-3i128 << 64) | 4), + ] + }; + assert_arrays_eq!( + DecimalArray::new(expected, dtype, Validity::NonNullable), + actual, + &mut ctx + ); + Ok(()) + } +} 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 deleted file mode 100644 index fa04aceaf4d..00000000000 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs +++ /dev/null @@ -1,289 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use rstest::rstest; -use vortex_array::VortexSessionExecute; -use vortex_array::array_session; -use vortex_array::arrays::BoolArray; -use vortex_array::arrays::Constant; -use vortex_array::arrays::DecimalArray; -use vortex_array::assert_arrays_eq; -use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::i256; -use vortex_array::match_each_decimal_value_type; -use vortex_array::validity::Validity; -use vortex_buffer::Buffer; -use vortex_buffer::buffer; -use vortex_error::VortexResult; - -use super::*; - -#[rstest] -#[case::empty_non_nullable(0, Validity::NonNullable)] -#[case::empty_nullable(0, Validity::AllValid)] -#[case::empty_all_null(0, Validity::AllInvalid)] -#[case::all_null(3, Validity::AllInvalid)] -#[case::all_null_array(3, Validity::Array(BoolArray::from_iter([false; 3]).into_array()))] -fn test_split_without_valid_rows( - #[case] len: usize, - #[case] validity: Validity, - #[values( - DecimalType::I8, - DecimalType::I16, - DecimalType::I32, - DecimalType::I64, - DecimalType::I128, - DecimalType::I256 - )] - values_type: DecimalType, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = match_each_decimal_value_type!(values_type, |T| { - DecimalArray::new( - Buffer::::zeroed(len), - DecimalDType::new(T::MAX_PRECISION, 0), - validity, - ) - }); - let parts = split_decimal(&decimal, &mut ctx)?; - assert!(parts.msp.is::()); - assert!(parts.lower_parts.iter().all(|part| part.is::())); - assert_eq!(parts.msp.len(), len); - assert_eq!( - parts.msp.dtype().nullability(), - decimal.dtype().nullability() - ); - let round_tripped = round_trip(decimal.clone())?; - assert_eq!(round_tripped.values_type(), values_type); - assert_arrays_eq!(decimal, round_tripped, &mut ctx); - Ok(()) -} - -#[rstest] -#[case::non_nullable(Validity::NonNullable)] -#[case::all_valid(Validity::AllValid)] -#[case::all_null(Validity::AllInvalid)] -#[case::mixed(Validity::from_iter((0..263).map(|i| i % 3 != 1)))] -#[case::sparse(Validity::from_iter((0..263).map(|i| i % 16 == 0)))] -#[case::null_prefix_and_suffix(Validity::from_iter((0..263).map(|i| (67..196).contains(&i))))] -fn test_split_zeroes_null_words( - #[case] validity: Validity, - #[values(false, true)] wide_256: bool, - #[values(0, 1, 63, 64, 65, 257)] len: usize, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = if wide_256 { - DecimalArray::new( - buffer![i256::from_i128(-1); 263], - DecimalDType::new(76, 2), - validity, - ) - } else { - DecimalArray::new(buffer![-1i128; 263], DecimalDType::new(38, 2), validity) - }; - let decimal = decimal - .slice(3..len + 3)? - .execute::(&mut ctx)?; - let mask = decimal.validity()?.execute_mask(len, &mut ctx)?; - let expected = PrimitiveArray::new( - mask.iter() - .map(|valid| if valid { u64::MAX } else { 0 }) - .collect::>(), - Validity::NonNullable, - ); - let parts = split_decimal(&decimal, &mut ctx)?; - assert_eq!(parts.lower_parts.len(), if wide_256 { 3 } else { 1 }); - assert_eq!( - parts.msp.dtype(), - &DType::Primitive(PType::I64, decimal.dtype().nullability()) - ); - for lower in parts.lower_parts { - assert_arrays_eq!(expected.clone(), lower, &mut ctx); - } - assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); - Ok(()) -} - -fn round_trip(decimal: DecimalArray) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); - let parts = split_decimal(&decimal, &mut ctx)?; - let msp = parts.msp.execute::(&mut ctx)?; - let lower = parts - .lower_parts - .into_iter() - .map(|part| part.execute::(&mut ctx)) - .collect::>>()?; - assemble_decimal(&msp, &lower, decimal.decimal_dtype()) -} - -#[rstest] -#[case::zero(0)] -#[case::one(1)] -#[case::minus_one(-1)] -#[case::limb_boundary(1i128 << 64)] -#[case::just_below_limb_boundary((1i128 << 64) - 1)] -#[case::negative_limb_boundary(-(1i128 << 64))] -#[case::max(i128::MAX)] -#[case::min(i128::MIN)] -fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { - let decimal = DecimalArray::new( - Buffer::from(vec![value]), - DecimalDType::new(38, 2), - Validity::NonNullable, - ); - let round_tripped = round_trip(decimal)?; - assert_eq!(round_tripped.buffer::().as_slice(), &[value]); - Ok(()) -} - -#[rstest] -#[case::zero(i256::ZERO)] -#[case::one(i256::ONE)] -#[case::minus_one(i256::ZERO - i256::ONE)] -#[case::max(i256::MAX)] -#[case::min(i256::MIN)] -#[case::word_1(i256::from_parts(1u128 << 64, 0))] -#[case::word_2(i256::from_parts(0, 1))] -#[case::word_3(i256::from_parts(0, 1i128 << 64))] -#[case::mixed(i256::from_parts(u128::MAX, -3))] -fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { - let decimal = DecimalArray::new( - Buffer::from(vec![value]), - DecimalDType::new(76, 2), - Validity::NonNullable, - ); - let round_tripped = round_trip(decimal)?; - assert_eq!(round_tripped.buffer::().as_slice(), &[value]); - Ok(()) -} - -#[rstest] -fn test_split_narrow_decimal_reuses_values( - #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] validity: Validity, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); - let parts = split_decimal(&decimal, &mut ctx)?; - assert!(parts.lower_parts.is_empty()); - assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); - let msp = parts.msp.execute::(&mut ctx)?; - assert_eq!( - msp.as_slice::().as_ptr(), - decimal.buffer::().as_ptr() - ); - assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); - Ok(()) -} - -#[test] -fn test_split_i256_part_count_and_types() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = DecimalArray::new( - Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), - DecimalDType::new(76, 0), - Validity::NonNullable, - ); - let parts = split_decimal(&decimal, &mut ctx)?; - assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); - assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); - for part in &parts.lower_parts { - assert_eq!(part.dtype(), &LOWER_PART_DTYPE); - } - Ok(()) -} - -#[rstest] -fn test_split_i256_part_order( - #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] validity: Validity, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = DecimalArray::new( - buffer![ - i256::from_parts((2u128 << 64) | 3, (1i128 << 64) | 4), - i256::ZERO, - i256::from_parts((6u128 << 64) | 7, (-2i128 << 64) | 5), - ], - DecimalDType::new(76, 0), - validity.clone(), - ); - let parts = split_decimal(&decimal, &mut ctx)?; - assert_arrays_eq!( - PrimitiveArray::new(buffer![1i64, 0, -2], validity), - parts.msp, - &mut ctx - ); - assert_eq!(parts.lower_parts.len(), 3); - for (part, expected) in parts.lower_parts.into_iter().zip([ - buffer![4u64, 0, 5], - buffer![2u64, 0, 6], - buffer![3u64, 0, 7], - ]) { - assert_arrays_eq!( - PrimitiveArray::new(expected, Validity::NonNullable), - part, - &mut ctx - ); - } - 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, - #[values(0, 1, 3)] lower_len: 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] = PrimitiveArray::new(buffer![0u64; lower_len], Validity::NonNullable); - let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); - assert!(assemble_decimal(&msp, &lower, dtype).is_err()); -} - -#[rstest] -fn test_assemble_i256_part_order_and_sign_extension( - #[values(false, true)] narrow_msp: bool, - #[values(2, 3)] lower_count: usize, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let msp = if narrow_msp { - PrimitiveArray::new(buffer![3i8, -3], Validity::NonNullable) - } else { - PrimitiveArray::new(buffer![3i64, -3], Validity::NonNullable) - }; - let lower = - [4u64, 1, 2].map(|word| PrimitiveArray::new(buffer![word; 2], Validity::NonNullable)); - let dtype = DecimalDType::new(76, 0); - let actual = assemble_decimal(&msp, &lower[3 - lower_count..], dtype)?; - let low = (1u128 << 64) | 2; - let expected = if lower_count == 2 { - buffer![i256::from_parts(low, 3), i256::from_parts(low, -3)] - } else { - buffer![ - i256::from_parts(low, (3i128 << 64) | 4), - i256::from_parts(low, (-3i128 << 64) | 4), - ] - }; - assert_arrays_eq!( - DecimalArray::new(expected, dtype, Validity::NonNullable), - actual, - &mut ctx - ); - Ok(()) -} 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 b9babd52aaa..113ce2bb4fd 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -1,827 +1,55 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt::Display; -use std::fmt::Formatter; -use std::hash::Hasher; +//! Decimal byte-parts encoding. +//! +//! A `DecimalByteParts` array stores each value as a signed most significant part (MSP) +//! followed by `k` unsigned 64-bit lower parts ordered most significant first. The encoded +//! value is +//! +//! ```text +//! msp * 2^(64k) + Σ_{i; - -impl ArrayHash for DecimalBytePartsData { - fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} -} - -impl ArrayEq for DecimalBytePartsData { - fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { - true - } -} - -#[derive(Clone, prost::Message)] -pub struct DecimalBytesPartsMetadata { - #[prost(enumeration = "PType", tag = "1")] - zeroth_child_ptype: i32, - #[prost(uint32, tag = "2")] - 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) - } -} - -#[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; - - type OperationsVTable = Self; - type ValidityVTable = ValidityVTableFromChild; - - fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); - *ID - } - - fn validate( - &self, - _data: &Self::TypedArrayData, - dtype: &DType, - len: usize, - slots: &[Option], - ) -> VortexResult<()> { - 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, - slots.lower_parts.iter(), - *decimal_dtype, - dtype, - len, - ) - } - - fn nbuffers(_array: ArrayView<'_, Self>) -> usize { - 0 - } - - fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - vortex_panic!("DecimalBytePartsArray buffer index {idx} out of bounds") - } - - fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") - } - - fn with_buffers( - &self, - array: ArrayView<'_, Self>, - buffers: &[BufferHandle], - ) -> VortexResult> { - vortex_array::vtable::with_empty_buffers(self, array, buffers) - } - - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - vortex_ensure!( - array.lower_parts().is_empty(), - "serializing DecimalByteParts with lower parts is not supported" - ); - Ok(Some( - DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = DecimalBytesPartsMetadata::decode(metadata)?; - vortex_ensure!( - metadata.lower_part_count()? == 0, - "vortex.decimal_byte_parts must not carry lower parts" - ); - metadata.into_array_parts(dtype, len, children) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - DecimalBytePartsSlots::slot_name(idx) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) - } - - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - // 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)) - } + pub use super::assemble::assemble_decimal; } -#[array_slots(DecimalByteParts)] -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 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 -/// -/// All parts live in slots, so the array carries no additional data. -#[derive(Clone, Debug)] -pub struct DecimalBytePartsData; - -impl Display for DecimalBytePartsData { - fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { - Ok(()) - } -} - -impl DecimalBytePartsData { - /// 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, - ) -> VortexResult<()> { - if !msp.dtype().is_signed_int() { - vortex_bail!("decimal bytes parts, first part must be a signed array") - } - - let expected_dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); - vortex_ensure!( - dtype == &expected_dtype, - "expected dtype {expected_dtype}, got {dtype}" - ); - vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); - - let lower_part_count = lower_parts.len(); - - vortex_ensure!( - lower_part_count <= MAX_LOWER_PARTS, - "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" - ); - 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(()) - } -} - -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") - } - - /// Rebuild the array by applying `f` to the MSP and every lower part, in slot order. - /// - /// 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 { - 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()) - } - - /// Rebuild the array with a replacement MSP, preserving its lower parts, precision and scale. - /// - /// 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(), - ) - } -} - -impl> DecimalBytePartsArrayExt for T {} - -impl OperationsVTable for DecimalByteParts { - fn scalar_at( - array: ArrayView<'_, DecimalByteParts>, - index: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - let scalar = array.msp().execute_scalar(index, ctx)?; - - // 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"); - - 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 = 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))) - } -} - -impl ValidityChild for DecimalByteParts { - fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { - // validity stored in 0th child - array.msp().clone() - } -} - -#[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::DecimalType; - 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; +/// The maximum number of 64-bit lower parts an encoded `i128` decimal can carry. +const MAX_I128_LOWER_PARTS: usize = 1; - 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; +/// The maximum number of 64-bit lower parts an encoded `i256` decimal can carry. +const MAX_I256_LOWER_PARTS: usize = 3; - #[test] - fn test_scalar_at_decimal_parts() { - let decimal_dtype = DecimalDType::new(8, 2); - let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); - let array = DecimalByteParts::try_new( - PrimitiveArray::new( - buffer![100i32, 200i32, 400i32], - Validity::Array(BoolArray::from_iter(vec![false, true, true]).into_array()), - ) - .into_array(), - decimal_dtype, - ) - .unwrap() - .into_array(); +/// The maximum number of 64-bit lower parts an encoded decimal can carry. +const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS; - assert_eq!( - Scalar::null(dtype.clone()), - array - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new( - dtype.clone(), - Some(ScalarValue::Decimal(DecimalValue::I64(200))) - ) - .unwrap(), - array - .execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), - array - .execute_scalar(2, &mut array_session().create_execution_ctx()) - .unwrap() - ); - } +/// Number of bits stored in each lower part. +const LOWER_PART_BITS: usize = 64; - /// 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() - ); - } - - #[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(); - - 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(()) - } -} +/// Every lower part is a non-nullable `u64` primitive, since the MSP carries the sign +/// and validity. +const LOWER_PART_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable); diff --git a/encodings/decimal-byte-parts/tests/props.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs similarity index 81% rename from encodings/decimal-byte-parts/tests/props.rs rename to encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs index e33606b2880..5630d1706a8 100644 --- a/encodings/decimal-byte-parts/tests/props.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs @@ -1,18 +1,7 @@ // 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)] +//! Property tests for decimal byte-parts round trips. use hegel::TestCase; use hegel::generators as gs; @@ -28,11 +17,12 @@ 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; +use super::DecimalByteParts; +use super::DecimalBytePartsArray; +use super::testing::encode; + /// Largest magnitude a `Decimal(38, _)` can hold: 38 nines. const MAX_I128: i128 = 10i128.pow(38) - 1; @@ -46,21 +36,10 @@ const MAX_LEN: usize = 48; fn ctx() -> ExecutionCtx { let session = array_session(); - vortex_decimal_byte_parts::initialize(&session); + crate::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 { @@ -160,7 +139,10 @@ 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); + let round_tripped = canonicalize( + encode(&decimal).vortex_expect("encode").into_array(), + &mut ctx, + ); assert_eq!(round_tripped.values_type(), decimal.values_type()); assert_arrays_eq!(decimal, round_tripped, &mut ctx); @@ -178,7 +160,10 @@ fn encoded_survives_decode_then_encode(tc: TestCase) { let mut ctx = ctx(); let decoded = canonicalize(array.into_array(), &mut ctx); - let re_decoded = canonicalize(encode(&decoded, &mut ctx).into_array(), &mut ctx); + let re_decoded = canonicalize( + encode(&decoded).vortex_expect("encode").into_array(), + &mut ctx, + ); assert_arrays_eq!(decoded, re_decoded, &mut ctx); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs similarity index 59% rename from encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs rename to encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs index 12c5c7064e1..82cb56ef239 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -1,21 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Splitting decimal values into 64-bit parts and reassembling them. -//! -//! A `DecimalByteParts` array stores each value as a signed most significant part (MSP) -//! followed by `k` unsigned 64-bit lower parts ordered most significant first. The encoded -//! value is -//! -//! ```text -//! msp * 2^(64k) + Σ_{i VortexResult { + let parts = split_decimal(decimal, exec_ctx)?; + // SAFETY: splitting produces a signed MSP and zero, one, or three non-nullable u64 lower + // parts, all with the decimal's length and in most-significant-first order. This also holds + // for the constant parts used for empty and all-null inputs. The decimal dtype is preserved. + Ok(unsafe { + DecimalByteParts::new_unchecked(parts.msp, parts.lower_parts, decimal.decimal_dtype()) + }) +} /// A decimal array decomposed into byte parts. pub struct DecimalParts { @@ -80,7 +64,7 @@ impl DecimalParts { } /// Construct decimal parts arrays from the buffers constituting a wide decimal (`i128` or `i256`). - /// Wide decimals have an `i64` MSP and up to [`MAX_LOWER_PARTS`] `u64` lower parts. + /// Wide decimals have an `i64` MSP and up to [`super::MAX_LOWER_PARTS`] `u64` lower parts. fn from_wide( msp: Buffer, lower_parts: impl IntoIterator>, @@ -251,115 +235,76 @@ const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_I256_LOWER_PARTS]) { ) } -/// Reassemble primitive arrays that constitute decimal byte parts into a canonical decimal array. -/// -/// The MSP must be signed. There must be between zero and three (inclusive) `u64` lower parts, ordered -/// most significant first. The lower parts must be non-nullable. Every input array must have the same length. -/// -/// With no lower parts, the MSP buffer is reused as the decimal values. One lower part -/// assembles into `i128`. Two or three lower parts assemble into `i256`. -/// -/// # Errors -/// -/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity -/// cannot be derived. -pub fn assemble_decimal( - msp: &PrimitiveArray, - lower_parts: &[PrimitiveArray], - decimal_dtype: DecimalDType, -) -> VortexResult { - let validity = msp.validity()?; - vortex_ensure!(msp.dtype().as_ptype().is_signed_int()); +#[cfg(test)] +mod tests { + use rstest::rstest; + 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::PType; + use vortex_array::dtype::i256; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::split_decimal; + use crate::decimal_byte_parts::LOWER_PART_DTYPE; + use crate::decimal_byte_parts::MAX_LOWER_PARTS; - if lower_parts.is_empty() { - return Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { - // SAFETY: the buffer is typed by the array's own ptype, the decimal dtype is the - // array's, and the validity is taken from the same array. - unsafe { DecimalArray::new_unchecked(msp.to_buffer::

(), decimal_dtype, validity) } - })); + #[test] + fn test_split_i256_part_count_and_types() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), + DecimalDType::new(76, 0), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); + for part in &parts.lower_parts { + assert_eq!(part.dtype(), &LOWER_PART_DTYPE); + } + Ok(()) } - let len = msp.len(); - let lower: Vec<&[u64]> = lower_parts - .iter() - .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, - "lower part has len {}, expected {len}", - part.len() + #[rstest] + fn test_split_i256_part_order( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] + validity: Validity, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + buffer![ + i256::from_parts((2u128 << 64) | 3, (1i128 << 64) | 4), + i256::ZERO, + i256::from_parts((6u128 << 64) | 7, (-2i128 << 64) | 5), + ], + DecimalDType::new(76, 0), + validity.clone(), + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_arrays_eq!( + PrimitiveArray::new(buffer![1i64, 0, -2], validity), + parts.msp, + &mut ctx + ); + assert_eq!(parts.lower_parts.len(), 3); + for (part, expected) in parts.lower_parts.into_iter().zip([ + buffer![4u64, 0, 5], + buffer![2u64, 0, 6], + buffer![3u64, 0, 7], + ]) { + assert_arrays_eq!( + PrimitiveArray::new(expected, Validity::NonNullable), + part, + &mut ctx ); - Ok(part) - }) - .collect::>()?; - - Ok(match lower.as_slice() { - [first] => DecimalArray::new( - assemble_wide_decimal::(msp, [first]), - decimal_dtype, - validity, - ), - [first, second] => DecimalArray::new( - assemble_wide_decimal::(msp, [first, second]), - decimal_dtype, - validity, - ), - [first, second, third] => DecimalArray::new( - assemble_wide_decimal::(msp, [first, second, third]), - decimal_dtype, - validity, - ), - _ => vortex_bail!( - "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", - lower.len() - ), - }) -} - -/// Assemble a column of wide decimal values from the MSP and `K` lower-part columns. -/// -/// 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, -{ - let mut out = BufferMut::::with_capacity(msp.len()); - match_each_signed_integer_ptype!(msp.ptype(), |P| { - out.extend_trusted(msp.as_slice::

().iter().enumerate().map(|(row, value)| { - #[allow( - clippy::useless_conversion, - reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" - )] - 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"); + } + Ok(()) } - value } - -#[cfg(test)] -mod tests; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs index d2ce68f3700..2dfe2a55b3c 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Test-only helpers for building byte-parts arrays. +//! Shared fixtures for decimal byte-parts tests. use vortex_array::VortexSessionExecute; use vortex_array::array_session; @@ -13,18 +13,12 @@ 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; +use super::DecimalBytePartsArray; +use super::dbp_encode; /// 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(), - ) + dbp_encode(decimal, &mut array_session().create_execution_ctx()) } /// An `i128`-backed decimal array, encoded as byte parts with one lower part. @@ -48,6 +42,6 @@ pub(crate) fn i256_parts(values: Vec, validity: Validity) -> DecimalBytePa } /// Build an `i256` from a signed high `i128` and unsigned low `u128`. -pub(crate) fn i256_of(high: i128, low: u128) -> i256 { +pub(super) fn i256_of(high: i128, low: u128) -> i256 { i256::from_parts(low, high) } From 18b8b996ca023180178c18eefdafdf5e50620f05 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 11 Sep 2026 14:31:24 -0400 Subject: [PATCH 9/9] comments Signed-off-by: Matt Katz --- .../src/decimal_byte_parts/assemble.rs | 10 ++++------ .../src/decimal_byte_parts/compute/cast.rs | 18 +++--------------- .../src/decimal_byte_parts/compute/compare.rs | 3 +++ .../src/decimal_byte_parts/compute/take.rs | 2 ++ vortex-array/src/dtype/bigint/mod.rs | 14 ++++++++++++++ 5 files changed, 26 insertions(+), 21 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs index 5077c98e17c..5fcf2bcdb60 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs @@ -14,7 +14,6 @@ use vortex_array::dtype::i256; use vortex_array::match_each_signed_integer_ptype; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -99,7 +98,7 @@ pub fn assemble_decimal( /// 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, + T: NativeDecimalType + From + From + Shl + BitOr, { let mut out = BufferMut::::with_capacity(msp.len()); match_each_signed_integer_ptype!(msp.ptype(), |P| { @@ -123,12 +122,11 @@ where #[inline] pub(crate) fn assemble_wide_decimal_value(msp: i64, lower: [u64; K]) -> T where - T: NativeDecimalType + Shl + BitOr, + T: NativeDecimalType + From + From + Shl + BitOr, { - let mut value = T::from(msp).vortex_expect("MSP fits in the output type"); + let mut value: T = msp.into(); for part in lower { - value = (value << LOWER_PART_BITS) - | T::from(part).vortex_expect("lower word fits in the output type"); + value = (value << LOWER_PART_BITS) | part.into(); } value } 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 0594f15a7c6..e4d81c274da 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 @@ -15,13 +15,13 @@ use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl CastReduce for DecimalByteParts { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { - // Check if this is just a nullability change + // Check if this is just a nullability change. + // TODO(mk): Support non-nullability changes as well, e.g. precision. if !dtype.eq_ignore_nullability(array.dtype()) { return Ok(None); } - // DecimalBytePartsArray can only have Decimal dtype, so we only handle decimal-to-decimal casts + // DecimalBytePartsArray can only have Decimal dtype. let DType::Decimal(_, target_nullability) = dtype else { - // Cannot cast decimal to non-decimal types - delegate to canonical form return Ok(None); }; @@ -48,14 +48,10 @@ 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() { @@ -120,14 +116,6 @@ 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 fe4d69801a3..65933865780 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 @@ -41,6 +41,9 @@ impl CompareKernel for DecimalByteParts { // The MSP alone only determines the ordering when it holds the whole value. With // lower parts present, fall back to comparing the canonical decimal. + // + // TODO(mk): Compare the signed MSP and then the unsigned lower parts in significance + // order to avoid canonicalizing wide decimals. if !lhs.lower_parts().is_empty() { return Ok(None); } 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 9e504433a2b..5018a99e3bf 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 @@ -16,6 +16,8 @@ impl TakeReduce for DecimalByteParts { // 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. + // + // TODO(mk): Support lower parts using fill_null for nullable indices. if indices.dtype().is_nullable() && !array.lower_parts().is_empty() { return Ok(None); } diff --git a/vortex-array/src/dtype/bigint/mod.rs b/vortex-array/src/dtype/bigint/mod.rs index 47195526b1f..b03bdef43c2 100644 --- a/vortex-array/src/dtype/bigint/mod.rs +++ b/vortex-array/src/dtype/bigint/mod.rs @@ -130,6 +130,20 @@ impl From for i256 { } } +impl From for i256 { + #[inline] + fn from(value: i64) -> Self { + Self::from_i128(value.into()) + } +} + +impl From for i256 { + #[inline] + fn from(value: u64) -> Self { + Self::from_i128(value.into()) + } +} + impl Display for i256 { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.0)