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 27cf5957fd4..b2622d48f29 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2046,6 +2046,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" @@ -3846,6 +3865,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.2" @@ -5947,6 +6011,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" @@ -9992,8 +10062,11 @@ dependencies = [ name = "vortex-decimal-byte-parts" version = "0.1.0" dependencies = [ + "codspeed-divan-compat", + "hegeltest", "num-traits", "prost 0.14.4", + "rand 0.10.2", "rstest", "vortex-array", "vortex-buffer", @@ -10119,6 +10192,7 @@ dependencies = [ "oneshot", "parking_lot", "pin-project-lite", + "rand 0.10.2", "rstest", "tokio", "tracing", diff --git a/Cargo.toml b/Cargo.toml index 00a45b03618..173186218b1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -170,6 +170,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 4934ec4fa27..5a9849266fb 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -26,5 +26,12 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } +hegeltest = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } + +[[bench]] +name = "decimal_assemble" +harness = false diff --git a/encodings/decimal-byte-parts/benches/decimal_assemble.rs b/encodings/decimal-byte-parts/benches/decimal_assemble.rs new file mode 100644 index 00000000000..193149a7c68 --- /dev/null +++ b/encodings/decimal-byte-parts/benches/decimal_assemble.rs @@ -0,0 +1,106 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Canonicalizing `DecimalByteParts` into `i128`/`i256` values. +//! +//! Reassembly walks a most significant part plus one (`i128`) or three (`i256`) unsigned +//! 64-bit lower parts and produces one wide value per row. This benchmark measures the +//! shipped path through the array API, so it tracks whatever shape the crate currently uses +//! and cannot drift away from it. +//! +//! Alternative shapes were compared while choosing that path and then removed, since keeping +//! hand-written copies of the assembly loop here means maintaining the same loop twice. At +//! 65,536 rows, `fastest` of three runs: +//! +//! - **`i256` is dominated by the part count being visible to the compiler.** Specializing it +//! to a constant is 1.85x (190 µs against 351 µs). How the output is written barely matters +//! at 32 bytes per row. +//! - **`i128` is dominated by the write.** Specializing the part count is worth only ~1.04x, +//! while storing into a pre-sized buffer instead of pushing into a reserved one is 1.6x +//! (83 µs against 138 µs) — the bounds-checked `push` is the whole cost at 16 bytes per row. +//! - **Columnar always loses.** For `i256` each lane store is strided by 32 bytes, 2.3x slower +//! than the row loop (438 µs); cache blocking the passes recovered part of that and was +//! still 1.6x slower; expressing them as whole-value `i256` shifts was 11x slower. For +//! `i128` the two-pass column shape (103 µs) beats the *pushing* row loop but still loses to +//! the single-pass write, so the second pass buys nothing once the push is gone. +//! - **Hand-written 64-bit words do not beat the `u128` packing.** `i256::from_parts` takes a +//! `u128` and an `i128`, so each row ends in `u128::from(w0) | (u128::from(w1) << 64)`. +//! Writing four `u64` lanes by hand instead ties it. Disassembly says why: neither emits a +//! single `shld`/`shrd`, and both compile to four plain 64-bit stores per row at offsets +//! 0x0/0x8/0x10/0x18. The `i128` loop is the same — `(i128::from(msp) << 64) | i128::from(p)` +//! becomes two 64-bit stores. A shift by a constant multiple of 64 followed by an or is pure +//! data movement and LLVM recognizes it; the 128-bit codegen worth avoiding is division and +//! remainder, which call into compiler-rt, and shifts by a runtime amount. Neither is here. +//! +//! The removed variants are recoverable from git history if a future change needs to re-run +//! the comparison rather than trust these numbers. + +#![allow(clippy::unwrap_used, clippy::cast_possible_truncation)] + +use divan::Bencher; +use divan::black_box; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_decimal_byte_parts::DecimalByteParts; + +fn main() { + divan::main(); +} + +/// Rows per benchmark: a typical scan chunk, and large enough that the output does not fit +/// in L2. +const LEN: usize = 65_536; + +/// Deterministic pseudo-random words, so no part is constant or a sequence. +fn words(seed: u64, len: usize) -> Buffer { + let mut rng = StdRng::seed_from_u64(seed); + (0..len).map(|_| rng.random()).collect() +} + +fn msp(seed: u64, len: usize) -> Buffer { + words(seed, len) + .iter() + .map(|w| (w >> 40).cast_signed()) + .collect() +} + +/// Canonicalizing through the public array API, so the child execution and validity handling +/// around the assembly loop are included. +#[divan::bench(args = [1, 3])] +fn canonicalize_byte_parts(bencher: Bencher, lower_parts: usize) { + let msp = PrimitiveArray::new(msp(1, LEN), Validity::NonNullable); + let lower = (0..lower_parts) + .map(|i| PrimitiveArray::new(words(7 + i as u64, LEN), Validity::NonNullable)) + .collect::>(); + + let dtype = if lower_parts == 1 { + DecimalDType::new(38, 2) + } else { + DecimalDType::new(76, 2) + }; + let array = DecimalByteParts::try_new_with_lower_parts( + msp.into_array(), + lower.into_iter().map(IntoArray::into_array).collect(), + dtype, + ) + .unwrap() + .into_array(); + + let session = array_session(); + bencher + .with_inputs(|| session.create_execution_ctx()) + .bench_refs(|ctx| { + black_box(array.clone()) + .execute::(ctx) + .unwrap() + }); +} 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 2abf347bc7a..665b0f26ae2 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 @@ -8,28 +8,22 @@ use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; use vortex_array::scalar_fn::EmptyOptions; 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( array.msp().len(), EmptyOptions, [array.msp().clone(), mask.clone()], )?; - 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.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs new file mode 100644 index 00000000000..a21d62d046d --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs.rs @@ -0,0 +1,457 @@ +// 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, +} + +/// The decimal storage type that reassembling the given parts produces. +/// +/// # Errors +/// +/// Returns an error if `msp_ptype` is not a signed integer, or if there are more than +/// [`MAX_LOWER_PARTS`] lower parts. +pub(crate) fn assembled_values_type( + msp_ptype: PType, + lower_part_count: usize, +) -> VortexResult { + if lower_part_count > MAX_LOWER_PARTS { + vortex_bail!("at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}"); + } + if lower_part_count == 0 { + return DecimalType::try_from(msp_ptype); + } + let bits = msp_ptype.bit_width() + LOWER_PART_BITS * lower_part_count; + Ok(if bits <= 128 { + DecimalType::I128 + } else { + DecimalType::I256 + }) +} + +/// Split a canonical decimal array into a signed most significant part and unsigned 64-bit +/// lower parts. +/// +/// Values narrower than 128 bits are already a single signed part, so they are returned +/// with no lower parts. `i128` values split into an `i64` MSP and one lower part, `i256` +/// values into an `i64` MSP and three lower parts. +/// +/// # Errors +/// +/// Returns an error if the array's validity cannot be derived. +pub fn split_decimal(decimal: &DecimalArray) -> VortexResult { + let validity = decimal.validity()?; + Ok(match decimal.values_type() { + DecimalType::I8 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I16 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I32 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I64 => DecimalParts::flat(decimal.buffer::(), validity), + DecimalType::I128 => { + let (msp, lower) = split_i128(&decimal.buffer::()); + DecimalParts::new(msp, [lower], validity) + } + DecimalType::I256 => { + let (msp, lower) = split_i256(&decimal.buffer::()); + DecimalParts::new(msp, lower, validity) + } + }) +} + +/// Reassemble decimal byte parts into a canonical decimal array. +/// +/// The parts must already be canonical primitive arrays: a signed MSP, and `u64` lower +/// parts ordered most significant first. +/// +/// # Errors +/// +/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity +/// cannot be derived. +pub(crate) fn assemble_decimal( + msp: &PrimitiveArray, + lower_parts: &[PrimitiveArray], + decimal_dtype: DecimalDType, +) -> VortexResult { + let validity = msp.validity()?; + 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) } + })); + } + + // Slice every part to the MSP's length up front: the assembly loops then index slices the + // compiler knows are long enough, so the per-row bounds checks fall away. + let len = msp.len(); + let lower: Vec<&[u64]> = lower_parts + .iter() + .map(|part| { + let part = part.as_slice::(); + vortex_ensure!( + part.len() >= len, + "lower part has len {}, expected at least {len}", + part.len() + ); + Ok(&part[..len]) + }) + .collect::>()?; + + // The part count is dispatched to a constant so every 64-bit word lands at a compile-time + // index. Leaving it dynamic costs 1.8x on the `i256` path — see `benches/decimal_assemble.rs`. + let values = match assembled_values_type(msp.ptype(), lower.len())? { + // A single lower part can never widen to an `i256`: the MSP is at most 64 bits, so + // 64 + 64 fits an `i128` and takes the branch below. + DecimalType::I256 => match lower.as_slice() { + [first, second] => assemble_i256(msp, [first, second]), + [first, second, third] => assemble_i256(msp, [first, second, third]), + _ => vortex_bail!("unsupported lower part count {}", lower.len()), + }, + _ => { + return Ok(DecimalArray::new( + assemble_i128(msp, lower[0]), + decimal_dtype, + validity, + )); + } + }; + Ok(DecimalArray::new(values, decimal_dtype, validity)) +} + +/// 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) + }) +} + +/// 64-bit words in an `i256`. +const VALUE_WORDS: usize = 4; + +/// The 64-bit words of an `i256`, ascending significance. +/// +/// An `i256` is exactly `{_0: u64, _1: u64, _2: u64, _3: i64}` — three unsigned words beneath +/// a single signed one — which is the same shape this encoding stores. That is why splitting +/// and reassembling are pure reinterpretation rather than arithmetic: no carry ever crosses a +/// word boundary, so each word can be compressed independently and put back verbatim. +/// +/// The sign lives in the most significant word alone. When the most significant part is +/// narrower than 64 bits, or sits below word 3, the words above it are its sign extension. +type ValueWords = [u64; VALUE_WORDS]; + +/// Reinterpret an `i256` as its 64-bit words. +#[inline] +const fn i256_to_words(value: i256) -> ValueWords { + let (low, high) = value.to_parts(); + #[expect( + clippy::cast_possible_truncation, + reason = "each cast takes the low 64 bits of a word pair by construction" + )] + [ + low as u64, + (low >> LOWER_PART_BITS) as u64, + high as u64, + (high >> LOWER_PART_BITS) as u64, + ] +} + +/// Reinterpret 64-bit words as an `i256`, with the most significant word carrying the sign. +#[inline] +const fn i256_from_words(words: ValueWords) -> i256 { + i256::from_parts( + (words[0] as u128) | ((words[1] as u128) << LOWER_PART_BITS), + ((words[2] as u128) | ((words[3] as u128) << LOWER_PART_BITS)) as i128, + ) +} + +/// The words of a value whose most significant part sits at `msp_word`, with every word above +/// it filled with the MSP's sign. +#[inline] +fn sign_extended_words(msp: i64, msp_word: usize) -> ValueWords { + let mut words = [if msp < 0 { u64::MAX } else { 0 }; VALUE_WORDS]; + words[msp_word] = msp.cast_unsigned(); + words +} + +/// Combine a single row's parts into an `i256`. +/// +/// The lower parts fill the least significant 64-bit words, the MSP the word above them, +/// and the remaining high words are the MSP's sign extension. +#[inline] +pub(crate) fn combine_i256(msp: i64, lower: impl ExactSizeIterator) -> i256 { + let count = lower.len(); + let mut words = sign_extended_words(msp, count); + for (i, part) in lower.enumerate() { + words[count - 1 - i] = part; + } + i256_from_words(words) +} + +impl DecimalParts { + /// Parts for a decimal already stored in a single signed integer. + fn flat(values: Buffer, validity: Validity) -> Self { + Self { + msp: PrimitiveArray::new(values, validity).into_array(), + lower_parts: Vec::new(), + } + } + + fn new( + msp: Buffer, + lower_parts: impl IntoIterator>, + validity: Validity, + ) -> Self { + Self { + msp: PrimitiveArray::new(msp, validity).into_array(), + lower_parts: lower_parts + .into_iter() + .map(|part| PrimitiveArray::new(part, Validity::NonNullable).into_array()) + .collect(), + } + } +} + +#[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "splitting a wide integer into 64-bit windows truncates by construction" +)] +fn split_i128(values: &Buffer) -> (Buffer, Buffer) { + let mut msp = BufferMut::::with_capacity(values.len()); + let mut lower = BufferMut::::with_capacity(values.len()); + for value in values.iter() { + msp.push((value >> LOWER_PART_BITS) as i64); + lower.push(*value as u64); + } + (msp.freeze(), lower.freeze()) +} + +/// The inverse of [`assemble_i256`] at `K == MAX_LOWER_PARTS`: word 3 becomes the signed MSP, +/// and words 2, 1, 0 become the lower parts, most significant first. +fn split_i256(values: &Buffer) -> (Buffer, [Buffer; MAX_LOWER_PARTS]) { + let mut msp = BufferMut::::with_capacity(values.len()); + let mut lower = std::array::from_fn::<_, MAX_LOWER_PARTS, _>(|_| { + BufferMut::::with_capacity(values.len()) + }); + for value in values.iter() { + let words = i256_to_words(*value); + msp.push(words[MAX_LOWER_PARTS].cast_signed()); + for (part, word) in lower + .iter_mut() + .zip(words.iter().take(MAX_LOWER_PARTS).rev()) + { + part.push(*word); + } + } + (msp.freeze(), lower.map(BufferMut::freeze)) +} + +/// Only one lower part can share 128 bits with a signed MSP, so this shape is fixed. +#[expect( + clippy::useless_conversion, + reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" +)] +fn assemble_i128(msp: &PrimitiveArray, lower: &[u64]) -> Buffer { + // Store into a pre-sized buffer rather than pushing into a reserved one: at 16 bytes per + // row the bounds-checked `push` dominates, and dropping it is 1.6x — see + // `i128_row_write` against `i128_row_const` in `benches/decimal_assemble.rs`. The same + // shape does not pay off for `i256`, where zeroing 32 bytes per row costs more than the + // push it saves. + let mut out = BufferMut::::zeroed(msp.len()); + match_each_signed_integer_ptype!(msp.ptype(), |P| { + for ((slot, value), part) in out + .as_mut_slice() + .iter_mut() + .zip(msp.as_slice::

()) + .zip(lower) + { + *slot = (i128::from(i64::from(*value)) << LOWER_PART_BITS) | i128::from(*part); + } + }); + out.freeze() +} + +/// The lower parts fill the least significant 64-bit words, the MSP the word above them, and +/// the remaining high words are the MSP's sign extension. +/// +/// `K` is a constant so the word indices are compile-time constants and the placement loop +/// unrolls; the same loop with a runtime part count is 1.8x slower. +#[expect( + clippy::useless_conversion, + reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" +)] +fn assemble_i256(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer { + let mut out = BufferMut::::with_capacity(msp.len()); + match_each_signed_integer_ptype!(msp.ptype(), |P| { + for (row, value) in msp.as_slice::

().iter().enumerate() { + // The MSP occupies word `K`, the lower parts the `K` words beneath it most + // significant first, and anything above word `K` is the MSP's sign. + let mut words = sign_extended_words(i64::from(*value), K); + for (i, part) in lower.iter().enumerate() { + words[K - 1 - i] = part[row]; + } + out.push(i256_from_words(words)); + } + }); + out.freeze() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + 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_buffer::buffer; + use vortex_error::VortexResult; + + use super::*; + + fn round_trip(decimal: DecimalArray) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal)?; + 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(()) + } + + #[test] + fn test_split_narrow_decimal_has_no_lower_parts() -> VortexResult<()> { + let decimal = DecimalArray::new( + buffer![1i32, 2, 3], + DecimalDType::new(9, 2), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal)?; + assert!(parts.lower_parts.is_empty()); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); + Ok(()) + } + + #[test] + fn test_split_i256_part_count_and_types() -> VortexResult<()> { + 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)?; + 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(()) + } + + #[test] + fn test_assembled_values_type() -> VortexResult<()> { + assert_eq!(assembled_values_type(PType::I32, 0)?, DecimalType::I32); + assert_eq!(assembled_values_type(PType::I64, 1)?, DecimalType::I128); + assert_eq!(assembled_values_type(PType::I8, 1)?, DecimalType::I128); + assert_eq!(assembled_values_type(PType::I8, 2)?, DecimalType::I256); + assert_eq!(assembled_values_type(PType::I64, 3)?, DecimalType::I256); + assert!(assembled_values_type(PType::I64, 4).is_err()); + 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 d5b0024f5b7..459529d44b6 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -9,31 +9,37 @@ use vortex_array::Array; use vortex_array::ArrayParts; use vortex_array::ArrayView; pub(crate) mod compute; +mod limbs; 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; use prost::Message as _; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; +use vortex_array::ArrayPlugin; 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; @@ -42,10 +48,16 @@ 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::assembled_values_type; +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. @@ -69,6 +81,23 @@ pub struct DecimalBytesPartsMetadata { lower_part_count: u32, } +impl DecimalBytesPartsMetadata { + /// 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; @@ -90,8 +119,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 { @@ -118,15 +153,30 @@ impl VTable for DecimalByteParts { array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { + let lower_part_count = u32::try_from(array.lower_parts().len()) + .map_err(|_| vortex_err!("lower part count exceeds u32"))?; Ok(Some( DecimalBytesPartsMetadata { zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, - lower_part_count: 0, + lower_part_count, } .encode_to_vec(), )) } + fn serialized_id(&self, array: ArrayView<'_, Self>) -> ArrayId { + // The frozen `vortex.decimal_byte_parts` format promises a single child: readers that + // predate lower parts require `lower_part_count == 0`, so an array carrying lower + // parts must serialize under the v2 format id instead. That id is what the writer's + // permitted-encoding check gates, and what a reader without v2 support rejects as an + // unknown encoding instead of misreading the children. + if array.lower_parts().is_empty() { + VTable::id(self) + } else { + ArrayPlugin::id(&DecimalBytePartsV2) + } + } + fn deserialize( &self, dtype: &DType, @@ -137,26 +187,41 @@ 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) - }; + vortex_ensure!( + dtype.as_decimal_opt().is_some(), + "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" + let lower_part_count = metadata.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 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)) + 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(self.clone(), dtype.clone(), len, DecimalBytePartsData) + .with_slots(slots), + ) } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - DecimalBytePartsSlots::NAMES[idx].to_string() + DecimalBytePartsSlots::slot_name(idx) } fn reduce_parent( @@ -177,20 +242,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 | 64..0 \] msp = 127..64 and lower_part\[0\] = 64..0 +/// 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 { - // 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 { @@ -198,13 +264,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, 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, @@ -219,24 +289,35 @@ 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(); + 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() + ); + } + // Rejects part combinations that cannot be reassembled into a decimal value. This also + // bounds the lower part count. + let values_type = assembled_values_type(msp.dtype().as_ptype(), lower_part_count)?; + + // The parts must not assemble into a wider value than the declared precision holds. + // Without this, a crafted array carrying more parts than its precision needs + // canonicalizes to out-of-precision values that then panic in the scalar path. + let widest = DecimalType::smallest_decimal_value_type(&decimal_dtype); vortex_ensure!( - msp_dtype.is_signed_int(), - "decimal bytes parts, first part must be a signed array" + values_type <= widest, + "parts assemble into {values_type:?}, wider than the {widest:?} required by \ + decimal precision {}", + decimal_dtype.precision() ); - let _ = msp_len; - drop(expected_dtype); - Ok(Self { - _lower_parts: Vec::new(), - }) + Ok(()) } } @@ -245,47 +326,157 @@ 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 { + // Building lower parts in memory is never gated — reading a file requires it. What is + // gated is the serialized form: an array carrying lower parts serializes under the + // `DecimalBytePartsV2` format id, which only editions that contain it may write. 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 `vortex.decimal_byte_parts_v2` serialized format: byte parts carrying lower parts. +/// +/// This is a serialized format, not a second in-memory encoding. `vortex.decimal_byte_parts` +/// froze promising a single child, so an array with lower parts serializes under this id +/// instead — see [`VTable::serialized_id`] on [`DecimalByteParts`] — and both ids deserialize +/// back into the same [`DecimalBytePartsArray`]. A reader that predates lower parts fails on +/// this id with an unknown-encoding error rather than misreading the children. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsV2; + +impl ArrayPlugin for DecimalBytePartsV2 { + fn id(&self) -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts_v2"); + *ID + } + + fn serialize( + &self, + array: &ArrayRef, + session: &VortexSession, + ) -> VortexResult>> { + // In-memory arrays always carry the `DecimalByteParts` encoding id, so metadata + // serialization is resolved through that plugin; both formats share it. + ArrayPlugin::serialize(&DecimalByteParts, array, session) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + buffers: &[BufferHandle], + children: &dyn ArrayChildren, + session: &VortexSession, + ) -> VortexResult { + Ok(Array::try_from_parts(VTable::deserialize( + &DecimalByteParts, + dtype, + len, + metadata, + buffers, + children, + session, + )?)? + .into_array()) + } + + fn is_supported_encoding(&self, id: &ArrayId) -> bool { + *id == ArrayPlugin::id(self) || *id == VTable::id(&DecimalByteParts) + } +} + +/// The decimal storage type this array canonicalizes to. +fn values_type(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { + assembled_values_type(array.msp().dtype().as_ptype(), array.lower_parts().len()) +} + +/// 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 { @@ -294,17 +485,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))) } } @@ -317,21 +525,39 @@ impl ValidityChild for DecimalByteParts { #[cfg(test)] mod tests { + use rstest::rstest; + use vortex_array::ArrayContext; + 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::serde::SerializeOptions; + use vortex_array::serde::SerializedArray; + use vortex_array::session::ArraySessionExt; use vortex_array::validity::Validity; + use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_session::registry::ReadContext; + use super::*; use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::encode; + 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() { @@ -371,4 +597,404 @@ 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(()) + } + + #[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(()) + } + + #[rstest] + #[case::one_lower_part(i128_parts(wide_i128_values(), Validity::NonNullable))] + #[case::three_lower_parts(i256_parts(wide_i256_values(), Validity::NonNullable))] + #[case::nullable_three_lower_parts(i256_parts(wide_i256_values(), Validity::AllValid))] + fn test_serde_round_trip_with_lower_parts( + #[case] array: DecimalBytePartsArray, + ) -> VortexResult<()> { + test_serde_round_trip(array) + } + + #[rstest] + #[case::no_lower_parts( + encode(&DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(9, 2), Validity::NonNullable)) + .vortex_expect("valid decimal byte parts") + )] + fn test_serde_round_trip_flat(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { + test_serde_round_trip(array) + } + + fn test_serde_round_trip(array: DecimalBytePartsArray) -> VortexResult<()> { + let session = array_session(); + // Both serialized formats must be registered: an array with lower parts comes back + // under the v2 format id. + crate::initialize(&session); + + let array = array.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + let lower_part_count = array + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len(); + + let array_ctx = ArrayContext::empty(); + let serialized = array.serialize(&array_ctx, &session, &SerializeOptions::default())?; + let mut concat = ByteBufferMut::empty(); + for buf in serialized { + concat.extend_from_slice(buf.as_ref()); + } + let parts = SerializedArray::try_from(concat.freeze())?; + let decoded = parts.decode(&dtype, len, &ReadContext::new(array_ctx.to_ids()), &session)?; + + assert_eq!( + decoded + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len(), + lower_part_count, + "lower parts must survive serde" + ); + + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(array, decoded, &mut ctx); + 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) + )] + // Parts assembling into an i256 under a precision that only needs i128 would canonicalize + // to values outside the declared precision. + #[case::wider_than_precision(vec![lower_part(), lower_part()], DecimalDType::new(38, 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() + ); + } + + fn deserialize_with( + lower_part_count: u32, + children: Vec, + ) -> VortexResult> { + let metadata = DecimalBytesPartsMetadata { + zeroth_child_ptype: PType::I64 as i32, + lower_part_count, + }; + VTable::deserialize( + &DecimalByteParts, + &DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable), + 3, + &metadata.encode_to_vec(), + &[], + &children, + &array_session(), + ) + } + + #[test] + fn test_deserialize_reads_lower_parts() -> VortexResult<()> { + let parts = deserialize_with(1, vec![msp(), lower_part()])?; + let array = Array::try_from_parts(parts)?; + assert_eq!(array.lower_parts().len(), 1); + + let mut ctx = array_session().create_execution_ctx(); + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!( + canonical.buffer::().as_slice(), + &[(1i128 << 64) | 1, (2i128 << 64) | 2, (3i128 << 64) | 3] + ); + Ok(()) + } + + /// An array read from a file can be handed straight back to a writer, bypassing both the + /// constructor and the compressor. Its serialized id must still be the v2 format, so a + /// writer whose permitted encodings predate the v2 format refuses it. + #[test] + fn read_lower_parts_serialize_under_the_wide_format() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + + let array = deserialize_with(1, vec![msp(), lower_part()]) + .and_then(Array::try_from_parts)? + .into_array(); + + assert_eq!( + session.array_serialized_id(&array)?, + ArrayPlugin::id(&DecimalBytePartsV2) + ); + + let restricted = ArrayContext::empty() + .with_allowed_ids([VTable::id(&DecimalByteParts)].into_iter().collect()); + let err = array + .serialize(&restricted, &session, &SerializeOptions::default()) + .expect_err("expected the permitted-encoding check to refuse the v2 format"); + assert!( + err.to_string().contains("not permitted"), + "error should name the permitted-encoding check, got: {err}" + ); + Ok(()) + } + + /// Reading back an array that already carries lower parts, and computing over it, must + /// always work: the v2 format only restricts which writers may emit it. If reading or + /// the rebuild that every compute kernel does were blocked, a session whose editions + /// predate the v2 format could not read a file written by one that includes it. + #[test] + fn compute_over_existing_lower_parts_is_not_gated() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + // Stands in for an array materialized from a file: the parts already exist. + let array = deserialize_with(1, vec![msp(), lower_part()]) + .and_then(Array::try_from_parts)? + .into_array(); + + let sliced = array.slice(0..2)?; + assert_eq!(sliced.execute::(&mut ctx)?.len(), 2); + Ok(()) + } + + /// A crafted file may declare more lower parts than its precision needs. Assembling those + /// parts would produce values outside the declared precision, so it must be rejected at + /// deserialization rather than panicking later in the scalar path. + #[test] + fn test_deserialize_rejects_parts_wider_than_precision() { + let result = deserialize_with(2, vec![msp(), lower_part(), lower_part()]) + .and_then(Array::try_from_parts); + assert!(result.is_err(), "expected rejection, got {result:?}"); + } + + #[test] + fn test_deserialize_rejects_child_count_mismatch() { + // Metadata claiming a lower part that was not serialized. + assert!(deserialize_with(1, vec![msp()]).is_err()); + // Metadata claiming fewer lower parts than there are children. + assert!(deserialize_with(0, vec![msp(), lower_part()]).is_err()); + // Metadata claiming more lower parts than the encoding supports. + assert!( + deserialize_with( + 4, + vec![ + msp(), + lower_part(), + lower_part(), + lower_part(), + lower_part() + ] + ) + .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(()) + } } 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..bdee6df36a8 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Test-only helpers for building byte-parts arrays. + +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)?; + 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/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 36a53c3a614..8b4bf8e8b10 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -23,6 +23,9 @@ use vortex_session::VortexSession; /// Initialize decimal-byte-parts encoding in the given session. pub fn initialize(session: &VortexSession) { session.arrays().register(DecimalByteParts); + // The v2 serialized format deserializes into the same in-memory encoding. Registration + // covers reading; writing it is gated by the editions enabled on the writer's session. + session.arrays().register(DecimalBytePartsV2); compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( diff --git a/encodings/decimal-byte-parts/tests/format_v2.rs b/encodings/decimal-byte-parts/tests/format_v2.rs new file mode 100644 index 00000000000..9b83fe6e337 --- /dev/null +++ b/encodings/decimal-byte-parts/tests/format_v2.rs @@ -0,0 +1,161 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The v2 serialized format. +//! +//! Lower parts can be built and computed over freely. What changes with them is the bytes: +//! an array carrying lower parts serializes under `vortex.decimal_byte_parts_v2` rather +//! than the frozen `vortex.decimal_byte_parts` format, so a writer restricted to editions +//! without the v2 format refuses it, and a reader that predates lower parts fails with an +//! unknown-encoding error instead of misreading the children. These tests pin all of that: +//! construction always works, the serialized id tracks the parts, and the permitted-encoding +//! check applies to the serialized id. + +#![expect(clippy::tests_outside_test_module)] + +use vortex_array::ArrayContext; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArrayVTable; +use vortex_array::IntoArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::serde::SerializeOptions; +use vortex_array::session::ArraySessionExt; +use vortex_buffer::buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsV2; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() +} + +fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() +} + +fn session() -> VortexSession { + let session = vortex_array::array_session(); + vortex_decimal_byte_parts::initialize(&session); + session +} + +/// A single-child array is the stable shape and is always constructible. +#[test] +fn single_child_is_always_allowed() { + assert!(DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2)).is_ok()); + assert!( + DecimalByteParts::try_new_with_lower_parts(msp(), vec![], DecimalDType::new(19, 2)).is_ok() + ); +} + +/// Building lower parts in memory is always allowed — reading a file requires it. What +/// changes is the serialized format, not what can be constructed. +#[test] +fn lower_parts_can_always_be_constructed() { + assert!( + DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + ) + .is_ok() + ); +} + +/// A single-child array keeps the frozen format id, byte-compatible with every reader since +/// the format froze; lower parts move the array onto the v2 format id. +#[test] +fn serialized_id_tracks_lower_parts() -> VortexResult<()> { + let session = session(); + + let flat = DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2))?.into_array(); + assert_eq!( + session.array_serialized_id(&flat)?, + ArrayVTable::id(&DecimalByteParts) + ); + + let wide = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + assert_eq!( + session.array_serialized_id(&wide)?, + ArrayPlugin::id(&DecimalBytePartsV2) + ); + + Ok(()) +} + +/// The permitted-encoding check applies to the serialized id. A context restricted to the +/// frozen format — a writer whose enabled editions predate the v2 format — must refuse an +/// array carrying lower parts, however it was obtained. +/// +/// `ArrayParts` is public and `DecimalBytePartsData` is a public unit struct, so a caller can +/// assemble slots by hand and go straight to `Array::try_from_parts`, bypassing +/// `try_new_with_lower_parts` entirely. That back door is left open on purpose — it is the +/// same path `deserialize` uses. What must hold is that the resulting array cannot become +/// bytes under the frozen id. +#[test] +fn wide_format_is_refused_where_not_permitted() -> VortexResult<()> { + use vortex_array::Array; + use vortex_array::ArrayParts; + use vortex_array::ArraySlots; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_decimal_byte_parts::DecimalBytePartsData; + + let session = session(); + + let mut slots = ArraySlots::with_capacity(2); + slots.push(Some(msp())); + slots.push(Some(lower_part())); + + // Assembling the array by hand succeeds: this is the shape a file read produces. + let array = Array::try_from_parts( + ArrayParts::new( + DecimalByteParts, + DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable), + 3, + DecimalBytePartsData, + ) + .with_slots(slots), + )? + .into_array(); + assert_eq!(array.nchildren(), 2, "expected two limbs"); + + // A context permitting only the frozen format refuses to write it. + let restricted = ArrayContext::empty() + .with_allowed_ids([ArrayVTable::id(&DecimalByteParts)].into_iter().collect()); + let err = array + .serialize(&restricted, &session, &SerializeOptions::default()) + .expect_err("expected the permitted-encoding check to refuse the v2 format"); + assert!( + err.to_string().contains("not permitted"), + "error should name the permitted-encoding check, got: {err}" + ); + + // Permitting the v2 format id is exactly what allows the same array through. + let permissive = ArrayContext::empty().with_allowed_ids( + [ + ArrayVTable::id(&DecimalByteParts), + ArrayPlugin::id(&DecimalBytePartsV2), + ArrayVTable::id(&vortex_array::arrays::Primitive), + ] + .into_iter() + .collect(), + ); + let serialized = array.serialize(&permissive, &session, &SerializeOptions::default())?; + assert!(!serialized.is_empty()); + assert!( + permissive + .to_ids() + .contains(&ArrayPlugin::id(&DecimalBytePartsV2)), + "the file's encoding table must carry the v2 format id" + ); + + Ok(()) +} diff --git a/encodings/decimal-byte-parts/tests/props.rs b/encodings/decimal-byte-parts/tests/props.rs new file mode 100644 index 00000000000..583dc17e103 --- /dev/null +++ b/encodings/decimal-byte-parts/tests/props.rs @@ -0,0 +1,199 @@ +// 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) -> DecimalBytePartsArray { + let parts = split_decimal(decimal).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).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).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 fill in `sign_extended_words` is caught by neither +// property here. diff --git a/vortex-array/src/array/plugin.rs b/vortex-array/src/array/plugin.rs index 66845eb9a0a..1ff8d2816dc 100644 --- a/vortex-array/src/array/plugin.rs +++ b/vortex-array/src/array/plugin.rs @@ -42,6 +42,16 @@ pub trait ArrayPlugin: 'static + Send + Sync { fn serialize(&self, array: &ArrayRef, session: &VortexSession) -> VortexResult>>; + /// The serialized format id written to a file or stream for this array. + /// + /// Like [`serialize`](Self::serialize), this is only called for arrays whose encoding ID + /// matches this plugin. Defaults to the plugin [`id`](Self::id); a plugin owning more + /// than one serialized format overrides this to pick the format able to represent the + /// given array. + fn serialized_id(&self, _array: &ArrayRef) -> ArrayId { + self.id() + } + /// Deserialize an array from serialized components. /// /// The returned array doesn't necessary have to match this plugin's encoding ID. This is @@ -89,6 +99,10 @@ impl ArrayPlugin for V { V::serialize(array.as_::(), session) } + fn serialized_id(&self, array: &ArrayRef) -> ArrayId { + VTable::serialized_id(self, array.as_::()) + } + fn deserialize( &self, dtype: &DType, diff --git a/vortex-array/src/array/vtable/mod.rs b/vortex-array/src/array/vtable/mod.rs index 3078a6ae9f5..e4ff4b28e07 100644 --- a/vortex-array/src/array/vtable/mod.rs +++ b/vortex-array/src/array/vtable/mod.rs @@ -125,6 +125,16 @@ pub trait VTable: 'static + Clone + Sized + Send + Sync + Debug { session: &VortexSession, ) -> VortexResult>>; + /// The serialized format id written for this array. + /// + /// Defaults to the encoding [`id`](Self::id). An encoding with more than one serialized + /// format overrides this to pick the format able to represent the given array. The + /// writer's permitted-encoding check applies to the returned id, so a newer format can be + /// gated by editions independently of the in-memory encoding that produces it. + fn serialized_id(&self, _array: ArrayView<'_, Self>) -> ArrayId { + self.id() + } + /// Deserialize an array from serialized metadata, buffers, and children. /// /// The returned [`ArrayParts`] are still validated by the generic adapter. diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index f84ec182269..a50d04ea9cc 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -191,12 +191,14 @@ impl<'a> ArrayNodeFlatBuffer<'a> { &self, fbb: &mut FlatBufferBuilder<'fb>, ) -> VortexResult>> { - let encoding_idx = self.ctx.intern(&self.array.encoding_id()).ok_or_else(|| { - vortex_err!( - "Array encoding {} not permitted by ctx", - self.array.encoding_id() - ) - })?; + // The id written to the file is the plugin's serialized format id, which may differ + // from the in-memory encoding id. The permitted-encoding check applies to it: what is + // gated is the bytes a reader will meet, not the in-memory representation. + let serialized_id = self.session.array_serialized_id(self.array)?; + let encoding_idx = self + .ctx + .intern(&serialized_id) + .ok_or_else(|| vortex_err!("Array encoding {} not permitted by ctx", serialized_id))?; let metadata_bytes = self.session.array_serialize(self.array)?.ok_or_else(|| { vortex_err!( diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 2f3fbb9e4e7..bc481fb79b1 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -123,6 +123,22 @@ pub trait ArraySessionExt: SessionExt { plugin.serialize(array, &self.session()) } + + /// Resolve the serialized format id an array will be written under. + /// + /// Most arrays serialize under their own encoding id, but a plugin may pick a different + /// format for arrays its original format cannot represent — see + /// [`ArrayPlugin::serialized_id`]. + fn array_serialized_id(&self, array: &ArrayRef) -> VortexResult { + let Some(plugin) = self.arrays().registry.get(&array.encoding_id()) else { + vortex_bail!( + "Array {} is not registered for serializations", + array.encoding_id() + ); + }; + + Ok(plugin.serialized_id(array)) + } } impl ArraySessionExt for S {} diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal/mod.rs similarity index 63% rename from vortex-btrblocks/src/schemes/decimal.rs rename to vortex-btrblocks/src/schemes/decimal/mod.rs index 1dff2171f60..21319510d82 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal/mod.rs @@ -10,12 +10,12 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; -use vortex_array::dtype::DecimalType; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsSlots; +use vortex_decimal_byte_parts::split_decimal; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -28,6 +28,10 @@ use crate::SchemeExt; /// /// Narrows the decimal to the smallest integer type, compresses the underlying primitive, and wraps /// the result in a `DecimalBytePartsArray`. +/// +/// Only decimals that fit a single signed part are compressed. Anything still wider than 64 +/// bits after narrowing would need lower parts, which cannot be serialized, so those are left +/// as the canonical decimal. #[derive(Debug, Copy, Clone, PartialEq, Eq)] pub struct DecimalScheme; @@ -44,9 +48,9 @@ impl Scheme for DecimalScheme { vec![DecimalByteParts.id()] } - /// Children: primitive=0. + /// Children: msp=0. This scheme never emits lower parts. fn num_children(&self) -> usize { - 1 + DecimalBytePartsSlots::FIXED_COUNT } fn expected_compression_ratio( @@ -66,22 +70,28 @@ impl Scheme for DecimalScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): add support splitting i128/256 buffers into chunks of primitive values - // for compression. 2 for i128 and 4 for i256. let decimal = data.array().clone().execute::(exec_ctx)?; let decimal = narrowed_decimal(decimal); - let validity = decimal.validity()?; - let prim = match decimal.values_type() { - DecimalType::I8 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I16 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I32 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I64 => PrimitiveArray::new(decimal.buffer::(), validity), - _ => return Ok(decimal.into_array()), - }; + let parts = split_decimal(&decimal)?; - let compressed = - compressor.compress_child(&prim.into_array(), &compress_ctx, self.id(), 0, exec_ctx)?; + // A value too wide for one signed part splits into lower parts, which serialize under + // the v2 format id — one this scheme does not declare in `produced_encodings`, so a + // writer restricted to its editions could refuse it. Leave it as the canonical decimal + // rather than build something outside the scheme's declared output. + if !parts.lower_parts.is_empty() { + return Ok(decimal.into_array()); + } - DecimalByteParts::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array()) + let msp = compressor.compress_child( + &parts.msp, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::MSP, + exec_ctx, + )?; + DecimalByteParts::try_new(msp, decimal.decimal_dtype()).map(|d| d.into_array()) } } + +#[cfg(test)] +mod tests; diff --git a/vortex-btrblocks/src/schemes/decimal/tests.rs b/vortex-btrblocks/src/schemes/decimal/tests.rs new file mode 100644 index 00000000000..0d40607db67 --- /dev/null +++ b/vortex-btrblocks/src/schemes/decimal/tests.rs @@ -0,0 +1,158 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::iter; +use std::sync::LazyLock; + +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::BtrBlocksCompressor; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +/// Number of values per array: above the 1024-value sampling threshold, so scheme selection +/// runs on sampled estimates as it does for real file chunks. +const N: usize = 16_384; + +fn ten_pow(exp: u32) -> i256 { + i256::from_i128(10).wrapping_pow(exp) +} + +/// Deterministic 24-bit noise, so the low part of each value is neither constant nor a +/// sequence — the realistic shape for a wide decimal column with a large fixed magnitude. +fn noise(seed: u64) -> impl Iterator { + let mut rng = StdRng::seed_from_u64(seed); + iter::repeat_with(move || i128::from(rng.random::() >> 8)) +} + +/// `i128`-backed values that need more than 64 bits, so the encoding must carry one lower +/// part. +fn wide_i128_array(validity: Validity) -> DecimalArray { + let base = 10i128.pow(25); + let values: Buffer = noise(7).take(N).map(|delta| base + delta).collect(); + DecimalArray::new(values, DecimalDType::new(38, 2), validity) +} + +/// `i256`-backed values that need more than 128 bits, so the encoding must carry three +/// lower parts. +fn wide_i256_array(validity: Validity) -> DecimalArray { + let base = ten_pow(40); + let values: Buffer = noise(11) + .take(N) + .map(|delta| base + i256::from_i128(delta)) + .collect(); + DecimalArray::new(values, DecimalDType::new(76, 2), validity) +} + +fn compress(array: &ArrayRef) -> VortexResult { + BtrBlocksCompressor::default().compress(array, &mut SESSION.create_execution_ctx()) +} + +fn byte_parts(array: &ArrayRef) -> &ArrayRef { + assert!( + array.is::(), + "expected DecimalByteParts, got {}", + array.encoding_id() + ); + array +} + +fn lower_part_count(array: &ArrayRef) -> usize { + byte_parts(array) + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len() +} + +/// Values too wide for a single signed part are left as the canonical decimal. Splitting them +/// would need lower parts, which serialize under the v2 format id the scheme does not +/// declare, so it declines rather than building an array outside its declared output. +#[rstest] +#[case::i128(wide_i128_array(Validity::NonNullable).into_array())] +#[case::i128_nullable(wide_i128_array(Validity::from_iter((0..N).map(|i| i % 3 != 0))).into_array())] +#[case::i256(wide_i256_array(Validity::NonNullable).into_array())] +#[case::i256_nullable(wide_i256_array(Validity::from_iter((0..N).map(|i| i % 5 != 0))).into_array())] +fn test_wide_decimals_are_left_canonical(#[case] array: ArrayRef) -> VortexResult<()> { + let compressed = compress(&array)?; + + assert!( + compressed.as_opt::().is_none(), + "expected the wide decimal to be left canonical, got {}", + compressed.encoding_id() + ); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[test] +fn test_i256_decimal_round_trips_extreme_values() -> VortexResult<()> { + // Every 64-bit window exercised, including the sign boundary of the most significant + // part. Bounded by the precision so the values are legal `Decimal(76, 0)` scalars. + let max = ten_pow(76) - i256::ONE; + let values: Buffer = (0..N) + .map(|i| match i % 8 { + 0 => i256::ZERO, + 1 => i256::ONE, + 2 => i256::ZERO - i256::ONE, + 3 => i256::from_parts(u128::MAX, 0), + 4 => i256::from_parts(0, 1), + 5 => i256::from_parts(0, -1), + 6 => max, + _ => i256::ZERO - max, + }) + .collect(); + let array = + DecimalArray::new(values, DecimalDType::new(76, 0), Validity::NonNullable).into_array(); + + let compressed = compress(&array)?; + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[test] +fn test_narrow_decimal_has_no_lower_parts() -> VortexResult<()> { + // Values that fit 64 bits are narrowed rather than split, even when the declared + // precision needs an i256. + let values: Buffer = (0..N as i128).map(|i| i256::from_i128(i * 3)).collect(); + let array = + DecimalArray::new(values, DecimalDType::new(76, 2), Validity::NonNullable).into_array(); + + let compressed = compress(&array)?; + assert_eq!(lower_part_count(&compressed), 0); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[test] +fn test_canonical_of_compressed_wide_decimal_keeps_storage_width() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + let array = wide_i128_array(Validity::NonNullable).into_array(); + let canonical = compress(&array)?.execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + + let array = wide_i256_array(Validity::NonNullable).into_array(); + let canonical = compress(&array)?.execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + Ok(()) +} diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index ee91ffab298..03471af54d5 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -396,7 +396,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 @@ -432,6 +432,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, diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 5992ff06434..8f2b6aaf379 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -61,6 +61,7 @@ vortex-zigzag = { workspace = true } vortex-zstd = { workspace = true, optional = true } [dev-dependencies] +rand = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..fadb42eb296 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -11,6 +11,9 @@ use flatbuffers::FlatBufferBuilder; use futures::StreamExt; use futures::TryStreamExt; use futures::pin_mut; +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -38,6 +41,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; +use vortex_array::dtype::i256; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::and; @@ -250,6 +254,75 @@ async fn test_round_trip_many_types() { assert_eq!(read.len(), 3); } +/// End-to-end check that decimals wider than 64 bits survive a write/read round trip. +/// +/// The compressor declines to split these — that would need lower parts, which cannot be +/// serialized — so they are written as canonical decimals. This pins that the wide path still +/// round trips through a file rather than being compressed. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_wide_decimal_round_trips_through_a_file() -> VortexResult<()> { + const N: usize = 16_384; + + /// Deterministic 24-bit noise, so the low bits of each value are neither constant nor a + /// sequence. + fn noise(seed: u64) -> impl Iterator { + let mut rng = StdRng::seed_from_u64(seed); + iter::repeat_with(move || i128::from(rng.random::() >> 8)) + } + + // Values that need more than 64 bits, so `i128` storage cannot be narrowed away. + let decimal_38 = DecimalArray::new( + noise(7) + .take(N) + .map(|delta| 10i128.pow(25) + delta) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array(); + + // Values that need more than 128 bits, so `i256` storage cannot be narrowed away. + let base = i256::from_i128(10).wrapping_pow(40); + let decimal_76 = DecimalArray::new( + noise(11) + .take(N) + .map(|delta| base + i256::from_i128(delta)) + .collect::>(), + DecimalDType::new(76, 4), + Validity::from_iter((0..N).map(|i| i % 9 != 0)), + ) + .into_array(); + + let st = StructArray::from_fields(&[ + ("decimal_38", decimal_38), + ("decimal_76_nullable", decimal_76), + ])? + .into_array(); + let dtype = st.dtype().clone(); + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, st.clone().to_array_stream()) + .await?; + + let chunks: Vec<_> = SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .try_collect() + .await?; + let read = ChunkedArray::try_new(chunks, dtype)?.into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(read.len(), N); + assert_arrays_eq!(st, read, &mut ctx); + + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_read_simple_with_spawn() { @@ -2800,3 +2873,67 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { assert_eq!(result.len(), 1); Ok(()) } + +/// A multi-limb array handed to the default file writer does not reach the file as one. +/// +/// Two independent mechanisms stand between lower parts and a file, and this pins the one +/// that applies here: the writer recompresses its input, and the decimal scheme declines to +/// split values too wide for a single signed part, so the column lands as a canonical decimal +/// with no children. The other mechanism — the permitted-encoding check refusing the +/// `vortex.decimal_byte_parts_v2` serialized format outside its edition — is the backstop +/// for a write strategy that does not recompress, and is covered in the encoding crate and +/// the `vortex` editions tests. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_lower_parts_recompress_to_canonical_through_the_default_writer() -> VortexResult<()> { + use vortex_decimal_byte_parts::DecimalByteParts; + use vortex_decimal_byte_parts::split_decimal; + + let decimal = DecimalArray::new( + (0..64i128) + .map(|i| (1i128 << 70) + i) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + + // Building the encoded array is allowed; only getting it into a file is restricted. + let parts = split_decimal(&decimal)?; + assert_eq!(parts.lower_parts.len(), 1, "expected a wide split"); + let encoded = DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + )? + .into_array(); + + let st = StructArray::from_fields(&[("wide", encoded)])?.into_array(); + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, st.clone().to_array_stream()) + .await?; + + let chunks: Vec<_> = SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .try_collect() + .await?; + + for chunk in &chunks { + for field in chunk.children_iter() { + assert!( + field.as_opt::().is_none(), + "expected a canonical decimal in the file, got {}", + field.encoding_id() + ); + } + } + + let mut ctx = SESSION.create_execution_ctx(); + let read = ChunkedArray::try_new(chunks, st.dtype().clone())?.into_array(); + assert_arrays_eq!(st, read, &mut ctx); + Ok(()) +} diff --git a/vortex-test/compat-gen/Cargo.toml b/vortex-test/compat-gen/Cargo.toml index 4a62aca3671..2a5fe657d9b 100644 --- a/vortex-test/compat-gen/Cargo.toml +++ b/vortex-test/compat-gen/Cargo.toml @@ -20,6 +20,11 @@ name = "vortex-compat" path = "src/main.rs" test = false +[features] +# Fixtures for encodings whose on-disk shape is not yet stable. Kept out of the default +# fixture set so a default build never publishes a file older readers cannot open. +unstable_encodings = ["vortex/unstable_encodings"] + [dependencies] # Vortex crates vortex = { workspace = true, features = ["files", "tokio", "zstd"] } diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs new file mode 100644 index 00000000000..8a7207785d0 --- /dev/null +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Wide `DecimalByteParts` fixtures: values that need lower parts. +//! +//! These live in their own fixture file rather than as extra columns on +//! `decimal_byte_parts.vortex` because a fixture's `build()` is immutable once published. +//! `check` compares files written by older releases against what `build()` produces today, +//! so changing an existing fixture's schema fails the check against every previously +//! published version — see "Fixture evolution" in `DESIGN.md`, which requires a new fixture +//! file with a new name for a new type, encoding, or structural pattern. +//! +//! So `decimal_byte_parts.vortex` keeps testing exactly what it always did, decimals whose +//! values fit a single signed part, and the MSP-plus-lower-parts layout added alongside it +//! is covered here instead. + +use vortex::array::ArrayId; +use vortex::array::ArrayRef; +use vortex::array::ArrayVTable; +use vortex::array::IntoArray; +use vortex::array::arrays::DecimalArray; +use vortex::array::arrays::StructArray; +use vortex::array::dtype::DecimalDType; +use vortex::array::dtype::FieldNames; +use vortex::array::dtype::i256; +use vortex::array::validity::Validity; +use vortex::buffer::Buffer; +use vortex::encodings::decimal_byte_parts::DecimalByteParts; +use vortex::encodings::decimal_byte_parts::DecimalBytePartsArray; +use vortex::encodings::decimal_byte_parts::split_decimal; +use vortex::error::VortexResult; +use vortex_array::ExecutionCtx; + +use super::N; +use crate::fixtures::FlatLayoutFixture; + +/// Encode a canonical decimal as byte parts, splitting wide values into lower parts. +fn encode_byte_parts(decimal: &DecimalArray) -> VortexResult { + let parts = split_decimal(decimal)?; + DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ) +} + +pub struct DecimalBytePartsV2Fixture; + +impl FlatLayoutFixture for DecimalBytePartsV2Fixture { + fn name(&self) -> &str { + "decimal_byte_parts_v2.vortex" + } + + fn description(&self) -> &str { + "Wide decimal arrays split into a most significant part plus 64-bit lower parts" + } + + fn expected_encodings(&self) -> Vec { + vec![DecimalByteParts.id()] + } + + fn build(&self, _ctx: &mut ExecutionCtx) -> VortexResult { + // An `i128` magnitude above 2^64, so the encoding must carry one lower part. + let wide_128_dtype = DecimalDType::new(38, 2); + let wide_128 = DecimalArray::new( + (0..N as i128) + .map(|i| 10i128.pow(25) + i * 7) + .collect::>(), + wide_128_dtype, + Validity::NonNullable, + ); + let wide_128_arr = encode_byte_parts(&wide_128)?; + + // Negative values, so the sign extension above the MSP is exercised on read back. + let wide_128_negative = DecimalArray::new( + (0..N as i128) + .map(|i| -(10i128.pow(25)) - i * 7) + .collect::>(), + wide_128_dtype, + Validity::NonNullable, + ); + let wide_128_negative_arr = encode_byte_parts(&wide_128_negative)?; + + // An `i256` magnitude beyond 128 bits, so all three lower parts are populated, with + // nulls to pin that validity is carried by the MSP alone. + let wide_256_dtype = DecimalDType::new(76, 2); + let base = i256::from_i128(10).wrapping_pow(40); + let wide_256 = DecimalArray::new( + (0..N as i128) + .map(|i| base + i256::from_i128(i * 7)) + .collect::>(), + wide_256_dtype, + Validity::from_iter((0..N).map(|i| i % 7 != 0)), + ); + let wide_256_arr = encode_byte_parts(&wide_256)?; + + let arr = StructArray::try_new( + FieldNames::from([ + "dec_wide_128", + "dec_wide_128_negative", + "dec_wide_256_nullable", + ]), + vec![ + wide_128_arr.into_array(), + wide_128_negative_arr.into_array(), + wide_256_arr.into_array(), + ], + N, + Validity::NonNullable, + )?; + Ok(arr.into_array()) + } +} diff --git a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs index 830b50450da..f49d03a1fae 100644 --- a/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/mod.rs @@ -12,6 +12,8 @@ mod bytebool; mod constant; mod datetimeparts; mod decimal_byte_parts; +#[cfg(feature = "unstable_encodings")] +mod decimal_byte_parts_v2; mod delta; mod dict; mod for_; @@ -31,7 +33,8 @@ pub(crate) const N: usize = 1024; /// All per-encoding fixtures. pub fn fixtures() -> Vec> { - vec![ + #[allow(unused_mut)] + let mut fixtures: Vec> = vec![ Box::new(alp::AlpFixture), Box::new(alprd::AlprdFixture), Box::new(bitpacked::BitPackedFixture), @@ -53,5 +56,10 @@ pub fn fixtures() -> Vec> { Box::new(zstd::ZstdFixture), Box::new(zigzag::ZigZagFixture), Box::new(constant::ConstantFixture), - ] + ]; + #[cfg(feature = "unstable_encodings")] + fixtures.push(Box::new( + decimal_byte_parts_v2::DecimalBytePartsV2Fixture, + )); + fixtures } diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 54aa3f1e515..b4475aae703 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -68,6 +68,7 @@ tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } vortex = { path = ".", features = ["tokio"] } +vortex-array = { workspace = true, features = ["_test-harness"] } [features] default = ["files", "zstd"] diff --git a/vortex/src/editions/mod.rs b/vortex/src/editions/mod.rs index 88e4f351d0d..bb67b74bcd1 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -38,13 +38,14 @@ pub use self::unstable::UNSTABLE_2025_05_0; pub use self::unstable::UNSTABLE_2026_02_0; pub use self::unstable::UNSTABLE_2026_04_0; pub use self::unstable::UNSTABLE_2026_06_0; +pub use self::unstable::UNSTABLE_2026_08_0; /// The `core` edition enabled for writing by the default Vortex session. pub const DEFAULT_CORE_EDITION: EditionId = CORE_2026_08; /// The `unstable` edition enabled for writing by the default Vortex session when the /// `unstable_encodings` feature is selected. -pub const DEFAULT_UNSTABLE_EDITION: EditionId = UNSTABLE_2026_06_0; +pub const DEFAULT_UNSTABLE_EDITION: EditionId = UNSTABLE_2026_08_0; /// The first-party Vortex edition declarations. pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ @@ -57,6 +58,7 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &unstable::v2026_02::DECLARATION, &unstable::v2026_04::DECLARATION, &unstable::v2026_06::DECLARATION, + &unstable::v2026_08::DECLARATION, ]; /// Register the Vortex edition declarations with the session's [`EditionSession`]. diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index b8608f4427a..2fd99730a17 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -44,6 +44,7 @@ use super::DEFAULT_CORE_EDITION; use super::DEFAULT_UNSTABLE_EDITION; use super::EDITION_DECLARATIONS; use super::UNSTABLE_2026_06_0; +use super::UNSTABLE_2026_08_0; fn session() -> Result { let session = EditionSession::empty(); @@ -135,6 +136,21 @@ fn encodings_in_editions_unions_families() { assert!(core_only.iter().all(|id| both.contains(id))); } +/// The wide byte-parts serialized format joins the `unstable` family at 2026.08: absent from +/// the June draft, present from August on. +#[test] +fn decimal_byte_parts_v2_joins_unstable_2026_08() { + let session = session().unwrap_or_else(|e| panic!("registering editions: {e}")); + let in_edition = |edition| { + session + .encodings_in(edition) + .iter() + .any(|inclusion| inclusion.encoding_id.as_str() == "vortex.decimal_byte_parts_v2") + }; + assert!(!in_edition(&UNSTABLE_2026_06_0)); + assert!(in_edition(&UNSTABLE_2026_08_0)); +} + #[test] fn earlier_editions_are_subsets() { let session = session().unwrap_or_else(|e| panic!("registering editions: {e}")); @@ -448,6 +464,111 @@ async fn configured_btrblocks_builder_uses_enabled_editions_in_either_order() -> Ok(()) } +/// Build a byte-parts array that must serialize under `vortex.decimal_byte_parts_v2`. +fn wide_byte_parts_column() -> VortexResult { + use vortex_array::arrays::DecimalArray; + use vortex_array::dtype::DecimalDType; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + + use crate::encodings::decimal_byte_parts::DecimalByteParts; + use crate::encodings::decimal_byte_parts::split_decimal; + + let decimal = DecimalArray::new( + (0..64i128) + .map(|i| (1i128 << 70) + i) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal)?; + let encoded = DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + )? + .into_array(); + Ok(StructArray::from_fields(&[("wide", encoded)])?.into_array()) +} + +/// The wide byte-parts serialized format is gated by the editions enabled for writing. +/// +/// The in-memory encoding is `vortex.decimal_byte_parts` either way; what the editions gate +/// is the serialized format id. A flat (non-recompressing) strategy hands the multi-limb +/// array straight to serialization, so the permitted-encoding check is the only thing +/// between it and the file: without the `unstable` edition the write is refused, and with it +/// the array round-trips back into the same in-memory encoding, lower parts intact. +#[tokio::test] +async fn decimal_byte_parts_v2_is_gated_by_the_unstable_edition() -> VortexResult<()> { + use crate::VortexSessionDefault; + + let st = wide_byte_parts_column()?; + let session = VortexSession::default(); + let mut buffer = ByteBufferMut::empty(); + let result = session + .write_options() + .with_strategy(Arc::new(FlatLayoutStrategy::default())) + .write(&mut buffer, st.clone().to_array_stream()) + .await; + + #[cfg(not(feature = "unstable_encodings"))] + { + let _ = &st; + let error = match result { + Ok(_) => { + return Err(vortex_err!( + "the v2 serialized format requires the unstable edition" + )); + } + Err(error) => error, + }; + let message = error.to_string(); + assert!( + message.contains("not permitted by ctx"), + "unexpected error: {message}" + ); + } + + #[cfg(feature = "unstable_encodings")] + { + use vortex_array::VortexSessionExecute; + use vortex_array::assert_arrays_eq; + + use crate::encodings::decimal_byte_parts::DecimalByteParts; + use crate::encodings::decimal_byte_parts::DecimalBytePartsArraySlotsExt; + + result?; + let read = session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + + let lower_part_counts: Vec = read + .depth_first_traversal() + .filter_map(|node| { + node.as_opt::() + .map(|array| array.lower_parts().len()) + }) + .collect(); + assert!( + !lower_part_counts.is_empty(), + "expected the v2 format to deserialize into vortex.decimal_byte_parts" + ); + assert!( + lower_part_counts.iter().all(|count| *count > 0), + "lower parts must survive the round trip" + ); + + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(st, read, &mut ctx); + } + + Ok(()) +} + #[tokio::test] async fn opaque_compressor_cannot_write_outside_enabled_editions() -> VortexResult<()> { let session = baseline_core_session()?; diff --git a/vortex/src/editions/unstable/mod.rs b/vortex/src/editions/unstable/mod.rs index 5544ba45c0f..7c271198e25 100644 --- a/vortex/src/editions/unstable/mod.rs +++ b/vortex/src/editions/unstable/mod.rs @@ -10,8 +10,10 @@ pub mod v2025_05; pub mod v2026_02; pub mod v2026_04; pub mod v2026_06; +pub mod v2026_08; pub use v2025_05::UNSTABLE_2025_05_0; pub use v2026_02::UNSTABLE_2026_02_0; pub use v2026_04::UNSTABLE_2026_04_0; pub use v2026_06::UNSTABLE_2026_06_0; +pub use v2026_08::UNSTABLE_2026_08_0; diff --git a/vortex/src/editions/unstable/v2026_08.rs b/vortex/src/editions/unstable/v2026_08.rs new file mode 100644 index 00000000000..73483357529 --- /dev/null +++ b/vortex/src/editions/unstable/v2026_08.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The August 2026 `unstable` encoding cohort. + +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; + +/// The August 2026 draft edition of the `unstable` family. +pub const UNSTABLE_2026_08_0: EditionId = EditionId::new("unstable", 2026, 8, 0); + +/// The declaration of [`UNSTABLE_2026_08_0`] and the encodings that join the family at it. +/// +/// `vortex.decimal_byte_parts_v2` is a serialized format of the `DecimalByteParts` +/// encoding: byte-parts arrays carrying lower parts serialize under this id because the +/// frozen `vortex.decimal_byte_parts` format promises a single child. +pub static DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: UNSTABLE_2026_08_0, + min_vortex_version: None, + }, + added: &[&"vortex.decimal_byte_parts_v2"], +};