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 275db7ef8ee..ca9cd626f1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1120,7 +1120,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f21ff1fc630079352bae9b024f85519bf1f641cf7f326623f4c0b59f7ea834fd" dependencies = [ "compact_str", - "miniz_oxide", + "miniz_oxide 0.9.1", "thiserror 2.0.20", ] @@ -2229,6 +2229,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" @@ -4071,7 +4090,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", - "miniz_oxide", + "miniz_oxide 0.9.1", "zlib-rs", ] @@ -4678,6 +4697,51 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hegeltest" +version = "0.28.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "100bcd6ef825f5b6a60e2f55c05bb626ebf254dd8a09d16e006c4bb7883e7f1c" +dependencies = [ + "crc32fast", + "dashu-int", + "hegeltest-c", + "hegeltest-macros", + "miniz_oxide 0.8.9", + "parking_lot", + "paste", + "rand 0.10.2", + "rustc-hash", + "tempfile", +] + +[[package]] +name = "hegeltest-c" +version = "0.30.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a672fd53360ca4122c1a145a85e8fef835508d7b40eb9de43498978e796c54b" +dependencies = [ + "dashu-int", + "hashbrown 0.17.1", + "libm", + "miniz_oxide 0.8.9", + "parking_lot", + "rand 0.10.2", + "rustc-hash", + "tempfile", +] + +[[package]] +name = "hegeltest-macros" +version = "0.28.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba792d78fa3740a7c1627085c34618b998b8aa0f63625721235234f525aad1aa" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "hermit-abi" version = "0.5.3" @@ -6592,6 +6656,15 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", +] + [[package]] name = "miniz_oxide" version = "0.9.1" @@ -6898,6 +6971,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" @@ -10983,10 +11062,14 @@ 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-bench-support", "vortex-buffer", "vortex-error", "vortex-mask", diff --git a/Cargo.toml b/Cargo.toml index 6a94cd1eb56..4d73870beb4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -176,6 +176,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/docs/specs/editions.md b/docs/specs/editions.md index ee5f494861f..1aa3cfe4e11 100644 --- a/docs/specs/editions.md +++ b/docs/specs/editions.md @@ -172,9 +172,9 @@ representation gains support for wide decimals, represented by a signed most-sig represented that way, it emits `vortex.decimal_byte_parts` with `lower_part_count = 0`, even if the current in-memory array has lower-part children. - An array that cannot be collapsed into that old form losslessly uses the new - `vortex.decimal_byte_parts_v2` component, initially staged in a draft edition. + `vortex.decimal_byte_parts.v2` component, initially staged in a draft edition. - A new reader deserializes both IDs into the same in-memory representation. An older reader reports - `vortex.decimal_byte_parts_v2` as unknown instead of trying to decode a wire format it does not support. + `vortex.decimal_byte_parts.v2` as unknown instead of trying to decode a wire format it does not support. - When targeting an edition that permits only the old ID, serializing a value that can be collapsed succeeds; an irreducibly multi-part value fails because no lossless downgrade exists. diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 4934ec4fa27..e5d15c2c4f7 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -26,5 +26,17 @@ 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"] } +vortex-bench-support = { workspace = true } + +[[bench]] +name = "dbp_assemble" +harness = false + +[[bench]] +name = "dbp_split" +harness = false diff --git a/encodings/decimal-byte-parts/benches/common/mod.rs b/encodings/decimal-byte-parts/benches/common/mod.rs new file mode 100644 index 00000000000..252431a58a3 --- /dev/null +++ b/encodings/decimal-byte-parts/benches/common/mod.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared inputs for splitting and assembly benchmarks. + +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::i256; +use vortex_buffer::Buffer; + +pub(super) fn cases() -> Vec<(DecimalType, usize)> { + [DecimalType::I128, DecimalType::I256] + .into_iter() + .flat_map(|values_type| [1_024, 8_192].map(|len| (values_type, len))) + .collect() +} + +pub(super) fn i128_values(len: usize) -> Buffer { + let mut rng = StdRng::seed_from_u64(42); + let max = 10i128.pow(38) - 1; + (0..len).map(|_| rng.random_range(-max..=max)).collect() +} + +pub(super) fn i256_values(len: usize) -> Buffer { + let mut rng = StdRng::seed_from_u64(42); + // Keep the magnitude below 10^76 while exercising all four signed/unsigned words. + (0..len) + .map(|_| i256::from_parts(rng.random(), rng.random::() >> 4)) + .collect() +} + +#[cfg(not(codspeed))] +pub(super) mod arrays { + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; + use vortex_array::arrays::DecimalArray; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::DecimalType; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_error::vortex_panic; + + use super::i128_values; + use super::i256_values; + + pub(crate) fn cases() -> Vec<(DecimalType, usize)> { + [DecimalType::I64, DecimalType::I128, DecimalType::I256] + .into_iter() + .flat_map(|values_type| [1_024, 8_192].map(|len| (values_type, len))) + .collect() + } + + pub(crate) fn decimal_array( + values_type: DecimalType, + len: usize, + validity: Validity, + ) -> DecimalArray { + match values_type { + DecimalType::I64 => { + let mut rng = StdRng::seed_from_u64(42); + let max = 10i64.pow(18) - 1; + let values: Buffer = (0..len).map(|_| rng.random_range(-max..=max)).collect(); + DecimalArray::new(values, DecimalDType::new(18, 2), validity) + } + DecimalType::I128 => { + DecimalArray::new(i128_values(len), DecimalDType::new(38, 2), validity) + } + DecimalType::I256 => { + DecimalArray::new(i256_values(len), DecimalDType::new(76, 2), validity) + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } + } +} diff --git a/encodings/decimal-byte-parts/benches/dbp_assemble.rs b/encodings/decimal-byte-parts/benches/dbp_assemble.rs new file mode 100644 index 00000000000..48161a2c664 --- /dev/null +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -0,0 +1,169 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reassembling integer slices, including output-buffer allocation. +//! Array benchmarks also include validation, part execution, and casts. + +mod common; + +use divan::Bencher; +use divan::black_box; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::i256; +use vortex_buffer::buffer; +use vortex_decimal_byte_parts::_benchmarking::assemble_wide_decimal; +use vortex_decimal_byte_parts::_benchmarking::i128_to_parts; +use vortex_decimal_byte_parts::_benchmarking::i256_to_parts; +use vortex_decimal_byte_parts::_benchmarking::split_wide; +use vortex_error::vortex_panic; +use vortex_mask::Mask; + +use crate::common::cases; +use crate::common::i128_values; +use crate::common::i256_values; + +fn main() { + divan::main(); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = cases())] +fn dbp_assemble_kernel(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let validity = Mask::new_true(len); + match values_type { + DecimalType::I128 => { + let values = i128_values(len); + let (msp, [lower]) = split_wide(values.as_slice(), &validity, i128_to_parts); + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(lower.as_slice()).iter().map(|&word| [word]), + ) + }); + } + DecimalType::I256 => { + let values = i256_values(len); + let (msp, [first, second, third]) = + split_wide(values.as_slice(), &validity, i256_to_parts); + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(first.as_slice()) + .iter() + .zip(black_box(second.as_slice())) + .zip(black_box(third.as_slice())) + .map(|((&a, &b), &c)| [a, b, c]), + ) + }); + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } +} + +// The kernel widens the MSP while assembling; lower parts are already `u64`. +#[vortex_bench_support::cpu_features] +#[divan::bench(args = cases())] +fn dbp_assemble_kernel_narrow_msp(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let msp = buffer![-1i8; len]; + match values_type { + DecimalType::I128 => { + let lower = buffer![u64::from(u8::MAX); len]; + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(lower.as_slice()).iter().map(|&word| [word]), + ) + }); + } + DecimalType::I256 => { + let first = buffer![u64::from(u8::MAX); len]; + let second = buffer![u64::from(u16::MAX); len]; + let third = buffer![u64::from(u32::MAX); len]; + bencher.bench(|| { + assemble_wide_decimal::( + black_box(msp.as_slice()), + black_box(first.as_slice()) + .iter() + .zip(black_box(second.as_slice())) + .zip(black_box(third.as_slice())) + .map(|((&a, &b), &c)| [a, b, c]), + ) + }); + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } +} + +/// Benchmarks of full end-to-end dbp assembly, including kernel. These +/// are more volatile and should not be run in CI due to noise. +#[cfg(not(codspeed))] +mod arrays { + use divan::Bencher; + use divan::black_box; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::DecimalType; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_decimal_byte_parts::_benchmarking::assemble_decimal; + use vortex_decimal_byte_parts::split_decimal; + use vortex_error::VortexExpect; + use vortex_error::vortex_panic; + + use crate::common::arrays::cases; + use crate::common::arrays::decimal_array; + + #[divan::bench(args = cases())] + fn dbp_assemble(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let decimal = decimal_array(values_type, len, Validity::NonNullable); + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal, &mut ctx).vortex_expect("split benchmark input"); + bench_parts( + bencher, + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + ); + } + + #[divan::bench(args = cases())] + fn dbp_assemble_narrowed(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let (precision, lower_parts) = match values_type { + DecimalType::I64 => (18, vec![]), + DecimalType::I128 => (38, vec![buffer![u8::MAX; len].into_array()]), + DecimalType::I256 => ( + 76, + vec![ + buffer![u8::MAX; len].into_array(), + buffer![u16::MAX; len].into_array(), + buffer![u32::MAX; len].into_array(), + ], + ), + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + }; + bench_parts( + bencher, + buffer![-1i8; len].into_array(), + lower_parts, + DecimalDType::new(precision, 2), + ); + } + + fn bench_parts( + bencher: Bencher, + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) { + let session = array_session(); + bencher + .with_inputs(|| session.create_execution_ctx()) + .bench_refs(|ctx| { + assemble_decimal(black_box(&msp), black_box(&lower_parts), decimal_dtype, ctx) + .vortex_expect("assemble decimal byte parts") + }); + } +} diff --git a/encodings/decimal-byte-parts/benches/dbp_split.rs b/encodings/decimal-byte-parts/benches/dbp_split.rs new file mode 100644 index 00000000000..dd5b9abcf13 --- /dev/null +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splitting integer slices, including output-buffer allocation. +//! Array benchmarks also include validity execution and array construction. + +mod common; + +use divan::Bencher; +use divan::black_box; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::dtype::DecimalType; +use vortex_decimal_byte_parts::_benchmarking::i128_to_parts; +use vortex_decimal_byte_parts::_benchmarking::i256_to_parts; +use vortex_decimal_byte_parts::_benchmarking::split_wide; +use vortex_error::vortex_panic; +use vortex_mask::Mask; + +use crate::common::cases; +use crate::common::i128_values; +use crate::common::i256_values; + +fn main() { + divan::main(); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = cases())] +fn dbp_split_kernel_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Mask::new_true(len)); +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(args = cases())] +fn dbp_split_kernel_mixed_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let mut rng = StdRng::seed_from_u64(42); + let validity = Mask::from_iter((0..len).map(|_| rng.random_bool(0.5))); + bench_split(bencher, values_type, len, validity); +} + +fn bench_split(bencher: Bencher, values_type: DecimalType, len: usize, validity: Mask) { + match values_type { + DecimalType::I128 => { + let values = i128_values(len); + bencher.bench(|| { + split_wide( + black_box(values.as_slice()), + black_box(&validity), + i128_to_parts, + ) + }); + } + DecimalType::I256 => { + let values = i256_values(len); + bencher.bench(|| { + split_wide( + black_box(values.as_slice()), + black_box(&validity), + i256_to_parts, + ) + }); + } + _ => vortex_panic!("unsupported benchmark storage type: {values_type}"), + } +} + +/// Benchmarks of full end-to-end dbp assembly, including kernel. These +/// are more volatile and should not be run in CI due to noise. +#[cfg(not(codspeed))] +mod arrays { + use divan::Bencher; + use divan::black_box; + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::dtype::DecimalType; + use vortex_array::validity::Validity; + use vortex_decimal_byte_parts::split_decimal; + use vortex_error::VortexExpect; + + use crate::common::arrays::cases; + use crate::common::arrays::decimal_array; + + #[divan::bench(args = cases())] + fn dbp_split_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Validity::AllValid); + } + + #[divan::bench(args = cases())] + fn dbp_split_all_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + bench_split(bencher, values_type, len, Validity::AllInvalid); + } + + #[divan::bench(args = cases())] + fn dbp_split_mixed_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { + let mut rng = StdRng::seed_from_u64(42); + let validity = Validity::from_iter((0..len).map(|_| rng.random_bool(0.5))); + bench_split(bencher, values_type, len, validity); + } + + fn bench_split(bencher: Bencher, values_type: DecimalType, len: usize, validity: Validity) { + let decimal = decimal_array(values_type, len, validity); + let session = array_session(); + bencher + .with_inputs(|| session.create_execution_ctx()) + .bench_refs(|ctx| { + split_decimal(black_box(&decimal), ctx).vortex_expect("split decimal array") + }); + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs new file mode 100644 index 00000000000..c1826ae46d4 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -0,0 +1,749 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal byte-parts array types, validation, and VTable implementations. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::serde::ArrayChildren; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityChild; +use vortex_array::vtable::ValidityVTableFromChild; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; + +use super::MAX_LOWER_PARTS; +use super::assemble::assemble_decimal; +use super::assemble::assemble_wide_decimal_value; +use super::decimal_byte_parts_v2_id; +use super::rules::PARENT_RULES; + +/// A [`DecimalByteParts`]-encoded Vortex array. +pub type DecimalBytePartsArray = Array; + +/// This array encodes decimals by splitting them between 1-4 columns of primitive typed children. +/// +/// The most significant part (MSP) stores the most significant decimal bits. It is signed and is +/// nullable iff the decimal is nullable. +/// +/// Every lower part is a non-nullable unsigned integer holding a 64-bit window of the value. +/// Parts may have narrower integer dtypes when their values fit; their positions stay 64 bits apart. +/// +/// e.g. for a decimal i128 \[ 127..64 | 63..0 \] msp = 127..64 and lower_part\[0\] = 63..0 +/// +/// All parts live in slots, so the array carries no additional data. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsData; + +impl Display for DecimalBytePartsData { + fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl ArrayHash for DecimalBytePartsData { + fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} +} + +impl ArrayEq for DecimalBytePartsData { + fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { + true + } +} + +impl DecimalBytePartsData { + /// Validate the parts of a [`DecimalBytePartsArray`]. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array of length `len`, if `dtype` + /// does not match the MSP's nullability, if there are more than `MAX_LOWER_PARTS` + /// lower parts, or if any lower part is not a non-nullable unsigned integer array of length `len`. + pub fn validate<'a>( + msp: &ArrayRef, + lower_parts: impl ExactSizeIterator, + decimal_dtype: DecimalDType, + dtype: &DType, + len: usize, + ) -> VortexResult<()> { + if !msp.dtype().is_signed_int() { + vortex_bail!("msp must be a signed integer array") + } + + let expected_dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + vortex_ensure!( + dtype == &expected_dtype, + "expected dtype {expected_dtype}, got {dtype}" + ); + vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); + + let lower_part_count = lower_parts.len(); + + vortex_ensure!( + lower_part_count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {lower_part_count}" + ); + for (idx, part) in lower_parts.enumerate() { + vortex_ensure!( + part.dtype().is_unsigned_int() && !part.dtype().is_nullable(), + "lower part {idx} must have a non-nullable unsigned integer dtype, got {}", + part.dtype() + ); + vortex_ensure!( + part.len() == len, + "lower part {idx} has len {}, expected {len}", + part.len() + ); + } + Ok(()) + } +} + +#[derive(Clone, Debug)] +pub struct DecimalByteParts; + +impl DecimalByteParts { + /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. + /// + /// # Errors + /// + /// Returns an error if the MSP is not a signed integer array. + pub fn try_new( + msp: ArrayRef, + decimal_dtype: DecimalDType, + ) -> VortexResult { + Self::try_new_with_lower_parts(msp, Vec::new(), decimal_dtype) + } + + /// Construct a new [`DecimalBytePartsArray`] from an MSP array, its lower parts, and a + /// decimal dtype. + /// + /// Lower parts are ordered most significant first and must each be a non-nullable unsigned integer + /// array of the same length as the MSP. See [`super::split_decimal`] for producing them from a + /// decimal array. + /// + /// # Errors + /// + /// Returns an error if the parts do not describe a valid decimal, see + /// [`DecimalBytePartsData::validate`]. + pub fn try_new_with_lower_parts( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> VortexResult { + let len = msp.len(); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + Array::try_from_parts( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + ) + } + + /// Construct a [`DecimalBytePartsArray`] from parts whose invariants are already established. + /// + /// # Safety + /// + /// The MSP must have a signed integer dtype (`i8`, `i16`, `i32`, or `i64`). There must be + /// at most `MAX_LOWER_PARTS` lower parts, each a non-nullable unsigned integer array with the same + /// length as the MSP. Lower parts are ordered most significant first. + pub(super) unsafe fn new_unchecked( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> DecimalBytePartsArray { + let len = msp.len(); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + // SAFETY: the caller guarantees the part types, lengths, and count. The slot builder + // fills every required slot, and the length and nullability come from the MSP. + unsafe { + Array::from_parts_unchecked( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData) + .with_slots(slots), + ) + } + } +} + +impl VTable for DecimalByteParts { + type TypedArrayData = DecimalBytePartsData; + + type OperationsVTable = Self; + type ValidityVTable = ValidityVTableFromChild; + + fn id(&self) -> ArrayId { + decimal_byte_parts_v2_id() + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let Some(decimal_dtype) = dtype.as_decimal_opt() else { + vortex_bail!("expected decimal dtype, got {}", dtype) + }; + + let min_slots = DecimalBytePartsSlots::FIXED_COUNT; + let max_slots = min_slots + MAX_LOWER_PARTS; + vortex_ensure!( + (min_slots..=max_slots).contains(&slots.len()), + "expected {min_slots}..={max_slots} slots, got {}", + slots.len() + ); + for (idx, slot) in slots.iter().enumerate() { + vortex_ensure!(slot.is_some(), "missing required slot {idx}"); + } + + let slots = DecimalBytePartsSlotsView::from_slots(slots); + DecimalBytePartsData::validate( + slots.msp, + slots.lower_parts.iter(), + *decimal_dtype, + dtype, + len, + ) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("DecimalBytePartsArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + _array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + vortex_bail!("DecimalByteParts serialization requires DecimalBytePartsPlugin") + } + + fn deserialize( + &self, + _dtype: &DType, + _len: usize, + _metadata: &[u8], + _buffers: &[BufferHandle], + _children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + vortex_bail!("DecimalByteParts deserialization requires DecimalBytePartsPlugin") + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + DecimalBytePartsSlots::slot_name(idx) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + let lower_parts = array.lower_parts().to_vec(); + let assembled = assemble_decimal(array.msp(), &lower_parts, array.decimal_dtype(), ctx)?; + + Ok(ExecutionResult::done(assembled)) + } +} + +#[array_slots(DecimalByteParts)] +pub struct DecimalBytePartsSlots { + /// The most significant parts of the decimal values. + #[slot(0)] + pub msp: ArrayRef, + /// The remaining 64-bit windows of the decimal values, most significant first. + #[slot(1..)] + pub lower_parts: Vec, +} + +pub(crate) trait DecimalBytePartsArrayExt: DecimalBytePartsArraySlotsExt { + /// The decimal dtype of this array. + fn decimal_dtype(&self) -> DecimalDType { + *self + .as_ref() + .dtype() + .as_decimal_opt() + .vortex_expect("must be a decimal dtype") + } + + /// Rebuild the array by applying `f` to the MSP and every lower part, in slot order. + /// + /// This applies row operations such as slicing and filtering to all parts together, + /// preserving the decimal precision and scale. + fn map_parts( + &self, + mut f: impl FnMut(&ArrayRef) -> VortexResult, + ) -> VortexResult { + let msp = f(self.msp())?; + let lower_parts = self + .lower_parts() + .iter() + .map(&mut f) + .collect::>>()?; + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, self.decimal_dtype()) + } + + /// Rebuild the array with a replacement MSP, preserving its lower parts, precision and scale. + /// + /// Use this for operations such as masking and nullability casts that only affect the MSP. + /// The replacement MSP determines the result's nullability. + fn with_msp(&self, msp: ArrayRef) -> VortexResult { + DecimalByteParts::try_new_with_lower_parts( + msp, + self.lower_parts().to_vec(), + self.decimal_dtype(), + ) + } +} + +impl> DecimalBytePartsArrayExt for T {} + +impl OperationsVTable for DecimalByteParts { + type ProbeState = (); + + fn scalar_at( + array: ArrayView<'_, DecimalByteParts>, + index: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let scalar = array.msp().execute_scalar(index, ctx)?; + + // Widen the MSP's signed value (i8/i16/i32/i64) to i64 for scalar reconstruction. + // The array retains its original MSP storage type. + let primitive_scalar = scalar.as_primitive(); + let msp = primitive_scalar.as_::().vortex_expect("non-null"); + + // Zero-extend each narrowed lower value to its 64-bit window. + let lower_parts = array + .lower_parts() + .iter() + .map(|part| { + Ok(part + .execute_scalar(index, ctx)? + .as_primitive() + .as_::() + .vortex_expect("lower parts are non-nullable")) + }) + .collect::>>()?; + + let value = match lower_parts.as_slice() { + [] => DecimalValue::I64(msp), + [first] => DecimalValue::I128(assemble_wide_decimal_value(msp, [*first])), + [first, second] => { + DecimalValue::I256(assemble_wide_decimal_value(msp, [*first, *second])) + } + [first, second, third] => { + DecimalValue::I256(assemble_wide_decimal_value(msp, [*first, *second, *third])) + } + _ => vortex_bail!( + "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", + lower_parts.len() + ), + }; + + Scalar::try_new(array.dtype().clone(), Some(ScalarValue::Decimal(value))) + } +} + +impl ValidityChild for DecimalByteParts { + fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { + // validity stored in 0th child + array.msp().clone() + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::Array; + use vortex_array::ArrayParts; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::builtins::ArrayBuiltins; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::DecimalType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::scalar::DecimalValue; + use vortex_array::scalar::Scalar; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::DecimalByteParts; + use super::DecimalBytePartsArraySlotsExt; + use super::DecimalBytePartsData; + use crate::decimal_byte_parts::LOWER_PART_DTYPE; + use crate::decimal_byte_parts::MAX_LOWER_PARTS; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_parts; + + #[test] + fn test_scalar_at_decimal_parts() { + let decimal_dtype = DecimalDType::new(8, 2); + let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); + let array = DecimalByteParts::try_new( + PrimitiveArray::new( + buffer![100i32, 200i32, 400i32], + Validity::Array(BoolArray::from_iter(vec![false, true, true]).into_array()), + ) + .into_array(), + decimal_dtype, + ) + .unwrap() + .into_array(); + + assert_eq!( + Scalar::null(dtype.clone()), + array + .execute_scalar(0, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Decimal(DecimalValue::I64(200))) + ) + .unwrap(), + array + .execute_scalar(1, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), + array + .execute_scalar(2, &mut array_session().create_execution_ctx()) + .unwrap() + ); + } + + #[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::from_parts((7u128 << 64) | 8, (5i128 << 64) | 6)], + 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] + fn test_scalar_at_matches_canonical_for_each_part_count( + #[values(false, true)] narrow_msp: bool, + #[values(0, 1, 2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let validity = Validity::from_iter([false, true, true]); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![0i8, 3, -3], validity) + } else { + PrimitiveArray::new(buffer![0i64, 3, -3], validity) + }; + let lower = [4u64, 1, 2] + .into_iter() + .take(lower_count) + .map(|word| PrimitiveArray::new(buffer![word; 3], Validity::NonNullable).into_array()) + .collect(); + let dtype = DecimalDType::new(if lower_count <= 1 { 38 } else { 76 }, 0); + let array = DecimalByteParts::try_new_with_lower_parts(msp.into_array(), lower, dtype)?; + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)?; + for row in 0..array.len() { + assert_eq!( + array.execute_scalar(row, &mut ctx)?, + canonical.execute_scalar(row, &mut ctx)? + ); + } + Ok(()) + } + + #[test] + fn test_scalar_at_null_with_lower_parts() -> VortexResult<()> { + let array = i128_parts( + vec![1i128 << 100, 2, 3], + Validity::Array(BoolArray::from_iter([false, true, true]).into_array()), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::decimal( + DecimalValue::I128(2), + DecimalDType::new(38, 2), + Nullability::Nullable + ) + ); + Ok(()) + } + + #[rstest] + #[case([PType::U8, PType::U16, PType::U32])] + #[case([PType::U16, PType::U32, PType::U64])] + #[case([PType::U32, PType::U64, PType::U8])] + #[case([PType::U64, PType::U8, PType::U16])] + fn test_independently_narrowed_parts( + #[case] lower_ptypes: [PType; 3], + #[values(PType::I8, PType::I16, PType::I32, PType::I64)] msp_ptype: PType, + #[values(1, 2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let validity = Validity::from_iter([false, true, true]); + let msp = PrimitiveArray::new(buffer![0i64, 1, -1], validity.clone()) + .into_array() + .cast(DType::Primitive(msp_ptype, Nullability::Nullable))?; + // Set the unsigned type's highest bit to catch accidental sign extension. + let words = lower_ptypes.map(|ptype| 1u64 << (ptype.byte_width() * 8 - 1)); + let lower = lower_ptypes + .into_iter() + .zip(words) + .take(lower_count) + .map(|(ptype, word)| { + buffer![0u64, word, word] + .into_array() + .cast(DType::Primitive(ptype, Nullability::NonNullable)) + }) + .collect::>>()?; + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + let encoded = DecimalByteParts::try_new_with_lower_parts(msp, lower, dtype)?; + let expected = match lower_count { + 1 => DecimalArray::new( + buffer![ + 0i128, + (1i128 << 64) | i128::from(words[0]), + (-1i128 << 64) | i128::from(words[0]) + ], + dtype, + validity, + ), + 2 => { + let low = (u128::from(words[0]) << 64) | u128::from(words[1]); + DecimalArray::new( + buffer![ + i256::ZERO, + i256::from_parts(low, 1), + i256::from_parts(low, -1) + ], + dtype, + validity, + ) + } + _ => { + let low = (u128::from(words[1]) << 64) | u128::from(words[2]); + DecimalArray::new( + buffer![ + i256::ZERO, + i256::from_parts(low, (1i128 << 64) | i128::from(words[0])), + i256::from_parts(low, (-1i128 << 64) | i128::from(words[0])), + ], + dtype, + validity, + ) + } + }; + let actual = encoded + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(expected.clone(), actual, &mut ctx); + for row in 0..expected.len() { + assert_eq!( + encoded.execute_scalar(row, &mut ctx)?, + expected.execute_scalar(row, &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) + )] + fn test_rejects_invalid_parts( + #[case] lower_parts: Vec, + #[case] decimal_dtype: DecimalDType, + ) { + assert!( + DecimalByteParts::try_new_with_lower_parts(msp(), lower_parts, decimal_dtype).is_err() + ); + } + + #[rstest] + #[case::no_slots(vec![])] + #[case::missing_msp(vec![None])] + #[case::missing_lower(vec![Some(msp()), None])] + #[case::gap_in_lower(vec![Some(msp()), None, Some(lower_part())])] + fn test_rejects_missing_slots(#[case] slots: Vec>) { + let parts = ArrayParts::new( + DecimalByteParts, + DType::Decimal(DecimalDType::new(76, 2), Nullability::NonNullable), + 3, + DecimalBytePartsData, + ) + .with_slots(slots.into_iter().collect()); + assert!(Array::try_from_parts(parts).is_err()); + } + + #[test] + fn test_wide_decimal_buffer_types() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let i128_array = i128_parts(vec![1i128 << 100], Validity::NonNullable); + let canonical = i128_array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + + let i256_array = i256_parts(vec![i256::from_parts(0, 1 << 100)], 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::from_parts(9, 1)] + ); + 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/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs new file mode 100644 index 00000000000..35d93d4da1e --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs @@ -0,0 +1,464 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reassembling decimal arrays and values from their parts. + +use std::ops::BitOr; +use std::ops::Shl; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::dtype::i256; +use vortex_array::match_each_signed_integer_ptype; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::trusted_len::TrustedLen; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use super::LOWER_PART_BITS; +use super::LOWER_PART_DTYPE; +use super::MAX_LOWER_PARTS; + +/// Reassemble decimal parts into a decimal array. +/// +/// The MSP must have a signed integer dtype, and every lower part must have a non-nullable +/// unsigned integer dtype. All parts must have the same length. +/// +/// With no lower parts, the MSP buffer is reused as the decimal values. One lower part +/// assembles into `i128`. Two or three lower parts assemble into `i256`. +/// +/// If there are lower parts, each part (including the MSP) is widened into a 64 bit array. +/// For example, parts consisting of a `i8` MSP and a single `u32` lower part is assembled into +/// a decimal array with `64 + 64 = 128` bit storage. +/// +/// # Errors +/// +/// Returns an error for invalid part dtypes, lengths, or counts, or if executing a part fails. +pub fn assemble_decimal( + msp: &ArrayRef, + lower_parts: &[ArrayRef], + decimal_dtype: DecimalDType, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + vortex_ensure!( + msp.dtype().is_signed_int(), + "MSP must have a signed integer dtype" + ); + + let validity = msp.validity()?; + + if lower_parts.is_empty() { + return assemble_narrow_decimal(msp, validity, decimal_dtype, exec_ctx); + } + + vortex_ensure!( + lower_parts.len() <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", + lower_parts.len() + ); + let len = msp.len(); + for (idx, part) in lower_parts.iter().enumerate() { + vortex_ensure!( + part.dtype().is_unsigned_int() && !part.dtype().is_nullable(), + "lower part {idx} must have a non-nullable unsigned integer dtype, got {}", + part.dtype() + ); + vortex_ensure!( + part.len() == len, + "lower part {idx} has len {}, expected {len}", + part.len() + ); + } + + assemble_wide_decimal_from_arrays(msp, lower_parts, validity, decimal_dtype, exec_ctx) +} + +fn assemble_narrow_decimal( + msp: &ArrayRef, + validity: Validity, + decimal_dtype: DecimalDType, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + // TODO(mk): Broadcast a constant MSP directly instead of materializing its buffer. + let msp = msp.clone().execute::(exec_ctx)?; + Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { + DecimalArray::new(msp.to_buffer::

(), decimal_dtype, validity).into_array() + })) +} + +/// Execute the MSP at its signed integer width and cast lower parts to `u64` before assembly. +/// The number of lower parts determines the decimal storage type: one produces `i128`, while +/// two or three produce `i256`. +fn assemble_wide_decimal_from_arrays( + msp: &ArrayRef, + lower_parts: &[ArrayRef], + validity: Validity, + decimal_dtype: DecimalDType, + exec_ctx: &mut ExecutionCtx, +) -> VortexResult { + // TODO(mk): Broadcast constant parts directly instead of materializing their buffers. + let msp = msp.clone().execute::(exec_ctx)?; + // TODO(mk): Revisit dispatching on lower-part dtypes and widening values during assembly. + // Casting narrowed parts allocates temporary buffers and adds passes over the data. + // Nested dtype dispatch is significantly in benchmarks, but adds code and generic instantiations. + let lower = lower_parts + .iter() + .map(|part| { + part.cast(LOWER_PART_DTYPE)? + .execute::(exec_ctx) + }) + .collect::>>()?; + + Ok(match_each_signed_integer_ptype!(msp.ptype(), |Msp| { + let msp = msp.as_slice::(); + match lower.as_slice() { + [first] => DecimalArray::new( + assemble_wide_decimal::( + msp, + first.as_slice::().iter().map(|&word| [word]), + ), + decimal_dtype, + validity, + ) + .into_array(), + [first, second] => DecimalArray::new( + assemble_wide_decimal::( + msp, + first + .as_slice::() + .iter() + .zip(second.as_slice::()) + .map(|(&a, &b)| [a, b]), + ), + decimal_dtype, + validity, + ) + .into_array(), + [first, second, third] => DecimalArray::new( + assemble_wide_decimal::( + msp, + first + .as_slice::() + .iter() + .zip(second.as_slice::()) + .zip(third.as_slice::()) + .map(|((&a, &b), &c)| [a, b, c]), + ), + decimal_dtype, + validity, + ) + .into_array(), + _ => vortex_bail!("expected between one and {MAX_LOWER_PARTS} lower parts"), + } + })) +} + +/// Assemble a signed MSP slice and one array of `K` lower words per row. +/// +/// The caller zips the `u64` lower-part slices into rows. The iterator must have the same +/// length as the MSP slice. MSP values are widened to `i64` as they are read. +pub fn assemble_wide_decimal( + msp: &[Msp], + lower: impl TrustedLen, +) -> Buffer +where + T: NativeDecimalType + From + From + Shl + BitOr, + Msp: Copy + Into, +{ + let mut out = BufferMut::::with_capacity(msp.len()); + out.extend_trusted( + msp.iter() + .zip(lower) + .map(|(&value, parts)| assemble_wide_decimal_value(value.into(), parts)), + ); + out.freeze() +} + +/// Reassemble a decimal's unscaled integer from its signed MSP and `K` lower words. +/// +/// Sign-extend the MSP to `T`, then append each lower word by shifting left 64 bits and +/// filling the low bits. Lower words are ordered most significant first. Callers select +/// `i128` for one lower word and `i256` for two or three. +#[inline] +pub(crate) fn assemble_wide_decimal_value(msp: i64, lower: [u64; K]) -> T +where + T: NativeDecimalType + From + From + Shl + BitOr, +{ + let mut value: T = msp.into(); + for part in lower { + value = (value << LOWER_PART_BITS) | part.into(); + } + value +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::Constant; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::DecimalType; + use vortex_array::dtype::NativeDecimalType; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::match_each_decimal_value_type; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::assemble_decimal; + use crate::decimal_byte_parts::split_decimal; + + #[rstest] + #[case::empty_non_nullable(0, Validity::NonNullable)] + #[case::empty_nullable(0, Validity::AllValid)] + #[case::empty_all_null(0, Validity::AllInvalid)] + #[case::all_null(3, Validity::AllInvalid)] + #[case::all_null_array(3, Validity::Array(BoolArray::from_iter([false; 3]).into_array()))] + fn test_split_without_valid_rows( + #[case] len: usize, + #[case] validity: Validity, + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + values_type: DecimalType, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = match_each_decimal_value_type!(values_type, |T| { + DecimalArray::new( + Buffer::::zeroed(len), + DecimalDType::new(T::MAX_PRECISION, 0), + validity, + ) + }); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.msp.is::()); + assert!(parts.lower_parts.iter().all(|part| part.is::())); + assert_eq!(parts.msp.len(), len); + assert_eq!( + parts.msp.dtype().nullability(), + decimal.dtype().nullability() + ); + let round_tripped = round_trip(decimal.clone())?; + assert_eq!(round_tripped.values_type(), values_type); + assert_arrays_eq!(decimal, round_tripped, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::non_nullable(Validity::NonNullable)] + #[case::all_valid(Validity::AllValid)] + #[case::all_null(Validity::AllInvalid)] + #[case::mixed(Validity::from_iter((0..263).map(|i| i % 3 != 1)))] + #[case::sparse(Validity::from_iter((0..263).map(|i| i % 16 == 0)))] + #[case::null_prefix_and_suffix(Validity::from_iter((0..263).map(|i| (67..196).contains(&i))))] + fn test_split_zeroes_null_words( + #[case] validity: Validity, + #[values(false, true)] wide_256: bool, + #[values(0, 1, 63, 64, 65, 257)] len: usize, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = if wide_256 { + DecimalArray::new( + buffer![i256::from_i128(-1); 263], + DecimalDType::new(76, 2), + validity, + ) + } else { + DecimalArray::new(buffer![-1i128; 263], DecimalDType::new(38, 2), validity) + }; + let decimal = decimal + .slice(3..len + 3)? + .execute::(&mut ctx)?; + let mask = decimal.validity()?.execute_mask(len, &mut ctx)?; + let expected = PrimitiveArray::new( + mask.iter() + .map(|valid| if valid { u64::MAX } else { 0 }) + .collect::>(), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), if wide_256 { 3 } else { 1 }); + assert_eq!( + parts.msp.dtype(), + &DType::Primitive(PType::I64, decimal.dtype().nullability()) + ); + for lower in parts.lower_parts { + assert_arrays_eq!(expected.clone(), lower, &mut ctx); + } + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + Ok(()) + } + + fn round_trip(decimal: DecimalArray) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal, &mut ctx)?; + assemble_decimal( + &parts.msp, + &parts.lower_parts, + decimal.decimal_dtype(), + &mut ctx, + )? + .execute::(&mut ctx) + } + + #[rstest] + #[case::zero(0)] + #[case::one(1)] + #[case::minus_one(-1)] + #[case::limb_boundary(1i128 << 64)] + #[case::just_below_limb_boundary((1i128 << 64) - 1)] + #[case::negative_limb_boundary(-(1i128 << 64))] + #[case::max(i128::MAX)] + #[case::min(i128::MIN)] + fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) + } + + #[rstest] + #[case::zero(i256::ZERO)] + #[case::one(i256::ONE)] + #[case::minus_one(i256::ZERO - i256::ONE)] + #[case::max(i256::MAX)] + #[case::min(i256::MIN)] + #[case::word_1(i256::from_parts(1u128 << 64, 0))] + #[case::word_2(i256::from_parts(0, 1))] + #[case::word_3(i256::from_parts(0, 1i128 << 64))] + #[case::mixed(i256::from_parts(u128::MAX, -3))] + fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) + } + + #[rstest] + fn test_split_narrow_decimal_reuses_values( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] + validity: Validity, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.lower_parts.is_empty()); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); + let msp = parts.msp.execute::(&mut ctx)?; + assert_eq!( + msp.as_slice::().as_ptr(), + decimal.buffer::().as_ptr() + ); + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::signed(PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable))] + #[case::float(PrimitiveArray::new(buffer![0f32; 2], Validity::NonNullable))] + #[case::nullable_all_valid(PrimitiveArray::new(buffer![0u64; 2], Validity::AllValid))] + #[case::nullable_all_null(PrimitiveArray::new(buffer![0u64; 2], Validity::AllInvalid))] + #[case::nullable_mixed(PrimitiveArray::new(buffer![0u64; 2], Validity::from_iter([true, false])))] + fn test_assemble_rejects_invalid_lower_dtype( + #[case] invalid_lower: PrimitiveArray, + #[values(1, 2, 3)] lower_count: usize, + ) { + let msp = PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable); + let mut lower = + vec![PrimitiveArray::new(buffer![0u64; 2], Validity::NonNullable); lower_count]; + lower[lower_count - 1] = invalid_lower; + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + let lower = lower + .into_iter() + .map(IntoArray::into_array) + .collect::>(); + let mut ctx = array_session().create_execution_ctx(); + assert!(assemble_decimal(&msp.into_array(), &lower, dtype, &mut ctx).is_err()); + } + + #[rstest] + fn test_assemble_rejects_mismatched_lower_lengths( + #[values(1, 2, 3)] lower_count: usize, + #[values(0, 1, 3)] lower_len: usize, + ) { + let msp = PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable); + let mut lower = + vec![PrimitiveArray::new(buffer![0u64; 2], Validity::NonNullable); lower_count]; + lower[lower_count - 1] = + PrimitiveArray::new(buffer![0u64; lower_len], Validity::NonNullable); + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + let lower = lower + .into_iter() + .map(IntoArray::into_array) + .collect::>(); + let mut ctx = array_session().create_execution_ctx(); + assert!(assemble_decimal(&msp.into_array(), &lower, dtype, &mut ctx).is_err()); + } + + #[rstest] + fn test_assemble_i256_part_order_and_sign_extension( + #[values(false, true)] narrow_msp: bool, + #[values(2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![3i8, -3], Validity::NonNullable) + } else { + PrimitiveArray::new(buffer![3i64, -3], Validity::NonNullable) + }; + let lower = [4u64, 1, 2].map(|word| buffer![word; 2].into_array()); + let dtype = DecimalDType::new(76, 0); + let actual = assemble_decimal( + &msp.into_array(), + &lower[3 - lower_count..], + dtype, + &mut ctx, + )?; + let low = (1u128 << 64) | 2; + let expected = if lower_count == 2 { + buffer![i256::from_parts(low, 3), i256::from_parts(low, -3)] + } else { + buffer![ + i256::from_parts(low, (3i128 << 64) | 4), + i256::from_parts(low, (-3i128 << 64) | 4), + ] + }; + assert_arrays_eq!( + DecimalArray::new(expected, dtype, Validity::NonNullable), + actual, + &mut ctx + ); + Ok(()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs index 5ae1bf0101e..e4d81c274da 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/cast.rs @@ -10,17 +10,18 @@ use vortex_array::scalar_fn::fns::cast::CastReduce; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl CastReduce for DecimalByteParts { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { - // Check if this is just a nullability change + // Check if this is just a nullability change. + // TODO(mk): Support non-nullability changes as well, e.g. precision. if !dtype.eq_ignore_nullability(array.dtype()) { return Ok(None); } - // DecimalBytePartsArray can only have Decimal dtype, so we only handle decimal-to-decimal casts - let DType::Decimal(target_decimal, target_nullability) = dtype else { - // Cannot cast decimal to non-decimal types - delegate to canonical form + // DecimalBytePartsArray can only have Decimal dtype. + let DType::Decimal(_, target_nullability) = dtype else { return Ok(None); }; @@ -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(), - )) + array.with_msp(new_msp).map(|a| Some(a.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..65933865780 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/compare.rs @@ -39,6 +39,15 @@ 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. + // + // TODO(mk): Compare the signed MSP and then the unsigned lower parts in significance + // order to avoid canonicalizing wide decimals. + if !lhs.lower_parts().is_empty() { + return Ok(None); + } + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); let scalar_type = lhs.msp().dtype().with_nullability(nullability); @@ -158,10 +167,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 +231,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..6921e5dfda4 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,17 @@ 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::DecimalBytePartsArrayExt; + 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())) + array + .map_parts(|part| part.filter(mask.clone())) + .map(|d| Some(d.into_array())) } } @@ -32,9 +27,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::dtype::i256; + 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_parts; #[test] fn test_filter_decimal_byte_parts() { @@ -59,4 +58,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::from_parts(0, 1), + i256::from_parts(5, -1), + i256::from_parts(u128::MAX, 0), + i256::from_parts(7, 1 << 64), + i256::from_parts(0, 0), + ], + Validity::from_iter([true, false, true, true, false]), + ); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs index 065bc5e0051..3fe59111f6e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::ArrayRef; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; @@ -15,7 +16,9 @@ use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; /// DecimalByteParts-specific is_constant kernel. /// -/// Delegates to checking if the MSP (most significant part) is constant. +/// Delegates to checking that every part is constant: the MSP (most significant part) plus +/// each lower part. An all-null array is constant regardless of the bits its lower parts +/// hold in null slots. #[derive(Debug)] pub(crate) struct DecimalBytePartsIsConstantKernel; @@ -34,7 +37,27 @@ impl DynAggregateKernel for DecimalBytePartsIsConstantKernel { return Ok(None); }; - let result = is_constant(array.msp(), ctx)?; + let result = is_constant_parts(array, ctx)?; Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } + +fn is_constant_parts( + array: ArrayView<'_, DecimalByteParts>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !is_constant(array.msp(), ctx)? { + return Ok(false); + } + // Null slots hold undefined bits in the lower parts, so they cannot make a constant + // (all-null) array non-constant. + if array.array().all_invalid(ctx)? { + return Ok(true); + } + for part in array.lower_parts().iter() { + if !is_constant(part, ctx)? { + return Ok(false); + } + } + Ok(true) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 5e8d28e3526..cb71ba7880c 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -1,9 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::ArrayVTable; -use vortex_array::arrays::Dict; -use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; @@ -19,9 +16,4 @@ pub(crate) fn initialize(session: &VortexSession) { DecimalByteParts, CompareExecuteAdaptor(DecimalByteParts), ); - kernels.register_execute_parent_kernel( - Dict.id(), - DecimalByteParts, - TakeExecuteAdaptor(DecimalByteParts), - ); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index e7dc95af84f..9a022ef34ce 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs @@ -6,24 +6,17 @@ use vortex_array::ArrayView; use vortex_array::IntoArray; use vortex_array::scalar_fn::fns::mask::Mask as MaskExpr; use vortex_array::scalar_fn::fns::mask::MaskReduce; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::DecimalByteParts; +use crate::decimal_byte_parts::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; impl MaskReduce for DecimalByteParts { fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult> { + // Validity lives in the MSP, so only that part needs masking: the lower parts hold + // undefined bits in null slots, which is exactly what a masked-out row is. let masked_msp = MaskExpr::try_new(array.msp().clone(), mask.clone())?.into_array(); - Ok(Some( - DecimalByteParts::try_new( - masked_msp, - *array - .dtype() - .as_decimal_opt() - .vortex_expect("must be a decimal dtype"), - )? - .into_array(), - )) + array.with_msp(masked_msp).map(|a| Some(a.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs index 6c2d0dabb31..c8385c2d6d6 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -7,6 +7,7 @@ mod filter; pub(crate) mod is_constant; pub(crate) mod kernel; mod mask; +mod slice; mod take; #[cfg(test)] @@ -19,10 +20,35 @@ 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_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::from_parts(0, 1), + i256::from_parts(0, -1), + i256::from_parts(u128::MAX, 0), + i256::from_parts(7, 1 << 64), + i256::from_parts(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/slice.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs similarity index 53% rename from encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs rename to encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs index 14807421c73..1a2efc9034e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/slice.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/slice.rs @@ -7,23 +7,15 @@ 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::DecimalBytePartsArrayExt; 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(), - )) + array + .map_parts(|part| part.slice(range.clone())) + .map(|d| Some(d.into_array())) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/take.rs index 7a18f7bf91b..a74915e9f22 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,106 @@ 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::DecimalBytePartsArrayExt; use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; -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 { + fn take(array: ArrayView<'_, Self>, indices: &ArrayRef) -> VortexResult> { + // Taking with nullable indices makes every taken part nullable, but lower parts must + // stay non-nullable — validity belongs to the MSP alone. Fall back to the + // canonical path rather than rebuilding parts we would have to strip nullability from. + // + // TODO(mk): Support lower parts using fill_null for nullable indices. + if indices.dtype().is_nullable() && !array.lower_parts().is_empty() { + return Ok(None); + } + + array + .map_parts(|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::dtype::i256; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::DecimalByteParts; + + /// 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 = DecimalByteParts::encode(&decimal, &mut session.create_execution_ctx())? + .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::from_parts(1 << 70, 1), i256::from_parts(2, 0), i256::from_parts(3, 0)]), + 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 = DecimalByteParts::encode(&decimal, &mut ctx)? + .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/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index e1d081ddaca..65bb8222f7a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -1,376 +1,63 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::fmt::Display; -use std::fmt::Formatter; -use std::hash::Hasher; +//! Decimal byte-parts encoding. +//! +//! A `DecimalByteParts` array stores each value as a signed most significant part (MSP) +//! followed by `k` unsigned 64-bit lower parts ordered most significant first. The encoded +//! value is +//! +//! ```text +//! msp * 2^(64k) + Σ_{i; - -impl ArrayHash for DecimalBytePartsData { - fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} -} - -impl ArrayEq for DecimalBytePartsData { - fn array_eq(&self, _other: &Self, _accuracy: EqMode) -> bool { - true - } -} - -#[derive(Clone, prost::Message)] -pub struct DecimalBytesPartsMetadata { - #[prost(enumeration = "PType", tag = "1")] - zeroth_child_ptype: i32, - #[prost(uint32, tag = "2")] - lower_part_count: u32, -} - -impl VTable for DecimalByteParts { - type TypedArrayData = DecimalBytePartsData; - - type OperationsVTable = Self; - type ValidityVTable = ValidityVTableFromChild; - - fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); - *ID - } - - fn validate( - &self, - _data: &Self::TypedArrayData, - dtype: &DType, - len: usize, - slots: &[Option], - ) -> VortexResult<()> { - let Some(decimal_dtype) = dtype.as_decimal_opt() else { - vortex_bail!("expected decimal dtype, got {}", dtype) - }; - let msp = DecimalBytePartsSlotsView::from_slots(slots).msp; - DecimalBytePartsData::validate(msp, *decimal_dtype, dtype, len) - } - - fn nbuffers(_array: ArrayView<'_, Self>) -> usize { - 0 - } - - fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - vortex_panic!("DecimalBytePartsArray buffer index {idx} out of bounds") - } - - fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") - } - - fn with_buffers( - &self, - array: ArrayView<'_, Self>, - buffers: &[BufferHandle], - ) -> VortexResult> { - vortex_array::vtable::with_empty_buffers(self, array, buffers) - } - - fn serialize( - array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - Ok(Some( - DecimalBytesPartsMetadata { - zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, - lower_part_count: 0, - } - .encode_to_vec(), - )) - } - - fn deserialize( - &self, - dtype: &DType, - len: usize, - metadata: &[u8], - _buffers: &[BufferHandle], - children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - let metadata = DecimalBytesPartsMetadata::decode(metadata)?; - let Some(decimal_dtype) = dtype.as_decimal_opt() else { - vortex_bail!("decoding decimal but given non decimal dtype {}", dtype) - }; - - let encoded_dtype = DType::Primitive(metadata.zeroth_child_ptype(), dtype.nullability()); - - let msp = children.get(0, &encoded_dtype, len)?; - - assert_eq!( - metadata.lower_part_count, 0, - "lower_part_count > 0 not currently supported" - ); - - 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)) - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - DecimalBytePartsSlots::NAMES[idx].to_string() - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - PARENT_RULES.evaluate(array, parent, child_idx) - } - fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - to_canonical_decimal(&array, ctx).map(ExecutionResult::done) - } -} - -#[array_slots(DecimalByteParts)] -pub struct DecimalBytePartsSlots { - /// The most significant parts of the decimal values. - #[slot(0)] - pub msp: ArrayRef, -} - -/// This array encodes decimals as between 1-4 columns of primitive typed children. -/// The most significant part (msp) sorting the most significant decimal bits. -/// This array must be signed and is nullable iff the decimal is nullable. -/// -/// e.g. for a decimal i128 \[ 127..64 | 64..0 \] msp = 127..64 and lower_part\[0\] = 64..0 -#[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, -} - -impl Display for DecimalBytePartsData { - fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result { - Ok(()) - } -} - -pub struct DecimalBytePartsDataParts { - pub msp: ArrayRef, -} - -impl DecimalBytePartsData { - pub fn validate( - msp: &ArrayRef, - decimal_dtype: DecimalDType, - dtype: &DType, - len: usize, - ) -> VortexResult<()> { - if !msp.dtype().is_signed_int() { - vortex_bail!("decimal bytes parts, first part must be a signed array") - } - - let expected_dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); - vortex_ensure!( - dtype == &expected_dtype, - "expected dtype {expected_dtype}, got {dtype}" - ); - vortex_ensure!(msp.len() == len, "expected len {len}, got {}", msp.len()); - 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()); - vortex_ensure!( - msp_dtype.is_signed_int(), - "decimal bytes parts, first part must be a signed array" - ); - let _ = msp_len; - drop(expected_dtype); - Ok(Self { - _lower_parts: Vec::new(), - }) - } -} - -#[derive(Clone, Debug)] -pub struct DecimalByteParts; - -impl DecimalByteParts { - /// Construct a new [`DecimalBytePartsArray`] from an MSP array and decimal dtype. - pub fn try_new( - msp: ArrayRef, - decimal_dtype: DecimalDType, - ) -> VortexResult { - 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), - ) - }) - } -} - -/// 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() - })) -} - -impl OperationsVTable for DecimalByteParts { - type ProbeState = (); +mod array; +mod assemble; +pub(crate) mod compute; +mod plugin; +#[cfg(test)] +mod prop_tests; +mod rules; +mod split; +#[cfg(test)] +mod testing; - fn scalar_at( - array: ArrayView<'_, DecimalByteParts>, - index: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - // TODO(joe): support parts len != 1 - let scalar = array.msp().execute_scalar(index, ctx)?; +pub use array::*; +pub use plugin::DecimalBytePartsPlugin; +pub use plugin::DecimalBytePartsV2Metadata; +pub use plugin::decimal_byte_parts_v1_id; +pub use plugin::decimal_byte_parts_v2_id; +pub use split::DecimalParts; +pub use split::split_decimal; - // 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))), - ) - } +#[doc(hidden)] +pub mod _benchmarking { + pub use super::assemble::assemble_decimal; + pub use super::assemble::assemble_wide_decimal; + pub use super::split::i128_to_parts; + pub use super::split::i256_to_parts; + pub use super::split::split_wide; } -impl ValidityChild for DecimalByteParts { - fn validity_child(array: ArrayView<'_, DecimalByteParts>) -> ArrayRef { - // validity stored in 0th child - array.msp().clone() - } -} +/// The maximum number of 64-bit lower parts an encoded `i128` decimal can carry. +const MAX_I128_LOWER_PARTS: usize = 1; -#[cfg(test)] -mod tests { - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::array_session; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::dtype::DType; - use vortex_array::dtype::DecimalDType; - use vortex_array::dtype::Nullability; - use vortex_array::scalar::DecimalValue; - use vortex_array::scalar::Scalar; - use vortex_array::scalar::ScalarValue; - use vortex_array::validity::Validity; - use vortex_buffer::buffer; +/// The maximum number of 64-bit lower parts an encoded `i256` decimal can carry. +const MAX_I256_LOWER_PARTS: usize = 3; - use crate::DecimalByteParts; +/// The maximum number of 64-bit lower parts an encoded decimal can carry. +const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS; - #[test] - fn test_scalar_at_decimal_parts() { - let decimal_dtype = DecimalDType::new(8, 2); - let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); - let array = DecimalByteParts::try_new( - PrimitiveArray::new( - buffer![100i32, 200i32, 400i32], - Validity::Array(BoolArray::from_iter(vec![false, true, true]).into_array()), - ) - .into_array(), - decimal_dtype, - ) - .unwrap() - .into_array(); +/// Number of bits stored in each lower part. +const LOWER_PART_BITS: usize = 64; - assert_eq!( - Scalar::null(dtype.clone()), - array - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new( - dtype.clone(), - Some(ScalarValue::Decimal(DecimalValue::I64(200))) - ) - .unwrap(), - array - .execute_scalar(1, &mut array_session().create_execution_ctx()) - .unwrap() - ); - assert_eq!( - Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), - array - .execute_scalar(2, &mut array_session().create_execution_ctx()) - .unwrap() - ); - } -} +/// Dtype of the lower parts produced by splitting, before any narrowing. +const LOWER_PART_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable); diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs new file mode 100644 index 00000000000..3896cd51c96 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! ArrayPlugin implementation for DBP that handles different wire formats. + +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayId; +use vortex_array::ArrayPlugin; +use vortex_array::ArrayRef; +use vortex_array::ArraySerialization; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::DecimalByteParts; +use super::DecimalBytePartsArraySlotsExt; + +#[cfg(test)] +mod tests; + +mod v1; +mod v2; + +pub use v2::DecimalBytePartsV2Metadata; + +/// The frozen single-child DBP serialized ID. +pub fn decimal_byte_parts_v1_id() -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); + *ID +} + +/// The current in-memory DBP ID and serialized ID for arrays with lower parts. +pub fn decimal_byte_parts_v2_id() -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts.v2"); + *ID +} + +/// Serde for the [`DecimalByteParts`] array using the frozen v1 and v2 wire formats. +/// +/// Each version owns its metadata schema and serde functions. The plugin writes v1 whenever an +/// array has no lower parts, so such arrays stay readable by older readers, and v2 otherwise. +/// The v2 format itself accepts any lower part count up to the maximum. +/// +/// Register this plugin, or call [`crate::initialize`], to enable both formats. Direct registration +/// of [`DecimalByteParts`] does not support serde. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsPlugin; + +impl ArrayPlugin for DecimalBytePartsPlugin { + fn id(&self) -> ArrayId { + VTable::id(&DecimalByteParts) + } + + fn serialized_ids(&self) -> Vec { + vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()] + } + + fn serialize( + &self, + array: &ArrayRef, + _session: &VortexSession, + ) -> VortexResult> { + let view = array.as_opt::().ok_or_else(|| { + vortex_err!( + "DecimalByteParts plugin cannot serialize {}", + array.encoding_id() + ) + })?; + let serialized = if view.lower_parts().is_empty() { + v1::serialize(view)? + } else { + v2::serialize(view)? + }; + Ok(Some(serialized)) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let array = if parts.serialized_id == decimal_byte_parts_v1_id() { + v1::deserialize(parts)? + } else if parts.serialized_id == decimal_byte_parts_v2_id() { + v2::deserialize(parts)? + } else { + vortex_bail!( + "DecimalByteParts plugin does not recognize serialized ID {}", + parts.serialized_id + ) + }; + Ok(array.into_array()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs new file mode 100644 index 00000000000..92c9d1edbb7 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs @@ -0,0 +1,324 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use prost::Message as _; +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayVTable; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::Primitive; +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::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::VortexExpect; +use vortex_session::registry::ReadContext; + +use super::*; +use crate::DecimalBytePartsArray; +use crate::DecimalBytePartsArraySlotsExt; +use crate::decimal_byte_parts::MAX_LOWER_PARTS; + +#[rstest] +#[case::no_lower_parts(DecimalByteParts::try_new( + buffer![1i32, 2, 3].into_array(), DecimalDType::new(9, 2), +))] +#[case::one_lower_part(DecimalByteParts::try_new_with_lower_parts( + msp(), vec![lower_part()], DecimalDType::new(38, 2), +))] +#[case::wider_i64_storage(DecimalByteParts::encode( + &DecimalArray::new(buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable), + &mut array_session().create_execution_ctx(), +))] +#[case::wider_i128_storage(DecimalByteParts::encode( + &DecimalArray::new(buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable), + &mut array_session().create_execution_ctx(), +))] +#[case::wider_i256_storage(DecimalByteParts::encode( + &DecimalArray::new( + buffer![i256::from_i128(-99), i256::ZERO, i256::from_i128(99)], + DecimalDType::new(2, 0), Validity::NonNullable, + ), + &mut array_session().create_execution_ctx(), +))] +#[case::redundant_lower_parts(DecimalByteParts::try_new_with_lower_parts( + buffer![0i64; 3].into_array(), + vec![buffer![0u64; 3].into_array(), buffer![0u64; 3].into_array(), lower_part()], + DecimalDType::new(38, 2), +))] +#[case::narrowed_lower_parts(DecimalByteParts::try_new_with_lower_parts( + buffer![-1i16, 0, 1].into_array(), + vec![ + buffer![u8::MAX, 128, 0].into_array(), + ConstantArray::new(u16::MAX, 3).into_array(), + buffer![0u32, 1 << 31, u32::MAX].into_array(), + ], + DecimalDType::new(76, 2), +))] +#[case::nullable_mixed_lower_parts(DecimalByteParts::try_new_with_lower_parts( + PrimitiveArray::new( + buffer![-1i8, 0, 1], Validity::from_iter([true, false, true]), + ).into_array(), + vec![ + buffer![u64::MAX, 1 << 63, 0].into_array(), + buffer![0u8, 128, u8::MAX].into_array(), + buffer![u32::MAX, 1 << 31, 0].into_array(), + ], + DecimalDType::new(76, 2), +))] +fn serde_round_trip(#[case] array: VortexResult) -> VortexResult<()> { + let session = session(); + let array = array?; + let lower_part_count = array.lower_parts().len(); + let lower_part_dtypes: Vec<_> = array + .lower_parts() + .iter() + .map(|part| part.dtype().clone()) + .collect(); + let array = array.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + + let expected_id = if lower_part_count == 0 { + decimal_byte_parts_v1_id() + } else { + decimal_byte_parts_v2_id() + }; + assert_eq!( + session + .array_serialize(&array)? + .vortex_expect("byte parts arrays are serializable") + .serialized_id, + expected_id + ); + + 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() + .iter() + .map(|part| part.dtype().clone()) + .collect::>(), + lower_part_dtypes, + "lower-part dtypes and order must survive serde" + ); + + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) +} + +#[test] +fn v1_metadata_is_unchanged() -> VortexResult<()> { + let session = session(); + let array = DecimalByteParts::try_new(msp(), DecimalDType::new(19, 2))?.into_array(); + let serialized = session + .array_serialize(&array)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialized.serialized_id, decimal_byte_parts_v1_id()); + // v1 metadata for an i64 MSP: field 1 = 7, with no lower-part fields emitted. + assert_eq!(serialized.metadata, [8, 7]); + Ok(()) +} + +#[rstest] +#[case::v1_lower_part_count( + decimal_byte_parts_v1_id(), + v1_metadata(1), + vec![msp(), lower_part()], + "must not carry lower parts" +)] +#[case::v1_extra_child( + decimal_byte_parts_v1_id(), + v1_metadata(0), + vec![msp(), lower_part()], + "exactly one child" +)] +#[case::v2_missing_child( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::U64 as i32]), + vec![msp()], + "expected 2 children, got 1" +)] +#[case::v2_extra_child( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::U64 as i32]), + vec![msp(), lower_part(), lower_part()], + "expected 2 children, got 3" +)] +#[case::v2_too_many_lower_parts( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::U64 as i32; MAX_LOWER_PARTS + 1]), + vec![msp(), lower_part(), lower_part(), lower_part(), lower_part()], + "lower parts, got 4" +)] +#[case::v2_signed_lower_part( + decimal_byte_parts_v2_id(), + v2_metadata(vec![PType::I64 as i32]), + vec![msp(), lower_part()], + "unsigned integer dtype" +)] +#[case::v2_unknown_ptype( + decimal_byte_parts_v2_id(), + v2_metadata(vec![i32::MAX]), + vec![msp(), lower_part()], + "invalid PType" +)] +fn decoder_rejects_malformed_payloads( + #[case] serialized_id: ArrayId, + #[case] metadata: Vec, + #[case] children: Vec, + #[case] expected_error: &str, +) { + let result = deserialize_with(serialized_id, &metadata, children); + assert!( + result + .as_ref() + .is_err_and(|err| err.to_string().contains(expected_error)), + "expected {expected_error}, got {result:?}" + ); +} + +#[test] +fn serialization_requires_v2_permission() -> VortexResult<()> { + let session = session(); + let array = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + + let restricted = ArrayContext::empty().with_allowed_ids( + [decimal_byte_parts_v1_id(), ArrayVTable::id(&Primitive)] + .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( + [ + decimal_byte_parts_v1_id(), + decimal_byte_parts_v2_id(), + ArrayVTable::id(&Primitive), + ] + .into_iter() + .collect(), + ); + array.serialize(&permissive, &session, &SerializeOptions::default())?; + assert!( + permissive.to_ids().contains(&decimal_byte_parts_v2_id()), + "the file's encoding table must carry the v2 format id" + ); + + Ok(()) +} + +#[test] +fn bare_vtable_refuses_serde() -> VortexResult<()> { + let session = array_session(); + session.arrays().register(DecimalByteParts); + let msp = msp(); + let array = DecimalByteParts::try_new(msp.clone(), DecimalDType::new(19, 2))?.into_array(); + let result = session.array_serialize(&array); + assert!( + result.as_ref().is_err_and(|err| err + .to_string() + .contains("DecimalByteParts serialization requires DecimalBytePartsPlugin")), + "expected unsupported VTable serialization, got {result:?}" + ); + + let id = VTable::id(&DecimalByteParts); + let plugin = session + .arrays() + .registry() + .get(&id) + .vortex_expect("registered"); + let children = vec![msp]; + let result = plugin.deserialize( + ArrayDeserialization::new(id, array.dtype(), array.len(), &[8, 7], &[], &children), + &session, + ); + assert!( + result.as_ref().is_err_and(|err| err + .to_string() + .contains("DecimalByteParts deserialization requires DecimalBytePartsPlugin")), + "expected unsupported VTable deserialization, got {result:?}" + ); + Ok(()) +} + +fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() +} + +fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() +} + +/// v1 metadata for an i64 MSP: field 1 = 7, then field 2 only when the count is non-zero, as +/// proto3 omits default values. +fn v1_metadata(lower_part_count: u8) -> Vec { + let mut metadata = vec![8, 7]; + if lower_part_count > 0 { + metadata.extend([16, lower_part_count]); + } + metadata +} + +fn v2_metadata(lower_part_ptypes: Vec) -> Vec { + DecimalBytePartsV2Metadata { + msp_ptype: PType::I64 as i32, + lower_part_ptypes, + } + .encode_to_vec() +} + +fn session() -> VortexSession { + let session = array_session(); + crate::initialize(&session); + session +} + +/// Decode a hand-built payload of three rows through the plugin. +fn deserialize_with( + serialized_id: ArrayId, + metadata: &[u8], + children: Vec, +) -> VortexResult { + let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); + DecimalBytePartsPlugin.deserialize( + ArrayDeserialization::new(serialized_id, &dtype, 3, metadata, &[], &children), + &array_session(), + ) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs new file mode 100644 index 00000000000..ec816c62736 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs @@ -0,0 +1,84 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serde for the frozen single-child DBP wire format. + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayParts; +use vortex_array::ArraySerialization; +use vortex_array::ArrayView; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::smallvec::smallvec; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::decimal_byte_parts_v1_id; +use crate::DecimalByteParts; +use crate::DecimalBytePartsArray; +use crate::DecimalBytePartsArraySlotsExt; +use crate::DecimalBytePartsData; + +#[derive(Clone, prost::Message)] +struct DecimalBytePartsMetadata { + #[prost(enumeration = "PType", tag = "1")] + zeroth_child_ptype: i32, + #[prost(uint32, tag = "2")] + lower_part_count: u32, +} + +pub(super) fn serialize( + array: ArrayView<'_, DecimalByteParts>, +) -> VortexResult { + vortex_ensure!( + array.lower_parts().is_empty(), + "v1 must not carry lower parts" + ); + let msp = array.msp(); + let metadata = DecimalBytePartsMetadata { + zeroth_child_ptype: PType::try_from(msp.dtype())? as i32, + lower_part_count: 0, + } + .encode_to_vec(); + Ok(ArraySerialization::new( + decimal_byte_parts_v1_id(), + metadata, + vec![], + vec![msp.clone()], + )) +} + +pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult { + vortex_ensure!( + parts.serialized_id == decimal_byte_parts_v1_id(), + "expected the v1 format" + ); + let metadata = DecimalBytePartsMetadata::decode(parts.metadata)?; + vortex_ensure!( + parts.dtype.as_decimal_opt().is_some(), + "expected a decimal dtype" + ); + vortex_ensure!( + metadata.lower_part_count == 0, + "v1 must not carry lower parts" + ); + vortex_ensure!(parts.children.len() == 1, "v1 must carry exactly one child"); + let ptype = PType::try_from(metadata.zeroth_child_ptype)?; + vortex_ensure!( + ptype.is_signed_int(), + "MSP must have a signed integer dtype" + ); + let encoded_dtype = DType::Primitive(ptype, parts.dtype.nullability()); + let msp = parts.children.get(0, &encoded_dtype, parts.len)?; + Array::try_from_parts( + ArrayParts::new( + DecimalByteParts, + parts.dtype.clone(), + parts.len, + DecimalBytePartsData, + ) + .with_slots(smallvec![Some(msp)]), + ) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs new file mode 100644 index 00000000000..c82db1cbce3 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serde for DBP values with unsigned lower parts. + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayDeserialization; +use vortex_array::ArrayParts; +use vortex_array::ArraySerialization; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use super::decimal_byte_parts_v2_id; +use crate::DecimalByteParts; +use crate::DecimalBytePartsArray; +use crate::DecimalBytePartsArraySlotsExt; +use crate::DecimalBytePartsData; +use crate::decimal_byte_parts::MAX_LOWER_PARTS; + +/// Metadata for decimal byte parts with lower parts. +#[derive(Clone, prost::Message)] +pub struct DecimalBytePartsV2Metadata { + /// Ptype of the most significant part. + #[prost(enumeration = "PType", tag = "1")] + pub(super) msp_ptype: i32, + /// Ptypes of the lower parts, ordered most significant first. + #[prost(enumeration = "PType", repeated, tag = "2")] + pub(super) lower_part_ptypes: Vec, +} + +pub(super) fn serialize( + array: ArrayView<'_, DecimalByteParts>, +) -> VortexResult { + let lower_parts = array.lower_parts(); + + let metadata = DecimalBytePartsV2Metadata { + msp_ptype: PType::try_from(array.msp().dtype())? as i32, + lower_part_ptypes: lower_parts + .iter() + .map(|part| PType::try_from(part.dtype()).map(|ptype| ptype as i32)) + .collect::>()?, + } + .encode_to_vec(); + + let mut children = Vec::with_capacity(1 + lower_parts.len()); + children.push(array.msp().clone()); + children.extend(lower_parts.iter().cloned()); + + Ok(ArraySerialization::new( + decimal_byte_parts_v2_id(), + metadata, + vec![], + children, + )) +} + +pub(super) fn deserialize(parts: ArrayDeserialization<'_>) -> VortexResult { + let metadata = DecimalBytePartsV2Metadata::decode(parts.metadata)?; + vortex_ensure!( + parts.dtype.as_decimal_opt().is_some(), + "expected a decimal dtype" + ); + + let lower_part_count = metadata.lower_part_ptypes.len(); + vortex_ensure!( + lower_part_count <= MAX_LOWER_PARTS, + "v2 carries at most {MAX_LOWER_PARTS} lower parts, got {lower_part_count}" + ); + vortex_ensure!( + parts.children.len() == 1 + lower_part_count, + "expected {} children, got {}", + 1 + lower_part_count, + parts.children.len() + ); + + let msp_ptype = PType::try_from(metadata.msp_ptype)?; + vortex_ensure!( + msp_ptype.is_signed_int(), + "MSP must have a signed integer dtype, got {msp_ptype}" + ); + let msp_dtype = DType::Primitive(msp_ptype, parts.dtype.nullability()); + + let mut slots = ArraySlots::with_capacity(parts.children.len()); + slots.push(Some(parts.children.get(0, &msp_dtype, parts.len)?)); + for (idx, raw_ptype) in metadata.lower_part_ptypes.into_iter().enumerate() { + let ptype = PType::try_from(raw_ptype) + .map_err(|_| vortex_err!("invalid PType {raw_ptype} for lower part {idx}"))?; + vortex_ensure!( + ptype.is_unsigned_int(), + "lower part {idx} must have an unsigned integer dtype, got {ptype}" + ); + slots.push(Some(parts.children.get( + 1 + idx, + &DType::Primitive(ptype, Nullability::NonNullable), + parts.len, + )?)); + } + Array::try_from_parts( + ArrayParts::new( + DecimalByteParts, + parts.dtype.clone(), + parts.len, + DecimalBytePartsData, + ) + .with_slots(slots), + ) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs new file mode 100644 index 00000000000..b1625485ee7 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs @@ -0,0 +1,178 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Property tests for decimal byte-parts round trips. + +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_error::VortexExpect; + +use super::DecimalByteParts; +use super::DecimalBytePartsArray; + +/// 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(); + crate::initialize(&session); + session.create_execution_ctx() +} + +/// 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 encoded = DecimalByteParts::encode(&decimal, &mut ctx).vortex_expect("encode"); + let round_tripped = canonicalize(encoded.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_encoded = DecimalByteParts::encode(&decoded, &mut ctx).vortex_expect("encode"); + let re_decoded = canonicalize(re_encoded.into_array(), &mut ctx); + + assert_arrays_eq!(decoded, re_decoded, &mut ctx); +} + +// TODO(joe): restore the coverage removed alongside these two round trips. Each of the +// following was a property here and caught mutations that the round trips do not: +// +// - `scalar_at` against bulk canonicalization. `combine_i128`/`combine_i256` are a second +// implementation of the assembly loops and can drift from them silently. +// - filter, slice and take against the same operation on the canonical array. These caught +// part-order and word-placement mutations, though the round trips catch those too. +// - a serialize/decode round trip, which is the only property that exercised the metadata +// carrying the lower part count. +// - sign extension above a most significant part below the top word, checked against an +// expectation computed independently of the assembly loop. This is the one real gap: a +// round trip compares decode against decode, so a decode-side sign-extension bug is +// invisible to it. Dropping the sign extension is caught by neither property here. diff --git a/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/split.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs new file mode 100644 index 00000000000..f33101e4a3a --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -0,0 +1,311 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splitting canonical decimal arrays into signed and unsigned parts. + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::PType; +use vortex_array::dtype::i256; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::DecimalByteParts; +use super::DecimalBytePartsArray; +use super::LOWER_PART_BITS; +use super::MAX_I128_LOWER_PARTS; +use super::MAX_I256_LOWER_PARTS; + +impl DecimalByteParts { + /// Encode a [`DecimalArray`] as byte parts, splitting wide values into lower parts. + /// + /// # Errors + /// + /// Returns an error if the decimal cannot be split. + pub fn encode( + decimal: &DecimalArray, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let parts = split_decimal(decimal, exec_ctx)?; + // SAFETY: splitting produces a signed MSP and zero, one, or three non-nullable u64 lower + // parts, all with the decimal's length and in most-significant-first order. This also + // holds for the constant parts used for empty and all-null inputs. The decimal dtype is + // preserved. + Ok(unsafe { Self::new_unchecked(parts.msp, parts.lower_parts, decimal.decimal_dtype()) }) + } +} + +/// A decimal array decomposed into byte parts. +pub struct DecimalParts { + /// The signed most significant part. This carries the validity of the whole array. + pub msp: ArrayRef, + /// The unsigned 64-bit lower parts, most significant first. + pub lower_parts: Vec, +} + +impl DecimalParts { + /// Construct decimal parts from the MSP buffer constituting a narrow decimal (`i64` or narrower). + /// Narrow decimals have an MSP at most as wide as `i64` and no lower parts. + fn from_narrow(values: Buffer, validity: Validity) -> Self { + Self { + msp: PrimitiveArray::new(values, validity).into_array(), + lower_parts: Vec::new(), + } + } + + /// Construct decimal parts arrays from the buffers constituting a wide decimal (`i128` or `i256`). + /// Wide decimals have an `i64` MSP and up to `MAX_LOWER_PARTS` `u64` lower parts. + fn from_wide( + msp: Buffer, + lower_parts: impl IntoIterator>, + validity: Validity, + ) -> 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(), + } + } +} + +/// Split a canonical decimal array into a signed most significant part (MSP) and unsigned 64-bit +/// lower parts. The MSP is at most 64 bits. +/// +/// 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 split into an `i64` MSP and three lower parts. +/// +/// The MSP retains the decimal's validity while lower parts are non-nullable. Lower parts +/// are constructed with zeroes at null positions instead of invalid bytes. +/// Empty and all-null arrays use constant parts, preserving the part types and MSP's nullability. +/// +/// # Errors +/// +/// Returns an error if the array's validity cannot be derived or executed. +pub fn split_decimal(decimal: &DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { + let validity = decimal.validity()?; + let len = decimal.len(); + let mask = validity.execute_mask(len, ctx)?; + + if mask.all_false() || decimal.is_empty() { + return Ok(split_no_valid_row(decimal, &validity)); + } + + Ok(match decimal.values_type() { + DecimalType::I8 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I16 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I32 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I64 => DecimalParts::from_narrow(decimal.buffer::(), validity), + DecimalType::I128 => { + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i128_to_parts); + DecimalParts::from_wide(msp, lower, validity) + } + DecimalType::I256 => { + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i256_to_parts); + DecimalParts::from_wide(msp, lower, validity) + } + }) +} + +/// Splits decimals with no valid rows (all null or empty) into constant decimal parts with the +/// corresponding nullability. +fn split_no_valid_row(decimal: &DecimalArray, validity: &Validity) -> DecimalParts { + let (msp_ptype, lower_part_count) = match decimal.values_type() { + DecimalType::I8 => (PType::I8, 0), + DecimalType::I16 => (PType::I16, 0), + DecimalType::I32 => (PType::I32, 0), + DecimalType::I64 => (PType::I64, 0), + DecimalType::I128 => (PType::I64, MAX_I128_LOWER_PARTS), + DecimalType::I256 => (PType::I64, MAX_I256_LOWER_PARTS), + }; + // Empty masks are also all-false. The default scalar is null for nullable inputs + // and zero for non-nullable empty inputs, preserving the MSP's nullability. + let msp = Scalar::default_value(&DType::Primitive(msp_ptype, validity.nullability())); + let len = decimal.len(); + DecimalParts { + msp: ConstantArray::new(msp, len).into_array(), + lower_parts: vec![ConstantArray::new(0u64, len).into_array(); lower_part_count], + } +} + +/// Split wide integers into a signed MSP and `N` unsigned lower parts. +/// +/// `to_parts` returns the MSP and lower words in most-significant-first order. +/// It is specialized for each input type: `i128` has one lower word and `i256` +/// has three. Null rows get zeros in every output buffer. The caller handles empty +/// and all-null arrays before calling this function. +pub fn split_wide( + values: &[T], + validity: &Mask, + to_parts: impl Fn(T) -> (i64, [u64; N]), +) -> (Buffer, [Buffer; N]) { + let len = values.len(); + let mut msp = BufferMut::::with_capacity(len); + let mut lower = std::array::from_fn::<_, N, _>(|_| BufferMut::::with_capacity(len)); + + // Allocate without zeroing, then initialize every part of each row together. + let msp_out = &mut msp.spare_capacity_mut()[..len]; + let mut lower_out = lower + .each_mut() + .map(|part| &mut part.spare_capacity_mut()[..len]); + + match validity { + Mask::AllTrue(_) => { + for row in 0..len { + let (high, words) = to_parts(values[row]); + msp_out[row].write(high); + for (part, word) in lower_out.iter_mut().zip(words) { + part[row].write(word); + } + } + } + Mask::Values(validity) => { + // A shorter bitmap would leave output slots uninitialized before set_len. + assert_eq!( + validity.bit_buffer().len(), + len, + "values and validity must have the same length" + ); + for (chunk_index, ((chunk, bits), msp)) in values + .chunks(64) + .zip(validity.bit_buffer().chunks().iter_padded()) + .zip(msp_out.chunks_mut(64)) + .enumerate() + { + for (i, (&value, msp)) in chunk.iter().zip(msp).enumerate() { + let mask = 0u64.wrapping_sub((bits >> i) & 1); + let (high, words) = to_parts(value); + msp.write(high & mask.cast_signed()); + for (part, word) in lower_out.iter_mut().zip(words) { + part[chunk_index * 64 + i].write(word & mask); + } + } + } + } + Mask::AllFalse(_) => unreachable!("all-null arrays are handled by split_decimal"), + } + + // SAFETY: the input and all output slices have len elements. Both branches + // initialize every slot, including null rows and the final partial chunk. + // The bitmap length check prevents the masked iteration from ending early. + unsafe { + msp.set_len(len); + for part in &mut lower { + part.set_len(len); + } + } + (msp.freeze(), lower.map(BufferMut::freeze)) +} + +/// Extract the high signed word and low unsigned word of an `i128`. +#[inline] +pub const fn i128_to_parts(value: i128) -> (i64, [u64; MAX_I128_LOWER_PARTS]) { + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "each cast preserves a 64-bit window of the original two's complement bits" + )] + ((value >> LOWER_PART_BITS) as i64, [value as u64]) +} + +/// Extract the signed MSP and three unsigned lower words of an `i256`. +#[inline] +pub const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_I256_LOWER_PARTS]) { + let (low, high) = value.to_parts(); + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "each cast preserves a 64-bit window of the original two's complement bits" + )] + ( + (high >> LOWER_PART_BITS) as i64, + [high as u64, (low >> LOWER_PART_BITS) as u64, low as u64], + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::PType; + use vortex_array::dtype::i256; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::split_decimal; + use crate::decimal_byte_parts::LOWER_PART_DTYPE; + use crate::decimal_byte_parts::MAX_LOWER_PARTS; + + #[test] + fn test_split_i256_part_count_and_types() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), + DecimalDType::new(76, 0), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); + for part in &parts.lower_parts { + assert_eq!(part.dtype(), &LOWER_PART_DTYPE); + } + Ok(()) + } + + #[rstest] + fn test_split_i256_part_order( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] + validity: Validity, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + buffer![ + i256::from_parts((2u128 << 64) | 3, (1i128 << 64) | 4), + i256::ZERO, + i256::from_parts((6u128 << 64) | 7, (-2i128 << 64) | 5), + ], + DecimalDType::new(76, 0), + validity.clone(), + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_arrays_eq!( + PrimitiveArray::new(buffer![1i64, 0, -2], validity), + parts.msp, + &mut ctx + ); + assert_eq!(parts.lower_parts.len(), 3); + for (part, expected) in parts.lower_parts.into_iter().zip([ + buffer![4u64, 0, 5], + buffer![2u64, 0, 6], + buffer![3u64, 0, 7], + ]) { + assert_arrays_eq!( + PrimitiveArray::new(expected, Validity::NonNullable), + part, + &mut ctx + ); + } + Ok(()) + } +} 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..af270bad407 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared fixtures for decimal byte-parts tests. + +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; + +use super::DecimalByteParts; +use super::DecimalBytePartsArray; + +/// An `i128`-backed decimal array, encoded as byte parts with one lower part. +pub(crate) fn i128_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { + DecimalByteParts::encode( + &DecimalArray::new(Buffer::from(values), DecimalDType::new(38, 2), validity), + &mut array_session().create_execution_ctx(), + ) + .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 { + DecimalByteParts::encode( + &DecimalArray::new(Buffer::from(values), DecimalDType::new(76, 2), validity), + &mut array_session().create_execution_ctx(), + ) + .vortex_expect("valid decimal byte parts") +} diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 36a53c3a614..2557555eac8 100644 --- a/encodings/decimal-byte-parts/src/lib.rs +++ b/encodings/decimal-byte-parts/src/lib.rs @@ -22,7 +22,9 @@ use vortex_session::VortexSession; /// Initialize decimal-byte-parts encoding in the given session. pub fn initialize(session: &VortexSession) { - session.arrays().register(DecimalByteParts); + // One plugin owns both serialized formats: registering it reads either ID and writes the + // one that fits the array. Which of them a writer may emit is decided by its editions. + session.arrays().register(DecimalBytePartsPlugin); compute::kernel::initialize(session); session.aggregate_fns().register_aggregate_kernel( diff --git a/vortex-array/src/dtype/bigint/mod.rs b/vortex-array/src/dtype/bigint/mod.rs index 3ebf01425d6..b03bdef43c2 100644 --- a/vortex-array/src/dtype/bigint/mod.rs +++ b/vortex-array/src/dtype/bigint/mod.rs @@ -130,6 +130,20 @@ impl From for i256 { } } +impl From for i256 { + #[inline] + fn from(value: i64) -> Self { + Self::from_i128(value.into()) + } +} + +impl From for i256 { + #[inline] + fn from(value: u64) -> Self { + Self::from_i128(value.into()) + } +} + impl Display for i256 { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.0) @@ -259,6 +273,7 @@ impl Shr for i256 { impl Shl for i256 { type Output = Self; + #[inline] fn shl(self, rhs: usize) -> Self::Output { use num_traits::ToPrimitive; Self( diff --git a/vortex-bench/src/lib.rs b/vortex-bench/src/lib.rs index 569dfc74a4d..0a48627aad0 100644 --- a/vortex-bench/src/lib.rs +++ b/vortex-bench/src/lib.rs @@ -278,7 +278,7 @@ pub fn retain_edition_encodings( .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - builder.retain_allowed_encodings(&allowed) + builder.allow_serialized_ids(&allowed) } /// Verify that local data has already been prepared for the requested benchmark formats. diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder/mod.rs similarity index 67% rename from vortex-btrblocks/src/builder.rs rename to vortex-btrblocks/src/builder/mod.rs index 341569f408c..3699964e702 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -18,7 +18,7 @@ use crate::schemes::integer; use crate::schemes::string; use crate::schemes::temporal; -/// All available compression schemes. +/// The newest versions of all available compression schemes. /// /// This list is order-sensitive: the builder preserves this order when constructing /// the final scheme list, so that tie-breaking is deterministic. @@ -90,12 +90,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + allowed_serialized_ids: Option>, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + allowed_serialized_ids: None, } } } @@ -107,13 +109,14 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + allowed_serialized_ids: None, } } /// Adds an external compression scheme not in [`ALL_SCHEMES`]. /// /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes - /// with the compressor. + /// with the compressor. Register only the newest version of a scheme. /// /// # Panics /// @@ -158,7 +161,7 @@ impl BtrBlocksCompressorBuilder { /// Both the array-level and the buffer-level Zstd schemes are added. Buffer-level /// compression preserves binary arrays' buffer layout for zero-conversion GPU decompression, /// but belongs to the opt-in `zstd` edition, so callers filter the two through - /// [`retain_allowed_encodings`](Self::retain_allowed_encodings). + /// [`allow_serialized_ids`](Self::allow_serialized_ids). /// /// This preset is intended for files that will be decoded by CUDA kernels. It may choose a /// larger encoded representation than the default compressor. @@ -194,122 +197,61 @@ impl BtrBlocksCompressorBuilder { } /// Removes the specified compression schemes by their [`SchemeId`]. + /// + /// An ID anywhere in a registered predecessor chain removes the entire chain. + /// + /// # Panics + /// + /// Panics if a traversed predecessor chain contains a cycle. pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { let ids: HashSet<_> = ids.into_iter().collect(); - self.schemes.retain(|s| !ids.contains(&s.id())); + self.schemes.retain(|scheme| { + let mut seen = HashSet::new(); + let mut candidate = Some(*scheme); + while let Some(version) = candidate { + assert!( + seen.insert(version.id()), + "cycle in scheme predecessor chain" + ); + if ids.contains(&version.id()) { + return false; + } + candidate = version.predecessor(); + } + true + }); self } - /// Retains only schemes whose produced serialized IDs all belong to `allowed`. + /// Restricts compression to the serialized IDs in `allowed`, intersecting with any earlier + /// call. /// - /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions - /// permit. - pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); + /// At build time, each scheme is replaced by the newest version in its predecessor chain + /// whose [`produced_encodings`](Scheme::produced_encodings) are all permitted. + /// Schemes with no eligible version are removed. This also applies to schemes added after + /// this call. The file writer passes the serialized IDs its enabled editions permit. + pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { + let allowed: HashSet = match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(allowed).copied().collect(), + None => allowed.clone(), + }; + self.allowed_serialized_ids = Some(allowed); self } /// Builds the configured [`BtrBlocksCompressor`]. + /// + /// # Panics + /// + /// Panics if predecessor chains contain a cycle or share a scheme ID. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + let compressor = CascadingCompressor::new(self.schemes); + BtrBlocksCompressor(match self.allowed_serialized_ids { + Some(allowed) => compressor.with_allowed_serialized_ids(allowed), + None => compressor, + }) } } #[cfg(test)] -mod tests { - use vortex_array::VTable; - use vortex_fastlanes::FoR; - - use super::*; - - #[test] - fn empty_starts_with_no_schemes() { - let builder = BtrBlocksCompressorBuilder::empty(); - assert!(builder.schemes.is_empty()); - } - - #[test] - fn default_includes_all_schemes() { - let builder = BtrBlocksCompressorBuilder::default(); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn retain_allowed_encodings_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - - let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - #[test] - fn cuda_compatible_excludes_alprd() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - !builder - .schemes - .iter() - .any(|s| s.id() == float::ALPRDScheme.id()) - ); - } - - /// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. - #[test] - fn cuda_compatible_excludes_every_sparse_scheme() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - for excluded in [ - integer::SparseScheme.id(), - float::NullDominatedSparseScheme.id(), - string::NullDominatedSparseScheme.id(), - ] { - assert!( - !builder.schemes.iter().any(|s| s.id() == excluded), - "{excluded} should be excluded" - ); - } - } - - #[test] - fn cuda_compatible_uses_fsst_for_strings() { - let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); - assert!( - builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::FSSTScheme.id()) - ); - #[cfg(feature = "zstd")] - assert!( - !builder - .schemes - .iter() - .any(|scheme| scheme.id() == string::ZstdScheme.id()) - ); - } - - #[test] - #[cfg(feature = "pco")] - fn cuda_compatible_excludes_pco() { - let builder = BtrBlocksCompressorBuilder::default() - .with_new_scheme(&integer::PcoScheme) - .with_new_scheme(&float::PcoScheme) - .only_cuda_compatible(); - for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { - assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); - } - } -} +mod tests; diff --git a/vortex-btrblocks/src/builder/tests.rs b/vortex-btrblocks/src/builder/tests.rs new file mode 100644 index 00000000000..4310b9416c8 --- /dev/null +++ b/vortex-btrblocks/src/builder/tests.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::VarBin; +use vortex_compressor::scheme::CompressionEstimate; +use vortex_compressor::scheme::EstimateVerdict; +use vortex_error::VortexResult; +use vortex_fastlanes::FoR; +use vortex_fsst::FSST; +use vortex_session::registry::CachedId; + +use super::*; +use crate::ArrayAndStats; +use crate::CompressorContext; + +#[test] +fn empty_starts_with_no_schemes() { + assert!(BtrBlocksCompressorBuilder::empty().schemes.is_empty()); +} + +#[test] +fn default_includes_all_schemes() { + assert_eq!( + BtrBlocksCompressorBuilder::default().schemes.len(), + ALL_SCHEMES.len() + ); +} + +#[test] +fn allowed_serialized_ids_filter_schemes_at_build() { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .build(); + for scheme in ALL_SCHEMES { + assert_eq!( + compressor.has_scheme(scheme.id()), + scheme.id() == integer::FoRScheme.id() + ); + } +} + +#[test] +fn allowing_all_declared_outputs_keeps_every_scheme() { + let allowed = ALL_SCHEMES + .iter() + .flat_map(|s| s.produced_encodings()) + .collect(); + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&allowed) + .build(); + for scheme in ALL_SCHEMES { + assert!(compressor.has_scheme(scheme.id())); + } +} + +#[rstest] +#[case::neither(vec![], false)] +#[case::fsst_only(vec![FSST.id()], false)] +#[case::varbin_only(vec![VarBin.id()], false)] +#[case::both(vec![FSST.id(), VarBin.id()], true)] +fn all_required_outputs_must_be_allowed(#[case] allowed: Vec, #[case] expected: bool) { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&allowed.into_iter().collect()) + .build(); + assert_eq!(compressor.has_scheme(string::FSSTScheme.id()), expected); +} + +#[rstest] +#[case::forbidden(HashSet::new(), false)] +#[case::permitted(HashSet::from([FoR.id()]), true)] +fn restriction_applies_to_schemes_added_later( + #[case] allowed: HashSet, + #[case] expected: bool, +) { + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_serialized_ids(&allowed) + .with_new_scheme(&integer::FoRScheme) + .build(); + assert_eq!(compressor.has_scheme(integer::FoRScheme.id()), expected); +} + +#[test] +fn repeated_restrictions_intersect() { + let compressor = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id(), FSST.id()])) + .allow_serialized_ids(&HashSet::from([FSST.id(), VarBin.id()])) + .build(); + assert!(!compressor.has_scheme(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(string::FSSTScheme.id())); +} + +#[test] +fn cuda_compatible_excludes_alprd() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + !builder + .schemes + .iter() + .any(|s| s.id() == float::ALPRDScheme.id()) + ); +} + +/// `vortex.sparse` has no CUDA decode kernel, so no sparse scheme may survive this preset. +#[test] +fn cuda_compatible_excludes_every_sparse_scheme() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + for excluded in [ + integer::SparseScheme.id(), + float::NullDominatedSparseScheme.id(), + string::NullDominatedSparseScheme.id(), + ] { + assert!( + !builder.schemes.iter().any(|s| s.id() == excluded), + "{excluded} should be excluded" + ); + } +} + +#[test] +fn cuda_compatible_uses_fsst_for_strings() { + let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); + assert!( + builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::FSSTScheme.id()) + ); + #[cfg(feature = "zstd")] + assert!( + !builder + .schemes + .iter() + .any(|scheme| scheme.id() == string::ZstdScheme.id()) + ); +} + +#[test] +#[cfg(feature = "pco")] +fn cuda_compatible_excludes_pco() { + let builder = BtrBlocksCompressorBuilder::default() + .with_new_scheme(&integer::PcoScheme) + .with_new_scheme(&float::PcoScheme) + .only_cuda_compatible(); + for scheme in [integer::PcoScheme.id(), float::PcoScheme.id()] { + assert!(!builder.schemes.iter().any(|s| s.id() == scheme)); + } +} + +static FOR_V2_ID: CachedId = CachedId::new("test.for_v2"); + +#[derive(Debug)] +struct NewFoRScheme; + +impl Scheme for NewFoRScheme { + fn scheme_name(&self) -> &'static str { + "test.for_v2" + } + + fn matches(&self, canonical: &Canonical) -> bool { + integer::FoRScheme.matches(canonical) + } + + fn produced_encodings(&self) -> Vec { + vec![*FOR_V2_ID] + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + Some(&integer::FoRScheme) + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Skip) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +#[test] +fn restrictions_select_predecessors_of_schemes_added_later() { + let compressor = BtrBlocksCompressorBuilder::empty() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .with_new_scheme(&NewFoRScheme) + .build(); + assert!(compressor.has_scheme(integer::FoRScheme.id())); + assert!(compressor.has_scheme_family(NewFoRScheme.id())); + assert!(!compressor.has_scheme(NewFoRScheme.id())); +} + +#[rstest] +#[case::old(integer::FoRScheme.id())] +#[case::new(NewFoRScheme.id())] +fn excluding_any_version_removes_the_chain(#[case] excluded: SchemeId) { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&NewFoRScheme) + .exclude_schemes([excluded]) + .build(); + assert!(!compressor.has_scheme_family(integer::FoRScheme.id())); + assert!(!compressor.has_scheme_family(NewFoRScheme.id())); +} diff --git a/vortex-btrblocks/src/canonical_compressor.rs b/vortex-btrblocks/src/canonical_compressor.rs index d93be365550..b9759821c63 100644 --- a/vortex-btrblocks/src/canonical_compressor.rs +++ b/vortex-btrblocks/src/canonical_compressor.rs @@ -294,7 +294,7 @@ mod tests { // survives. let compressor = BtrBlocksCompressorBuilder::default() .only_cuda_compatible() - .retain_allowed_encodings(&HashSet::from([allowed])) + .allow_serialized_ids(&HashSet::from([allowed])) .build(); let mut ctx = SESSION.create_execution_ctx(); let compressed = compressor.compress(&array.clone().into_array(), &mut ctx)?; diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 1dff2171f60..f77a77d8c50 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -8,7 +8,6 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; 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; @@ -16,6 +15,7 @@ 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::decimal_byte_parts_v1_id; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -41,7 +41,9 @@ impl Scheme for DecimalScheme { } fn produced_encodings(&self) -> Vec { - vec![DecimalByteParts.id()] + // This scheme only builds single-part arrays, which serialize under the frozen v1 ID. + // The in-memory ID is the v2 wire ID, which no edition permits yet. + vec![decimal_byte_parts_v1_id()] } /// Children: primitive=0. diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 21e440a7436..fa6dca51bb4 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -250,7 +250,7 @@ fn trace_scan_compare_on_compressed_quantity() -> VortexResult<()> { optimize root=vortex.binary(bool, len=4096) session=false reduce_parent static:DictionaryScalarFnValuesPushDownRule slot=0 parent=vortex.binary(bool, len=4096) child=vortex.dict(i16, len=4096) -> vortex.dict(bool, len=4096) done output=vortex.dict(bool, len=4096) - child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.dict(bool, len=4096) + child_execute_parent session[0]:execute_parent_fn slot=0 parent=vortex.binary(bool, len=4096) child=vortex.decimal_byte_parts.v2(decimal(15,2), len=4096) -> vortex.dict(bool, len=4096) iter 1 current=vortex.dict(bool, len=4096) builder_active=false ExecuteSlot slot=0 parent=vortex.dict(bool, len=4096) child=fastlanes.bitpacked(u8, len=4096) iter 2 current=fastlanes.bitpacked(u8, len=4096) stack_parent=vortex.dict(bool, len=4096) slot=0 builder_active=false @@ -403,8 +403,8 @@ 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) - done output=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.v2(decimal(15,2), len=4096) -> vortex.decimal_byte_parts.v2(decimal(15,2), len=43) + done output=vortex.decimal_byte_parts.v2(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 reduce_parent static:FoRFilterPushDownRule slot=0 parent=vortex.filter(i32, len=43) child=fastlanes.for(i32, len=4096) -> fastlanes.for(i32, len=43) @@ -439,6 +439,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.v2(decimal(15,2), len=4096) -> vortex.decimal_byte_parts.v2(decimal(15,2), len=64) + done output=vortex.decimal_byte_parts.v2(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-btrblocks/tests/golden.rs b/vortex-btrblocks/tests/golden.rs index fc636252c27..b3671131a4d 100644 --- a/vortex-btrblocks/tests/golden.rs +++ b/vortex-btrblocks/tests/golden.rs @@ -423,7 +423,7 @@ fn compressor_for_session( .into_iter() .collect(); without_onpair(builder) - .retain_allowed_encodings(&allowed) + .allow_serialized_ids(&allowed) .build() } @@ -436,7 +436,7 @@ fn compressor_with_onpair( .enabled_component_ids(ComponentKind::Array) .into_iter() .collect(); - builder.retain_allowed_encodings(&allowed).build() + builder.allow_serialized_ids(&allowed).build() } #[test] diff --git a/vortex-btrblocks/tests/scheme_versions.rs b/vortex-btrblocks/tests/scheme_versions.rs new file mode 100644 index 00000000000..9c68725ae0a --- /dev/null +++ b/vortex-btrblocks/tests/scheme_versions.rs @@ -0,0 +1,101 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayId; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VTable; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::assert_arrays_eq; + use vortex_btrblocks::ArrayAndStats; + use vortex_btrblocks::CascadingCompressor; + use vortex_btrblocks::CompressorContext; + use vortex_btrblocks::Scheme; + use vortex_btrblocks::SchemeExt; + use vortex_btrblocks::schemes::integer::DeltaScheme; + use vortex_btrblocks::schemes::integer::IntRLEScheme; + use vortex_compressor::scheme::CompressionEstimate; + use vortex_compressor::scheme::EstimateVerdict; + use vortex_error::VortexResult; + use vortex_fastlanes::Delta; + use vortex_fastlanes::RLE; + use vortex_session::registry::CachedId; + + static DELTA_V2_ID: CachedId = CachedId::new("test.delta_v2"); + static DELTA_V1: DeltaScheme = DeltaScheme::new(1.25); + + #[derive(Debug)] + struct DeltaV2; + + impl Scheme for DeltaV2 { + fn scheme_name(&self) -> &'static str { + "test.delta_v2" + } + + fn matches(&self, canonical: &Canonical) -> bool { + DELTA_V1.matches(canonical) + } + + fn produced_encodings(&self) -> Vec { + vec![*DELTA_V2_ID] + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + Some(&DELTA_V1) + } + + fn num_children(&self) -> usize { + 2 + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Skip) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } + } + + #[rstest] + #[case::predecessor(Delta.id(), true)] + #[case::replacement(*DELTA_V2_ID, false)] + fn rle_respects_selected_delta_version( + #[case] allowed_delta: ArrayId, + #[case] expect_delta: bool, + ) -> VortexResult<()> { + let session = array_session(); + vortex_fastlanes::initialize(&session); + let compressor = CascadingCompressor::new(vec![&IntRLEScheme, &DeltaV2]) + .with_allowed_serialized_ids([RLE.id(), allowed_delta].into_iter().collect()); + assert!(compressor.has_scheme_family(DELTA_V1.id())); + let array = PrimitiveArray::from_iter((0..65_536u32).map(|i| (i / 64) % 100)).into_array(); + let mut ctx = session.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut ctx)?; + assert_eq!(compressed.encoding_id(), RLE.id()); + let has_delta = compressed + .depth_first_traversal() + .any(|array| array.encoding_id() == Delta.id()); + assert_eq!(has_delta, expect_delta); + assert_arrays_eq!(compressed, array, &mut ctx); + Ok(()) + } +} diff --git a/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap index 780ef30a6b0..fad1b4d2b4e 100644 --- a/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__compact__decimal_prices.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: decimal(12,2), len=16384, nbytes=131072 -root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=47666 +root: vortex.decimal_byte_parts.v2(decimal(12,2), len=16384) nbytes=47666 metadata: msp: vortex.pco(i32, len=16384) nbytes=47666 metadata: ptype: i32, nrows: 16384, slice: 0..16384 diff --git a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap index f669755e4b1..6eb4ccfed8d 100644 --- a/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap +++ b/vortex-btrblocks/tests/snapshots/golden__regular__decimal_prices.snap @@ -3,7 +3,7 @@ source: vortex-btrblocks/tests/golden.rs expression: rendered --- input: decimal(12,2), len=16384, nbytes=131072 -root: vortex.decimal_byte_parts(decimal(12,2), len=16384) nbytes=49152 +root: vortex.decimal_byte_parts.v2(decimal(12,2), len=16384) nbytes=49152 metadata: msp: fastlanes.bitpacked(i32, len=16384) nbytes=49152 metadata: bit_width: 24, offset: 0 diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..ecfd3c2c542 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -59,7 +59,7 @@ impl CascadingCompressor { let canonical = array.clone().execute::(exec_ctx)?.0; let compact = canonical.compact(exec_ctx)?; - let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + let compressed = self.compress_canonical(compact, self.root_context(), exec_ctx)?; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); @@ -93,7 +93,7 @@ impl CascadingCompressor { let child_ctx = parent_ctx .clone() - .descend_with_scheme(parent_id, child_index); + .descend_with_scheme(self.resolve_scheme_id(parent_id), child_index); self.compress_canonical(compact, child_ctx, exec_ctx) } diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..3f88debbc24 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,8 +9,13 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_map::HashMap; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; +use crate::scheme::CompressorContext; use crate::scheme::DescendantExclusion; use crate::scheme::Scheme; use crate::scheme::SchemeExt; @@ -46,13 +51,39 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// Maps every registered version to the version selected for compression. + scheme_aliases: HashMap, + + /// Configuration only: retained so repeated restrictions intersect exactly. + allowed_serialized_ids: Option>, } impl CascadingCompressor { /// Creates a new compressor with the given schemes. /// + /// Register only the newest version of each scheme. Predecessor IDs are aliases for the + /// selected version in exclusions and [`has_scheme_family`](Self::has_scheme_family) checks. /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. + /// + /// # Panics + /// + /// Panics if predecessor chains contain a cycle or share a scheme ID, including when multiple + /// versions of the same scheme are registered separately. pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { + let mut scheme_aliases = HashMap::new(); + for &scheme in &schemes { + let mut candidate = Some(scheme); + while let Some(version) = candidate { + assert!( + scheme_aliases.insert(version.id(), scheme.id()).is_none(), + "scheme {} appears more than once in the registered predecessor chains", + version.id(), + ); + candidate = version.predecessor(); + } + } + // Root exclusion: exclude IntDict from list/listview offsets (monotonically // increasing data where dictionary encoding is wasteful). let root_exclusions = vec![DescendantExclusion { @@ -63,14 +94,76 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + scheme_aliases, + allowed_serialized_ids: None, } } - /// Returns whether the compressor was configured with `scheme`. + /// Selects the newest eligible version of each scheme, intersecting with any earlier call. + /// + /// A version is eligible only when all of its [`Scheme::produced_encodings`] are allowed. + /// Otherwise its predecessors are tried in order; the scheme is removed if none is eligible. + /// Selection preserves registration order and happens before any compression or estimation. + pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { + let allowed = match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(&allowed).copied().collect(), + None => allowed, + }; + let mut replacements = HashMap::new(); + self.schemes = self + .schemes + .into_iter() + .filter_map(|scheme| { + let mut candidate = Some(scheme); + while let Some(version) = candidate { + if version + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)) + { + replacements.insert(scheme.id(), version.id()); + return Some(version); + } + candidate = version.predecessor(); + } + None + }) + .collect(); + self.scheme_aliases.retain(|_, selected| { + if let Some(replacement) = replacements.get(selected) { + *selected = *replacement; + true + } else { + false + } + }); + self.allowed_serialized_ids = Some(allowed); + self + } + + /// The context a compress call starts from. + pub(crate) fn root_context(&self) -> CompressorContext { + CompressorContext::new() + } + + /// Returns whether any version in the scheme's family is enabled. + /// + /// A family is a registered scheme and its predecessor chain. `scheme` may name any version + /// in that chain. Use [`Self::has_scheme`] to check the exact selected version. + pub fn has_scheme_family(&self, scheme: SchemeId) -> bool { + self.scheme_aliases.contains_key(&scheme) + } + + /// Returns whether this exact scheme version is selected for compression. + /// + /// Use this before invoking a specific implementation directly. pub fn has_scheme(&self, scheme: SchemeId) -> bool { - self.schemes - .iter() - .any(|candidate| candidate.id() == scheme) + self.scheme_aliases.get(&scheme) == Some(&scheme) + } + + /// Resolves a registered version to the selected version, leaving unknown IDs unchanged. + fn resolve_scheme_id(&self, scheme: SchemeId) -> SchemeId { + self.scheme_aliases.get(&scheme).copied().unwrap_or(scheme) } } @@ -78,3 +171,6 @@ impl CascadingCompressor { #[cfg(test)] mod tests; + +#[cfg(test)] +mod version_tests; diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..c492729f77c 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -152,10 +152,9 @@ impl CascadingCompressor { // The root entry is always first in the history (if present). Check if the root has // excluded us. if let Some((_, child_idx)) = iter.next_if(|&(sid, _)| sid == ROOT_SCHEME_ID) - && self - .root_exclusions - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + && self.root_exclusions.iter().any(|rule| { + self.resolve_scheme_id(rule.excluded) == id && rule.children.contains(child_idx) + }) { return true; } @@ -163,10 +162,9 @@ impl CascadingCompressor { // Push rules: Check if any of our ancestors have excluded us. for (ancestor_id, child_idx) in iter { if let Some(ancestor) = self.schemes.iter().find(|s| s.id() == ancestor_id) - && ancestor - .descendant_exclusions() - .iter() - .any(|rule| rule.excluded == id && rule.children.contains(child_idx)) + && ancestor.descendant_exclusions().iter().any(|rule| { + self.resolve_scheme_id(rule.excluded) == id && rule.children.contains(child_idx) + }) { return true; } @@ -174,10 +172,9 @@ impl CascadingCompressor { // Pull rules: Check if we have excluded ourselves because of our ancestors. for rule in candidate.ancestor_exclusions() { - if history - .iter() - .any(|(sid, cidx)| *sid == rule.ancestor && rule.children.contains(*cidx)) - { + if history.iter().any(|(sid, cidx)| { + *sid == self.resolve_scheme_id(rule.ancestor) && rule.children.contains(*cidx) + }) { return true; } } diff --git a/vortex-compressor/src/compressor/version_tests.rs b/vortex-compressor/src/compressor/version_tests.rs new file mode 100644 index 00000000000..1fc6dfbac99 --- /dev/null +++ b/vortex-compressor/src/compressor/version_tests.rs @@ -0,0 +1,272 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_error::VortexResult; +use vortex_session::registry::CachedId; + +use super::*; +use crate::scheme::AncestorExclusion; +use crate::scheme::CompressionEstimate; +use crate::scheme::EstimateVerdict; +use crate::stats::ArrayAndStats; +use crate::stats::GenerateStatsOptions; + +static V1_ID: CachedId = CachedId::new("test.version_1"); +static V2_ID: CachedId = CachedId::new("test.version_2"); +static V3_ID: CachedId = CachedId::new("test.version_3"); +static AUX_ID: CachedId = CachedId::new("test.auxiliary"); + +#[derive(Debug)] +struct TestScheme { + name: &'static str, + version: u8, + predecessor: Option<&'static dyn Scheme>, + push: Option<&'static dyn Scheme>, + pull: Option<&'static dyn Scheme>, +} + +impl TestScheme { + const fn new( + name: &'static str, + version: u8, + predecessor: Option<&'static dyn Scheme>, + ) -> Self { + Self { + name, + version, + predecessor, + push: None, + pull: None, + } + } +} + +impl Scheme for TestScheme { + fn scheme_name(&self) -> &'static str { + self.name + } + + fn matches(&self, canonical: &Canonical) -> bool { + canonical.dtype().is_int() + } + + fn produced_encodings(&self) -> Vec { + match self.version { + 1 => vec![*V1_ID], + 2 => vec![*V2_ID, *AUX_ID], + 3 => vec![*V3_ID], + _ => vec![], + } + } + + fn predecessor(&self) -> Option<&'static dyn Scheme> { + self.predecessor + } + + fn num_children(&self) -> usize { + 2 + } + + fn descendant_exclusions(&self) -> Vec { + self.push + .map(|scheme| DescendantExclusion { + excluded: scheme.id(), + children: ChildSelection::One(1), + }) + .into_iter() + .collect() + } + + fn ancestor_exclusions(&self) -> Vec { + self.pull + .map(|scheme| AncestorExclusion { + ancestor: scheme.id(), + children: ChildSelection::One(1), + }) + .into_iter() + .collect() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // Older versions would beat newer versions if they reached estimation together. + CompressionEstimate::Verdict(EstimateVerdict::Ratio(5.0 - f64::from(self.version))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(data.array().clone()) + } +} + +static V1: TestScheme = TestScheme::new("test.scheme_v1", 1, None); +static V2: TestScheme = TestScheme::new("test.scheme_v2", 2, Some(&V1)); +static V3: TestScheme = TestScheme::new("test.scheme_v3", 3, Some(&V2)); +static OTHER: TestScheme = TestScheme::new("test.other", 0, None); + +#[test] +fn newest_eligible_version_is_selected_before_estimation() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut exec_ctx = session.create_execution_ctx(); + let data = ArrayAndStats::new( + PrimitiveArray::from_iter(0..128i32).into_array(), + GenerateStatsOptions::default(), + ); + for (allowed, expected) in [ + (None, V3.id()), + ( + Some(HashSet::from([*V1_ID, *V2_ID, *AUX_ID, *V3_ID])), + V3.id(), + ), + (Some(HashSet::from([*V1_ID, *V2_ID, *AUX_ID])), V2.id()), + (Some(HashSet::from([*V2_ID, *AUX_ID])), V2.id()), + (Some(HashSet::from([*V1_ID, *V2_ID])), V1.id()), + (Some(HashSet::from([*V1_ID])), V1.id()), + ] { + let mut compressor = CascadingCompressor::new(vec![&V3]); + if let Some(allowed) = allowed { + compressor = compressor.with_allowed_serialized_ids(allowed); + } + assert_eq!(compressor.schemes.len(), 1); + let winner = compressor.choose_best_scheme( + &compressor.schemes, + &data, + compressor.root_context(), + &mut exec_ctx, + )?; + assert_eq!(winner.map(|(scheme, _)| scheme.id()), Some(expected)); + for version in [&V1, &V2, &V3] { + assert!(compressor.has_scheme_family(version.id())); + assert_eq!( + compressor.has_scheme(version.id()), + version.id() == expected + ); + } + assert!(!compressor.has_scheme(OTHER.id())); + } + Ok(()) +} + +#[test] +fn no_eligible_version_removes_the_entire_chain() { + for allowed in [HashSet::new(), HashSet::from([*V2_ID])] { + let compressor = CascadingCompressor::new(vec![&V3]).with_allowed_serialized_ids(allowed); + assert!(compressor.schemes.is_empty()); + for version in [&V1, &V2, &V3] { + assert!(!compressor.has_scheme_family(version.id())); + assert!(!compressor.has_scheme(version.id())); + } + } +} + +#[test] +fn fallback_preserves_registration_order() { + let compressor = CascadingCompressor::new(vec![&V3, &OTHER]) + .with_allowed_serialized_ids(HashSet::from([*V1_ID])); + assert_eq!( + compressor + .schemes + .iter() + .map(|s| s.id()) + .collect::>(), + vec![V1.id(), OTHER.id()] + ); +} + +#[test] +fn successive_restrictions_keep_aliases_and_intersect_wire_ids() { + let compressor = CascadingCompressor::new(vec![&V3]) + .with_allowed_serialized_ids(HashSet::from([*V1_ID, *V2_ID, *AUX_ID])) + .with_allowed_serialized_ids(HashSet::from([*V1_ID])); + assert_eq!(compressor.schemes[0].id(), V1.id()); + assert_eq!(compressor.resolve_scheme_id(V3.id()), V1.id()); + + let compressor = compressor.with_allowed_serialized_ids(HashSet::from([*V2_ID, *AUX_ID])); + assert!(compressor.schemes.is_empty()); + assert!(!compressor.has_scheme_family(V3.id())); +} + +static PUSH_OLD: TestScheme = TestScheme { + push: Some(&V1), + ..TestScheme::new("test.push_old", 0, None) +}; +static PUSH_NEW: TestScheme = TestScheme { + push: Some(&V3), + ..TestScheme::new("test.push_new", 0, None) +}; +static PULL_OLD: TestScheme = TestScheme { + pull: Some(&V1), + ..TestScheme::new("test.pull_old", 0, None) +}; +static PULL_NEW: TestScheme = TestScheme { + pull: Some(&V3), + ..TestScheme::new("test.pull_new", 0, None) +}; + +#[test] +fn exclusions_follow_upgrades_and_fallbacks() { + for allowed in [HashSet::from([*V1_ID]), HashSet::from([*V3_ID])] { + let compressor = + CascadingCompressor::new(vec![&V3, &PUSH_OLD, &PUSH_NEW, &PULL_OLD, &PULL_NEW]) + .with_allowed_serialized_ids(allowed); + let selected = compressor.schemes[0]; + for child in [0, 1] { + for pusher in [&PUSH_OLD, &PUSH_NEW] { + let ctx = compressor + .root_context() + .descend_with_scheme(pusher.id(), child); + assert_eq!(compressor.is_excluded(selected, &ctx), child == 1); + } + let ctx = compressor + .root_context() + .descend_with_scheme(selected.id(), child); + for puller in [&PULL_OLD, &PULL_NEW] { + assert_eq!(compressor.is_excluded(puller, &ctx), child == 1); + } + assert!(compressor.is_excluded(selected, &ctx)); + } + } +} + +#[test] +fn root_exclusions_follow_new_versions() { + static DICT_V2: TestScheme = TestScheme::new("test.dict_v2", 3, Some(&IntDictScheme)); + let compressor = CascadingCompressor::new(vec![&DICT_V2]); + let ctx = compressor + .root_context() + .descend_with_scheme(ROOT_SCHEME_ID, structural::root_list_children::OFFSETS); + assert!(compressor.is_excluded(&DICT_V2, &ctx)); + let ctx = compressor + .root_context() + .descend_with_scheme(ROOT_SCHEME_ID, structural::root_list_children::SIZES); + assert!(!compressor.is_excluded(&DICT_V2, &ctx)); +} + +#[test] +#[should_panic(expected = "appears more than once")] +fn predecessor_cycles_are_rejected() { + static CYCLE: TestScheme = TestScheme::new("test.cycle", 1, Some(&CYCLE)); + CascadingCompressor::new(vec![&CYCLE]); +} + +#[test] +#[should_panic(expected = "appears more than once")] +fn registering_multiple_versions_is_rejected() { + CascadingCompressor::new(vec![&V3, &V1]); +} diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..0b9d8e3d4b8 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -41,7 +41,7 @@ pub struct CompressorContext { } impl CompressorContext { - /// Creates a new `CompressorContext`. + /// Creates a new root `CompressorContext`. /// /// This should **only** be created by the compressor. pub(crate) fn new() -> Self { diff --git a/vortex-compressor/src/scheme/exclusion.rs b/vortex-compressor/src/scheme/exclusion.rs index 2dba6b85046..46ca12d7735 100644 --- a/vortex-compressor/src/scheme/exclusion.rs +++ b/vortex-compressor/src/scheme/exclusion.rs @@ -34,7 +34,8 @@ impl ChildSelection { /// `ZigZag` excludes `Dict` from all its children. #[derive(Debug, Clone, Copy)] pub struct DescendantExclusion { - /// The scheme to exclude from descendants. + /// The scheme to exclude from descendants. Any version in its registered predecessor chain + /// refers to the selected version. pub excluded: SchemeId, /// Which children of the declaring scheme this rule applies to. pub children: ChildSelection, @@ -47,7 +48,8 @@ pub struct DescendantExclusion { /// `Sequence` excludes itself when `IntDict` is an ancestor on its codes child. #[derive(Debug, Clone, Copy)] pub struct AncestorExclusion { - /// The ancestor scheme that makes the declaring scheme ineligible. + /// The ancestor scheme that makes the declaring scheme ineligible. Any version in its + /// registered predecessor chain refers to the selected version. pub ancestor: SchemeId, /// Which children of the ancestor this rule applies to. pub children: ChildSelection, diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..f00cde3abd9 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,17 +124,31 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// The serialized IDs this scheme itself may write into its compressed output. + /// The serialized IDs this scheme may write its output under. Every ID must be permitted before + /// this scheme can be selected. /// - /// Every declared ID must be permitted for the scheme to be used. Cascaded children are - /// compressed by other schemes, which declare their own IDs, so only arrays constructed - /// directly by [`compress`](Scheme::compress) belong here. Canonical arrays the scheme - /// merely rearranges do not need to be declared. + /// Cascaded children are compressed by other schemes, which declare their own IDs, so only + /// arrays constructed directly by [`compress`](Scheme::compress) belong here. Canonical + /// arrays the scheme merely rearranges do not need to be declared. /// - /// For most encodings this is the in-memory encoding ID. An encoding with several wire - /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. + /// Alternative versions belong in the [`predecessor`](Scheme::predecessor) chain, rather than + /// in this list. Once selected, a scheme must produce output compatible with these IDs without + /// consulting the writer's configuration. fn produced_encodings(&self) -> Vec; + /// The preceding version of this scheme, used when this version's serialized IDs are unavailable. + /// + /// Register only the newest version. The compressor selects the first eligible version in + /// this chain during configuration, before matching, generating statistics, or estimating. + /// A predecessor is a compatibility fallback, not an alternative compression candidate. + /// + /// Versions must have distinct scheme IDs and form an acyclic chain. They must support the + /// same input types and preserve child indices, because exclusions and scheme dependencies + /// referring to any version in the registered chain apply to the selected version. + fn predecessor(&self) -> Option<&'static dyn Scheme> { + None + } + /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies /// every scheme. diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 44fc7d87938..915d41048b9 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -154,7 +154,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file_block_rows( .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() .only_cuda_compatible() - .retain_allowed_encodings(&allowed_encodings), + .allow_serialized_ids(&allowed_encodings), ) .with_flat_strategy(Arc::new(CudaFlatLayoutStrategy::default())); if block_rows > 0 { diff --git a/vortex-cuda/gpu-scan-cli/src/main.rs b/vortex-cuda/gpu-scan-cli/src/main.rs index 2aa32729ef3..515f14c211e 100644 --- a/vortex-cuda/gpu-scan-cli/src/main.rs +++ b/vortex-cuda/gpu-scan-cli/src/main.rs @@ -103,7 +103,7 @@ fn cuda_write_strategy(session: &VortexSession) -> Arc VortexResult<()> { Ok(()) } +#[rstest] +#[case::sparse(PrimitiveArray::from_iter( + (0..4096i32).map(|i| if i % 100 == 0 { i + 1 } else { 0 }), +).into_array())] +#[case::fsst(VarBinViewArray::from_iter( + (0..4096).map(|i| Some(format!("this_is_a_common_prefix_with_some_variation_{i}_and_a_common_suffix_pattern"))), + DType::Utf8(Nullability::NonNullable), +).into_array())] +#[tokio::test] +async fn test_writer_excludes_schemes_with_unavailable_outputs( + #[case] array: ArrayRef, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::() + .with::(); + // Permit Constant and VarBin, but not the subsequently registered Sparse and FSST. + crate::enable_all_registered_array_encodings(&session); + crate::register_default_encodings(&session); + let mut buf = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buf, array.clone().to_array_stream()) + .await?; + let read = session + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + assert_arrays_eq!(read, array, &mut session.create_execution_ctx()); + Ok(()) +} + #[tokio::test] async fn test_writer_empty_chunks() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 6c0d01dfa05..e6bb9be4ff8 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -254,7 +254,7 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_serialized_ids), + .allow_serialized_ids(&allowed_serialized_ids), ) .build(), }; diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 7288fc5d1e3..db032cd79ce 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -389,7 +389,7 @@ impl PyVortexWriteOptions { compressor = compressor.with_compact(); } let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); + .with_btrblocks_builder(compressor.allow_serialized_ids(&allowed_encodings)); let strategy = strategy.build(); current_runtime().block_on(async move { match resolve_store(path, store.map(|x| x.into_inner()))? { 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..166d84f5ac6 --- /dev/null +++ b/vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `DecimalByteParts` fixture for wide decimal values that need lower parts. + +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, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let parts = split_decimal(decimal, ctx)?; + 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, ctx)?; + + // 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, ctx)?; + + // 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, ctx)?; + + 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..4d799e33e74 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,7 @@ mod bytebool; mod constant; mod datetimeparts; mod decimal_byte_parts; +mod decimal_byte_parts_v2; mod delta; mod dict; mod for_; @@ -38,6 +39,7 @@ pub fn fixtures() -> Vec> { Box::new(bytebool::ByteBoolFixture), Box::new(datetimeparts::DateTimePartsFixture), Box::new(decimal_byte_parts::DecimalBytePartsFixture), + Box::new(decimal_byte_parts_v2::DecimalBytePartsV2Fixture), // Re-enable this once delta is stable // Box::new(delta::DeltaFixture), Box::new(dict::DictFixture), diff --git a/vortex-tui/src/convert.rs b/vortex-tui/src/convert.rs index ab316982b27..aba7ccc6b35 100644 --- a/vortex-tui/src/convert.rs +++ b/vortex-tui/src/convert.rs @@ -107,7 +107,7 @@ pub async fn exec_convert(session: &VortexSession, flags: ConvertArgs) -> anyhow compressor = compressor.with_compact(); } let strategy = WriteStrategyBuilder::default() - .with_btrblocks_builder(compressor.retain_allowed_encodings(&allowed_encodings)); + .with_btrblocks_builder(compressor.allow_serialized_ids(&allowed_encodings)); let mut file = File::create(output_path).await?; session