From ea5f87899fea94ebc50e290f67626f290b2b88e8 Mon Sep 17 00:00:00 2001 From: Matthew Katz <87445739+mhk197@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:15:20 -0400 Subject: [PATCH 01/12] Add wide decimal byte-part splitting and assembly logic (#9808) Signed-off-by: "Matt Katz" Signed-off-by: Matt Katz --- Cargo.lock | 2 + encodings/decimal-byte-parts/Cargo.toml | 10 + .../decimal-byte-parts/benches/common/mod.rs | 51 +++ .../benches/dbp_assemble.rs | 48 +++ .../decimal-byte-parts/benches/dbp_split.rs | 47 +++ .../src/decimal_byte_parts/limbs/mod.rs | 324 ++++++++++++++++++ .../src/decimal_byte_parts/limbs/tests.rs | 226 ++++++++++++ .../src/decimal_byte_parts/mod.rs | 9 + vortex-array/src/dtype/bigint/mod.rs | 1 + 9 files changed, 718 insertions(+) create mode 100644 encodings/decimal-byte-parts/benches/common/mod.rs create mode 100644 encodings/decimal-byte-parts/benches/dbp_assemble.rs create mode 100644 encodings/decimal-byte-parts/benches/dbp_split.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 275db7ef8ee..1343e6dda00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10983,8 +10983,10 @@ dependencies = [ name = "vortex-decimal-byte-parts" version = "0.1.0" dependencies = [ + "codspeed-divan-compat", "num-traits", "prost 0.14.4", + "rand 0.10.2", "rstest", "vortex-array", "vortex-buffer", diff --git a/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 4934ec4fa27..9f2e387a4da 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -26,5 +26,15 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } vortex-array = { path = "../../vortex-array", features = ["_test-harness"] } + +[[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..3e0e60730e4 --- /dev/null +++ b/encodings/decimal-byte-parts/benches/common/mod.rs @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared decimal inputs for splitting and assembly benchmarks. + +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::dtype::i256; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::vortex_panic; + +pub(super) 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(super) fn decimal_array( + values_type: DecimalType, + len: usize, + validity: Validity, +) -> DecimalArray { + let mut rng = StdRng::seed_from_u64(42); + + macro_rules! decimal { + ($T:ty, $precision:literal) => {{ + let max = <$T>::pow(10, $precision) - 1; + let values: Buffer<$T> = (0..len).map(|_| rng.random_range(-max..=max)).collect(); + DecimalArray::new(values, DecimalDType::new($precision, 2), validity) + }}; + } + + match values_type { + DecimalType::I64 => decimal!(i64, 18), + DecimalType::I128 => decimal!(i128, 38), + DecimalType::I256 => { + // Keep the magnitude below 10^76 while exercising all four signed/unsigned words. + let values: Buffer = (0..len) + .map(|_| i256::from_parts(rng.random(), rng.random::() >> 4)) + .collect(); + DecimalArray::new(values, 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..327899fa640 --- /dev/null +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Reassembling primitive decimal parts across storage widths and lengths. + +mod common; + +use divan::Bencher; +use divan::black_box; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DecimalType; +use vortex_array::validity::Validity; +use vortex_decimal_byte_parts::_benchmarking::assemble_decimal; +use vortex_decimal_byte_parts::split_decimal; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::common::cases; +use crate::common::decimal_array; + +fn main() { + divan::main(); +} + +#[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"); + let msp = parts + .msp + .execute::(&mut ctx) + .vortex_expect("execute benchmark MSP"); + let lower_parts = parts + .lower_parts + .into_iter() + .map(|part| part.execute::(&mut ctx)) + .collect::>>() + .vortex_expect("execute benchmark lower parts"); + let decimal_dtype = decimal.decimal_dtype(); + + bencher.bench(|| { + assemble_decimal(black_box(&msp), black_box(&lower_parts), decimal_dtype) + .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..ba716d1ba5b --- /dev/null +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -0,0 +1,47 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splitting decimal arrays across storage widths, lengths, and validity paths. + +mod common; + +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::cases; +use crate::common::decimal_array; + +fn main() { + divan::main(); +} + +#[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_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/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs new file mode 100644 index 00000000000..1e561b149fe --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs @@ -0,0 +1,324 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Splitting decimal values into 64-bit parts and reassembling them. +//! +//! A `DecimalByteParts` array stores each value as a signed most significant part (MSP) +//! followed by `k` unsigned 64-bit lower parts ordered most significant first. The encoded +//! value is +//! +//! ```text +//! msp * 2^(64k) + Σ_{i, +} + +impl DecimalParts { + /// Construct decimal parts from an MSP with no lower parts. + fn from_msp(values: Buffer, validity: Validity) -> Self { + Self { + msp: PrimitiveArray::new(values, validity).into_array(), + lower_parts: Vec::new(), + } + } + + fn new( + msp: Buffer, + lower_parts: impl IntoIterator>, + validity: Validity, + ) -> Self { + Self { + msp: PrimitiveArray::new(msp, validity).into_array(), + lower_parts: lower_parts + .into_iter() + .map(|part| PrimitiveArray::new(part, Validity::NonNullable).into_array()) + .collect(), + } + } +} + +/// 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. +/// +/// # 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()?; + Ok(match decimal.values_type() { + DecimalType::I8 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I16 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I32 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I64 => DecimalParts::from_msp(decimal.buffer::(), validity), + DecimalType::I128 => { + let mask = validity.execute_mask(decimal.len(), ctx)?; + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i128_to_parts); + DecimalParts::new(msp, lower, validity) + } + DecimalType::I256 => { + let mask = validity.execute_mask(decimal.len(), ctx)?; + let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i256_to_parts); + DecimalParts::new(msp, lower, validity) + } + }) +} + +/// 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. +fn split_wide( + values: &Buffer, + 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)); + + // Zero out all parts if all null + if validity.all_false() { + msp.push_n(0, len); + for part in &mut lower { + part.push_n(0, len); + } + return (msp.freeze(), lower.map(BufferMut::freeze)); + } + + // Allocate without zeroing, then initialize every part of each row together. + let msp_out = &mut msp.spare_capacity_mut()[..len]; + let mut lower_out = lower + .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!("AllFalse case addressed above"), + } + + // 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] +const fn i128_to_parts(value: i128) -> (i64, [u64; 1]) { + #[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] +const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_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], + ) +} + +/// Reassemble primitive arrays that constitute decimal byte parts into a canonical decimal array. +/// +/// The MSP must be signed. There must be between zero and three (inclusive) `u64` lower parts, ordered +/// most significant first. The lower parts must be non-nullable. Every input array must have the same length. +/// +/// With no lower parts, the MSP buffer is reused as the decimal values. One lower part +/// assembles into `i128`. Two or three lower parts assemble into `i256`. +/// +/// # Errors +/// +/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity +/// cannot be derived. +pub fn assemble_decimal( + msp: &PrimitiveArray, + lower_parts: &[PrimitiveArray], + decimal_dtype: DecimalDType, +) -> VortexResult { + let validity = msp.validity()?; + vortex_ensure!(msp.dtype().as_ptype().is_signed_int()); + + if lower_parts.is_empty() { + return Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { + // SAFETY: the buffer is typed by the array's own ptype, the decimal dtype is the + // array's, and the validity is taken from the same array. + unsafe { DecimalArray::new_unchecked(msp.to_buffer::

(), decimal_dtype, validity) } + })); + } + + let len = msp.len(); + let lower: Vec<&[u64]> = lower_parts + .iter() + .map(|part| { + vortex_ensure!( + part.dtype() == &LOWER_PART_DTYPE, + "lower part must be non-nullable u64" + ); + let part = part.as_slice::(); + vortex_ensure!( + part.len() == len, + "lower part has len {}, expected {len}", + part.len() + ); + Ok(part) + }) + .collect::>()?; + + Ok(match lower.as_slice() { + [first] => DecimalArray::new( + assemble_wide::(msp, [first]), + decimal_dtype, + validity, + ), + [first, second] => DecimalArray::new( + assemble_wide::(msp, [first, second]), + decimal_dtype, + validity, + ), + [first, second, third] => DecimalArray::new( + assemble_wide::(msp, [first, second, third]), + decimal_dtype, + validity, + ), + _ => vortex_bail!( + "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", + lower.len() + ), + }) +} + +/// Reassemble a signed MSP and `K` unsigned lower parts into wide integers. +/// +/// Each row starts with the MSP sign-extended to `T`. Appending a lower word shifts the +/// accumulated value left by 64 bits and fills the low bits with that word. Lower parts +/// are appended most significant first. +/// +/// The callers select `i128` for one lower part and `i256` for two or three. Since `K` +/// is constant, the compiler can unroll the loop that appends the lower words. +fn assemble_wide(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer +where + T: NativeDecimalType + Shl + BitOr, +{ + let mut out = BufferMut::::with_capacity(msp.len()); + match_each_signed_integer_ptype!(msp.ptype(), |P| { + out.extend_trusted(msp.as_slice::

().iter().enumerate().map(|(row, value)| { + #[allow( + clippy::useless_conversion, + reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" + )] + let mut value = T::from(i64::from(*value)).vortex_expect("MSP fits in the output type"); + for part in lower { + value = (value << LOWER_PART_BITS) + | T::from(part[row]).vortex_expect("lower word fits in the output type"); + } + value + })); + }); + out.freeze() +} + +#[cfg(test)] +mod tests; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs new file mode 100644 index 00000000000..3e3de06c44e --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::DecimalArray; +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 super::*; + +#[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 expected = PrimitiveArray::new( + decimal + .validity()? + .execute_mask(len, &mut ctx)? + .iter() + .map(|valid| if valid { u64::MAX } else { 0 }) + .collect::>(), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + for lower in parts.lower_parts { + assert_arrays_eq!(expected.clone(), lower, &mut ctx); + } + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + Ok(()) +} + +fn round_trip(decimal: DecimalArray) -> VortexResult { + let mut ctx = array_session().create_execution_ctx(); + let parts = split_decimal(&decimal, &mut ctx)?; + let msp = parts.msp.execute::(&mut ctx)?; + let lower = parts + .lower_parts + .into_iter() + .map(|part| part.execute::(&mut ctx)) + .collect::>>()?; + assemble_decimal(&msp, &lower, decimal.decimal_dtype()) +} + +#[rstest] +#[case::zero(0)] +#[case::one(1)] +#[case::minus_one(-1)] +#[case::limb_boundary(1i128 << 64)] +#[case::just_below_limb_boundary((1i128 << 64) - 1)] +#[case::negative_limb_boundary(-(1i128 << 64))] +#[case::max(i128::MAX)] +#[case::min(i128::MIN)] +fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) +} + +#[rstest] +#[case::zero(i256::ZERO)] +#[case::one(i256::ONE)] +#[case::minus_one(i256::ZERO - i256::ONE)] +#[case::max(i256::MAX)] +#[case::min(i256::MIN)] +#[case::word_1(i256::from_parts(1u128 << 64, 0))] +#[case::word_2(i256::from_parts(0, 1))] +#[case::word_3(i256::from_parts(0, 1i128 << 64))] +#[case::mixed(i256::from_parts(u128::MAX, -3))] +fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { + let decimal = DecimalArray::new( + Buffer::from(vec![value]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ); + let round_tripped = round_trip(decimal)?; + assert_eq!(round_tripped.buffer::().as_slice(), &[value]); + Ok(()) +} + +#[rstest] +fn test_split_narrow_decimal_has_no_lower_parts( + #[values(Validity::NonNullable, Validity::AllInvalid, Validity::from_iter([true, false, true]))] + validity: Validity, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); + let parts = split_decimal(&decimal, &mut ctx)?; + assert!(parts.lower_parts.is_empty()); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); + let msp = parts.msp.execute::(&mut ctx)?; + assert_eq!( + msp.as_slice::().as_ptr(), + decimal.buffer::().as_ptr() + ); + assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); + Ok(()) +} + +#[test] +fn test_split_i256_part_count_and_types() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), + DecimalDType::new(76, 0), + Validity::NonNullable, + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); + assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); + for part in &parts.lower_parts { + assert_eq!(part.dtype(), &LOWER_PART_DTYPE); + } + Ok(()) +} + +#[rstest] +fn test_split_i256_part_order( + #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] validity: Validity, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal = DecimalArray::new( + buffer![ + i256::from_parts((2u128 << 64) | 3, (1i128 << 64) | 4), + i256::ZERO, + i256::from_parts((6u128 << 64) | 7, (-2i128 << 64) | 5), + ], + DecimalDType::new(76, 0), + validity.clone(), + ); + let parts = split_decimal(&decimal, &mut ctx)?; + assert_arrays_eq!( + PrimitiveArray::new(buffer![1i64, 0, -2], validity), + parts.msp, + &mut ctx + ); + assert_eq!(parts.lower_parts.len(), 3); + for (part, expected) in parts.lower_parts.into_iter().zip([ + buffer![4u64, 0, 5], + buffer![2u64, 0, 6], + buffer![3u64, 0, 7], + ]) { + assert_arrays_eq!( + PrimitiveArray::new(expected, Validity::NonNullable), + part, + &mut ctx + ); + } + Ok(()) +} + +#[rstest] +fn test_assemble_rejects_mismatched_lower_lengths( + #[values(1, 2, 3)] lower_count: usize, + #[values(0, 1, 3)] lower_len: usize, +) { + let msp = PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable); + let mut lower = vec![PrimitiveArray::new(buffer![0u64; 2], Validity::NonNullable); lower_count]; + lower[lower_count - 1] = PrimitiveArray::new(buffer![0u64; lower_len], Validity::NonNullable); + let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); + assert!(assemble_decimal(&msp, &lower, dtype).is_err()); +} + +#[rstest] +fn test_assemble_i256_part_order_and_sign_extension( + #[values(false, true)] narrow_msp: bool, + #[values(2, 3)] lower_count: usize, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![3i8, -3], Validity::NonNullable) + } else { + PrimitiveArray::new(buffer![3i64, -3], Validity::NonNullable) + }; + let lower = + [4u64, 1, 2].map(|word| PrimitiveArray::new(buffer![word; 2], Validity::NonNullable)); + let dtype = DecimalDType::new(76, 0); + let actual = assemble_decimal(&msp, &lower[3 - lower_count..], dtype)?; + let low = (1u128 << 64) | 2; + let expected = if lower_count == 2 { + buffer![i256::from_parts(low, 3), i256::from_parts(low, -3)] + } else { + buffer![ + i256::from_parts(low, (3i128 << 64) | 4), + i256::from_parts(low, (-3i128 << 64) | 4), + ] + }; + assert_arrays_eq!( + DecimalArray::new(expected, dtype, Validity::NonNullable), + actual, + &mut ctx + ); + Ok(()) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index e1d081ddaca..951a893a785 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -9,9 +9,18 @@ use vortex_array::Array; use vortex_array::ArrayParts; use vortex_array::ArrayView; pub(crate) mod compute; +mod limbs; +pub use limbs::DecimalParts; +pub use limbs::MAX_LOWER_PARTS; +pub use limbs::split_decimal; mod rules; mod slice; +#[doc(hidden)] +pub mod _benchmarking { + pub use super::limbs::assemble_decimal; +} + use prost::Message as _; use vortex_array::ArrayEq; use vortex_array::ArrayHash; diff --git a/vortex-array/src/dtype/bigint/mod.rs b/vortex-array/src/dtype/bigint/mod.rs index 3ebf01425d6..47195526b1f 100644 --- a/vortex-array/src/dtype/bigint/mod.rs +++ b/vortex-array/src/dtype/bigint/mod.rs @@ -259,6 +259,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( From 8321696cbf7e6be525f8c495646227459388f9e2 Mon Sep 17 00:00:00 2001 From: Matthew Katz <87445739+mhk197@users.noreply.github.com> Date: Tue, 15 Sep 2026 09:05:12 -0400 Subject: [PATCH 02/12] Support wide decimals in `DecimalBytePartsArray` and kernels (#9809) `DecimalBytePartsArray` previously stored the entire unscaled decimal value in one signed integer child, limiting it to values that fit in 64 bits. It now supports wide decimals by representing each value as integer parts that can be compressed independently, while preserving the decimal's logical precision, scale, and nullability. The array has a signed most significant part (MSP) and up to three unsigned 64-bit lower parts, ordered most significant first. Splitting canonical decimal storage produces: | Decimal storage | Children | | --- | --- | | `i8` / `i16` / `i32` / `i64` | Signed MSP only; shares the original value buffer | | `i128` | `i64` MSP + one `u64` lower part | | `i256` | `i64` MSP + three `u64` lower parts | Only the MSP carries validity. Every lower part must be a non-nullable `u64` array with the same length, and splitting wide decimals zeroes the parts at null positions. All children remain `ArrayRef`s, so their individual encodings are independent of the decimal representation. `execute::` reassembles a `DecimalArray` from the MSP and lower parts children. `take` with nullable indices is not yet supported by the DecimalByteParts kernel for arrays with lower parts; it falls back to canonical execution. Taking each part directly would make the lower parts nullable, violating the representation's invariant. The frozen serializer also continues to reject arrays with lower parts. This PR also refines the splitting and assembly modules. * Assembly takes `ArrayRef`s instead of `PrimitiveArray`s so that in the future, we can add special fast paths for constant arrays * Assembly casts lower parts to `u64`s, allowing for assembly of narrowed lower parts. * Assembly loop is optimized such that it vectorizes for `i256` assembly on local runs. --------- Signed-off-by: Matt Katz --- .gitignore | 2 + Cargo.lock | 85 +- Cargo.toml | 1 + encodings/decimal-byte-parts/Cargo.toml | 2 + .../decimal-byte-parts/benches/common/mod.rs | 80 +- .../benches/dbp_assemble.rs | 181 +++- .../decimal-byte-parts/benches/dbp_split.rs | 105 +- .../src/decimal_byte_parts/array.rs | 922 ++++++++++++++++++ .../src/decimal_byte_parts/assemble.rs | 464 +++++++++ .../src/decimal_byte_parts/compute/cast.rs | 13 +- .../src/decimal_byte_parts/compute/compare.rs | 50 + .../src/decimal_byte_parts/compute/filter.rs | 46 +- .../decimal_byte_parts/compute/is_constant.rs | 27 +- .../src/decimal_byte_parts/compute/kernel.rs | 8 - .../src/decimal_byte_parts/compute/mask.rs | 15 +- .../src/decimal_byte_parts/compute/mod.rs | 34 + .../decimal_byte_parts/{ => compute}/slice.rs | 16 +- .../src/decimal_byte_parts/compute/take.rs | 109 ++- .../src/decimal_byte_parts/limbs/mod.rs | 324 ------ .../src/decimal_byte_parts/limbs/tests.rs | 226 ----- .../src/decimal_byte_parts/mod.rs | 414 +------- .../src/decimal_byte_parts/prop_tests.rs | 183 ++++ .../src/decimal_byte_parts/rules.rs | 42 +- .../src/decimal_byte_parts/split.rs | 310 ++++++ .../src/decimal_byte_parts/testing.rs | 47 + vortex-array/src/dtype/bigint/mod.rs | 14 + vortex-btrblocks/src/trace_tests.rs | 5 +- .../kernel/encodings/decimal_byte_parts.rs | 7 + 28 files changed, 2626 insertions(+), 1106 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs rename encodings/decimal-byte-parts/src/decimal_byte_parts/{ => compute}/slice.rs (53%) delete mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs delete mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs diff --git a/.gitignore b/.gitignore index f9613807332..6db14ce5f6a 100644 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,8 @@ coverage.xml *.cover *.py,cover .hypothesis/ +# hegeltest's example database, the Rust equivalent of .hypothesis/ +.hegel/ .pytest_cache/ cover/ diff --git a/Cargo.lock b/Cargo.lock index 1343e6dda00..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" @@ -10984,11 +11063,13 @@ 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/encodings/decimal-byte-parts/Cargo.toml b/encodings/decimal-byte-parts/Cargo.toml index 9f2e387a4da..e5d15c2c4f7 100644 --- a/encodings/decimal-byte-parts/Cargo.toml +++ b/encodings/decimal-byte-parts/Cargo.toml @@ -27,9 +27,11 @@ 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" diff --git a/encodings/decimal-byte-parts/benches/common/mod.rs b/encodings/decimal-byte-parts/benches/common/mod.rs index 3e0e60730e4..252431a58a3 100644 --- a/encodings/decimal-byte-parts/benches/common/mod.rs +++ b/encodings/decimal-byte-parts/benches/common/mod.rs @@ -1,51 +1,77 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Shared decimal inputs for splitting and assembly benchmarks. +//! Shared inputs for splitting and assembly benchmarks. 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::dtype::i256; -use vortex_array::validity::Validity; use vortex_buffer::Buffer; -use vortex_error::vortex_panic; pub(super) fn cases() -> Vec<(DecimalType, usize)> { - [DecimalType::I64, DecimalType::I128, DecimalType::I256] + [DecimalType::I128, DecimalType::I256] .into_iter() .flat_map(|values_type| [1_024, 8_192].map(|len| (values_type, len))) .collect() } -pub(super) fn decimal_array( - values_type: DecimalType, - len: usize, - validity: Validity, -) -> DecimalArray { +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; - macro_rules! decimal { - ($T:ty, $precision:literal) => {{ - let max = <$T>::pow(10, $precision) - 1; - let values: Buffer<$T> = (0..len).map(|_| rng.random_range(-max..=max)).collect(); - DecimalArray::new(values, DecimalDType::new($precision, 2), validity) - }}; + 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() } - match values_type { - DecimalType::I64 => decimal!(i64, 18), - DecimalType::I128 => decimal!(i128, 38), - DecimalType::I256 => { - // Keep the magnitude below 10^76 while exercising all four signed/unsigned words. - let values: Buffer = (0..len) - .map(|_| i256::from_parts(rng.random(), rng.random::() >> 4)) - .collect(); - DecimalArray::new(values, DecimalDType::new(76, 2), validity) + 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}"), } - _ => 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 index 327899fa640..48161a2c664 100644 --- a/encodings/decimal-byte-parts/benches/dbp_assemble.rs +++ b/encodings/decimal-byte-parts/benches/dbp_assemble.rs @@ -1,48 +1,169 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Reassembling primitive decimal parts across storage widths and lengths. +//! 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::VortexSessionExecute; -use vortex_array::array_session; -use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DecimalType; -use vortex_array::validity::Validity; -use vortex_decimal_byte_parts::_benchmarking::assemble_decimal; -use vortex_decimal_byte_parts::split_decimal; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; +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::decimal_array; +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(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"); - let msp = parts - .msp - .execute::(&mut ctx) - .vortex_expect("execute benchmark MSP"); - let lower_parts = parts - .lower_parts - .into_iter() - .map(|part| part.execute::(&mut ctx)) - .collect::>>() - .vortex_expect("execute benchmark lower parts"); - let decimal_dtype = decimal.decimal_dtype(); - - bencher.bench(|| { - assemble_decimal(black_box(&msp), black_box(&lower_parts), decimal_dtype) - .vortex_expect("assemble decimal byte parts") - }); +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 index ba716d1ba5b..dd5b9abcf13 100644 --- a/encodings/decimal-byte-parts/benches/dbp_split.rs +++ b/encodings/decimal-byte-parts/benches/dbp_split.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Splitting decimal arrays across storage widths, lengths, and validity paths. +//! Splitting integer slices, including output-buffer allocation. +//! Array benchmarks also include validity execution and array construction. mod common; @@ -10,38 +11,104 @@ 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 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::decimal_array; +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_all_valid(bencher: Bencher, (values_type, len): (DecimalType, usize)) { - bench_split(bencher, values_type, len, Validity::AllValid); +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_mixed_null(bencher: Bencher, (values_type, len): (DecimalType, usize)) { +fn dbp_split_kernel_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))); + 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: 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") - }); +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..bcdb3579f53 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -0,0 +1,922 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal byte-parts array types, validation, and VTable implementations. + +use std::fmt::Display; +use std::fmt::Formatter; +use std::hash::Hasher; + +use prost::Message as _; +use vortex_array::Array; +use vortex_array::ArrayEq; +use vortex_array::ArrayHash; +use vortex_array::ArrayId; +use vortex_array::ArrayParts; +use vortex_array::ArrayRef; +use vortex_array::ArraySlots; +use vortex_array::ArrayView; +use vortex_array::EqMode; +use vortex_array::ExecutionCtx; +use vortex_array::ExecutionResult; +use vortex_array::TypedArrayRef; +use vortex_array::array_slots; +use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::PType; +use vortex_array::scalar::DecimalValue; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::serde::ArrayChildren; +use vortex_array::vtable::OperationsVTable; +use vortex_array::vtable::VTable; +use vortex_array::vtable::ValidityChild; +use vortex_array::vtable::ValidityVTableFromChild; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::LOWER_PART_DTYPE; +use super::MAX_LOWER_PARTS; +use super::assemble::assemble_decimal; +use super::assemble::assemble_wide_decimal_value; +use super::rules::PARENT_RULES; + +/// A [`DecimalByteParts`]-encoded Vortex array. +pub type DecimalBytePartsArray = Array; + +#[derive(Clone, prost::Message)] +pub struct DecimalBytesPartsMetadata { + #[prost(enumeration = "PType", tag = "1")] + zeroth_child_ptype: i32, + #[prost(uint32, tag = "2")] + lower_part_count: u32, +} + +impl DecimalBytesPartsMetadata { + fn from_array(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { + Ok(Self { + zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, + lower_part_count: u32::try_from(array.lower_parts().len()) + .map_err(|_| vortex_err!("lower part count exceeds u32"))?, + }) + } + + fn into_array_parts( + self, + dtype: &DType, + len: usize, + children: &dyn ArrayChildren, + ) -> VortexResult> { + vortex_ensure!( + dtype.as_decimal_opt().is_some(), + "decoding decimal but given non decimal dtype {dtype}" + ); + + let encoded_dtype = DType::Primitive(self.zeroth_child_ptype(), dtype.nullability()); + + let lower_part_count = self.lower_part_count()?; + vortex_ensure!( + children.len() == DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, + "expected {} children, got {}", + DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, + children.len() + ); + + let msp = children.get(DecimalBytePartsSlots::MSP, &encoded_dtype, len)?; + + let mut slots = ArraySlots::with_capacity(children.len()); + slots.push(Some(msp)); + for idx in 0..lower_part_count { + slots.push(Some(children.get( + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx, + &LOWER_PART_DTYPE, + len, + )?)); + } + + Ok( + ArrayParts::new(DecimalByteParts, dtype.clone(), len, DecimalBytePartsData) + .with_slots(slots), + ) + } + + /// The number of lower parts encoded in this array. + /// + /// # Errors + /// + /// Returns an error if the count exceeds [`MAX_LOWER_PARTS`]. + fn lower_part_count(&self) -> VortexResult { + let count = usize::try_from(self.lower_part_count) + .map_err(|_| vortex_err!("lower part count {} out of range", self.lower_part_count))?; + vortex_ensure!( + count <= MAX_LOWER_PARTS, + "at most {MAX_LOWER_PARTS} lower parts are supported, got {count}" + ); + Ok(count) + } +} + +/// This array encodes decimals by splitting them between 1-4 columns of primitive typed children. +/// +/// The most significant part (MSP) stores the most significant decimal bits. It is signed and is +/// nullable iff the decimal is nullable. +/// +/// Every lower part is a non-nullable 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 + /// canonical decimal array. + /// + /// # Errors + /// + /// Returns an error if the parts do not describe a valid decimal, see + /// [`DecimalBytePartsData::validate`]. + pub fn try_new_with_lower_parts( + msp: ArrayRef, + lower_parts: Vec, + decimal_dtype: DecimalDType, + ) -> VortexResult { + let len = msp.len(); + let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); + let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); + Array::try_from_parts( + ArrayParts::new(DecimalByteParts, dtype, len, DecimalBytePartsData).with_slots(slots), + ) + } + + /// Construct a [`DecimalBytePartsArray`] from parts whose invariants are already established. + /// + /// # Safety + /// + /// The MSP must have a signed integer dtype (`i8`, `i16`, `i32`, or `i64`). There must be + /// at most [`MAX_LOWER_PARTS`] lower parts, each a non-nullable 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 { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); + *ID + } + + fn validate( + &self, + _data: &Self::TypedArrayData, + dtype: &DType, + len: usize, + slots: &[Option], + ) -> VortexResult<()> { + let Some(decimal_dtype) = dtype.as_decimal_opt() else { + vortex_bail!("expected decimal dtype, got {}", dtype) + }; + + let min_slots = DecimalBytePartsSlots::FIXED_COUNT; + let max_slots = min_slots + MAX_LOWER_PARTS; + vortex_ensure!( + (min_slots..=max_slots).contains(&slots.len()), + "expected {min_slots}..={max_slots} slots, got {}", + slots.len() + ); + for (idx, slot) in slots.iter().enumerate() { + vortex_ensure!(slot.is_some(), "missing required slot {idx}"); + } + + let slots = DecimalBytePartsSlotsView::from_slots(slots); + DecimalBytePartsData::validate( + slots.msp, + slots.lower_parts.iter(), + *decimal_dtype, + dtype, + len, + ) + } + + fn nbuffers(_array: ArrayView<'_, Self>) -> usize { + 0 + } + + fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { + vortex_panic!("DecimalBytePartsArray buffer index {idx} out of bounds") + } + + fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { + vortex_panic!("DecimalBytePartsArray buffer_name index {idx} out of bounds") + } + + fn with_buffers( + &self, + array: ArrayView<'_, Self>, + buffers: &[BufferHandle], + ) -> VortexResult> { + vortex_array::vtable::with_empty_buffers(self, array, buffers) + } + + fn serialize( + array: ArrayView<'_, Self>, + _session: &VortexSession, + ) -> VortexResult>> { + vortex_ensure!( + array.lower_parts().is_empty(), + "serializing DecimalByteParts with lower parts is not supported" + ); + Ok(Some( + DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), + )) + } + + fn deserialize( + &self, + dtype: &DType, + len: usize, + metadata: &[u8], + _buffers: &[BufferHandle], + children: &dyn ArrayChildren, + _session: &VortexSession, + ) -> VortexResult> { + let metadata = DecimalBytesPartsMetadata::decode(metadata)?; + vortex_ensure!( + metadata.lower_part_count()? == 0, + "vortex.decimal_byte_parts must not carry lower parts" + ); + metadata.into_array_parts(dtype, len, children) + } + + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + DecimalBytePartsSlots::slot_name(idx) + } + + fn reduce_parent( + array: ArrayView<'_, Self>, + parent: &ArrayRef, + child_idx: usize, + ) -> VortexResult> { + PARENT_RULES.evaluate(array, parent, child_idx) + } + + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + 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_array::vtable::VTable; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::DecimalByteParts; + use super::DecimalBytePartsArray; + use super::DecimalBytePartsArraySlotsExt; + use super::DecimalBytePartsData; + use crate::decimal_byte_parts::LOWER_PART_DTYPE; + use crate::decimal_byte_parts::MAX_LOWER_PARTS; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_of; + use crate::decimal_byte_parts::testing::i256_parts; + + #[test] + fn test_scalar_at_decimal_parts() { + let decimal_dtype = DecimalDType::new(8, 2); + let dtype = DType::Decimal(decimal_dtype, Nullability::Nullable); + let array = DecimalByteParts::try_new( + PrimitiveArray::new( + buffer![100i32, 200i32, 400i32], + Validity::Array(BoolArray::from_iter(vec![false, true, true]).into_array()), + ) + .into_array(), + decimal_dtype, + ) + .unwrap() + .into_array(); + + assert_eq!( + Scalar::null(dtype.clone()), + array + .execute_scalar(0, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new( + dtype.clone(), + Some(ScalarValue::Decimal(DecimalValue::I64(200))) + ) + .unwrap(), + array + .execute_scalar(1, &mut array_session().create_execution_ctx()) + .unwrap() + ); + assert_eq!( + Scalar::try_new(dtype, Some(ScalarValue::Decimal(DecimalValue::I64(400)))).unwrap(), + array + .execute_scalar(2, &mut array_session().create_execution_ctx()) + .unwrap() + ); + } + + /// The largest unscaled value a `Decimal(38, _)` can hold: `10^38 - 1`. + const MAX_PRECISION_38: i128 = 99_999_999_999_999_999_999_999_999_999_999_999_999; + + /// The largest unscaled value a `Decimal(76, _)` can hold: `10^76 - 1`. + fn max_precision_76() -> i256 { + i256::from_i128(10).wrapping_pow(76) - i256::ONE + } + + /// Values that exercise every 64-bit window of an `i128`, both signs, and the boundaries + /// where a lower part carries into the MSP. + fn wide_i128_values() -> Vec { + vec![ + 0, + 1, + -1, + (1 << 64) - 1, + 1 << 64, + -(1 << 64), + -((1 << 64) + 1), + MAX_PRECISION_38, + -MAX_PRECISION_38, + 1 << 100, + ] + } + + /// Values that exercise every 64-bit window of an `i256`. + fn wide_i256_values() -> Vec { + vec![ + i256::ZERO, + i256::ONE, + i256::ZERO - i256::ONE, + i256_of(0, u128::MAX), + i256_of(1, 0), + i256_of(-1, 0), + i256_of(-1, u128::MAX - 1), + i256_of(1 << 64, 12345), + max_precision_76(), + i256::ZERO - max_precision_76(), + ] + } + + #[rstest] + #[case::i128_non_nullable(i128_parts(wide_i128_values(), Validity::NonNullable))] + #[case::i256_non_nullable(i256_parts(wide_i256_values(), Validity::NonNullable))] + fn test_canonical_decimal_round_trips( + #[case] array: DecimalBytePartsArray, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)?; + assert_arrays_eq!(array, canonical, &mut ctx); + Ok(()) + } + + #[test] + fn test_lower_part_layout_i128() -> VortexResult<()> { + let array = i128_parts(vec![(3i128 << 64) | 7], Validity::NonNullable); + assert_eq!(array.lower_parts().len(), 1); + assert_eq!(array.msp().dtype().as_ptype(), PType::I64); + assert_eq!(array.lower_parts()[0].dtype(), &LOWER_PART_DTYPE); + + let mut ctx = array_session().create_execution_ctx(); + let msp = array.msp().clone().execute::(&mut ctx)?; + let lower = array.lower_parts()[0] + .clone() + .execute::(&mut ctx)?; + assert_eq!(msp.as_slice::(), &[3]); + assert_eq!(lower.as_slice::(), &[7]); + Ok(()) + } + + #[test] + fn test_lower_part_layout_i256() -> VortexResult<()> { + let array = i256_parts( + vec![i256_of((5i128 << 64) | 6, (7u128 << 64) | 8)], + Validity::NonNullable, + ); + assert_eq!(array.lower_parts().len(), MAX_LOWER_PARTS); + + let mut ctx = array_session().create_execution_ctx(); + let msp = array.msp().clone().execute::(&mut ctx)?; + assert_eq!(msp.as_slice::(), &[5]); + for (part, expected) in array.lower_parts().iter().zip([6u64, 7, 8]) { + let part = part.clone().execute::(&mut ctx)?; + assert_eq!(part.as_slice::(), &[expected]); + } + Ok(()) + } + + #[rstest] + #[case::i128(i128_parts(wide_i128_values(), Validity::AllValid))] + #[case::i256(i256_parts(wide_i256_values(), Validity::AllValid))] + fn test_scalar_at_matches_canonical(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)? + .into_array(); + let array = array.into_array(); + for idx in 0..array.len() { + assert_eq!( + array.execute_scalar(idx, &mut ctx)?, + canonical.execute_scalar(idx, &mut ctx)?, + "scalar mismatch at index {idx}" + ); + } + Ok(()) + } + + #[rstest] + fn test_scalar_at_matches_canonical_for_each_part_count( + #[values(false, true)] narrow_msp: bool, + #[values(0, 1, 2, 3)] lower_count: usize, + ) -> VortexResult<()> { + let validity = Validity::from_iter([false, true, true]); + let msp = if narrow_msp { + PrimitiveArray::new(buffer![0i8, 3, -3], validity) + } else { + PrimitiveArray::new(buffer![0i64, 3, -3], validity) + }; + let lower = [4u64, 1, 2] + .into_iter() + .take(lower_count) + .map(|word| PrimitiveArray::new(buffer![word; 3], Validity::NonNullable).into_array()) + .collect(); + let dtype = DecimalDType::new(if lower_count <= 1 { 38 } else { 76 }, 0); + let array = DecimalByteParts::try_new_with_lower_parts(msp.into_array(), lower, dtype)?; + let mut ctx = array_session().create_execution_ctx(); + let canonical = array + .clone() + .into_array() + .execute::(&mut ctx)?; + for row in 0..array.len() { + assert_eq!( + array.execute_scalar(row, &mut ctx)?, + canonical.execute_scalar(row, &mut ctx)? + ); + } + Ok(()) + } + + #[test] + fn test_scalar_at_null_with_lower_parts() -> VortexResult<()> { + let array = i128_parts( + vec![1i128 << 100, 2, 3], + Validity::Array(BoolArray::from_iter([false, true, true]).into_array()), + ) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + assert_eq!( + array.execute_scalar(1, &mut ctx)?, + Scalar::decimal( + DecimalValue::I128(2), + DecimalDType::new(38, 2), + Nullability::Nullable + ) + ); + Ok(()) + } + + #[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_of(1 << 100, 0)], Validity::NonNullable); + let canonical = i256_array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + + // A narrow MSP with a single lower part still fits 128 bits. + let array = DecimalByteParts::try_new_with_lower_parts( + buffer![1i8, -1, 0].into_array(), + vec![buffer![7u64, 7, 7].into_array()], + DecimalDType::new(38, 2), + )?; + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + assert_eq!( + canonical.buffer::().as_slice(), + &[(1i128 << 64) | 7, (-1i128 << 64) | 7, 7] + ); + + // Two lower parts under a narrow MSP overflow 128 bits, so the value widens. + let array = DecimalByteParts::try_new_with_lower_parts( + buffer![1i8].into_array(), + vec![buffer![0u64].into_array(), buffer![9u64].into_array()], + DecimalDType::new(76, 2), + )?; + let canonical = array.into_array().execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + assert_eq!(canonical.buffer::().as_slice(), &[i256_of(1, 9)]); + Ok(()) + } + + #[test] + fn test_unused_buffer_of_values_is_ignored_for_null_rows() -> VortexResult<()> { + // Null rows may hold arbitrary bits in the lower parts; they must stay null. + let array = DecimalByteParts::try_new_with_lower_parts( + PrimitiveArray::new( + buffer![0i64, 0, 0], + Validity::Array(BoolArray::from_iter([false, false, true]).into_array()), + ) + .into_array(), + vec![buffer![7u64, 9, 11].into_array()], + DecimalDType::new(38, 2), + )? + .into_array(); + + let mut ctx = array_session().create_execution_ctx(); + assert_eq!( + array.execute_scalar(0, &mut ctx)?, + Scalar::null(array.dtype().clone()) + ); + let canonical = array.clone().execute::(&mut ctx)?; + assert_arrays_eq!(array, canonical.into_array(), &mut ctx); + Ok(()) + } + #[test] + fn test_frozen_serializer_rejects_lower_parts() -> VortexResult<()> { + let session = array_session(); + let array = i128_parts(vec![1i128 << 70], Validity::NonNullable); + assert!(VTable::serialize(array.as_view(), &session).is_err()); + Ok(()) + } +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/assemble.rs new file mode 100644 index 00000000000..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..49c4021dd18 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs @@ -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::validity::Validity; use vortex_buffer::buffer; use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_of; + use crate::decimal_byte_parts::testing::i256_parts; #[test] fn test_filter_decimal_byte_parts() { @@ -59,4 +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_of(1, 0), + i256_of(-1, 5), + i256_of(0, u128::MAX), + i256_of(1 << 64, 7), + i256_of(0, 0), + ], + Validity::from_iter([true, false, true, true, false]), + ); + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs index 065bc5e0051..3fe59111f6e 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/is_constant.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::ArrayRef; +use vortex_array::ArrayView; use vortex_array::ExecutionCtx; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::fns::is_constant::IsConstant; @@ -15,7 +16,9 @@ use crate::decimal_byte_parts::DecimalBytePartsArraySlotsExt; /// DecimalByteParts-specific is_constant kernel. /// -/// Delegates to checking if the MSP (most significant part) is constant. +/// Delegates to checking that every part is constant: the MSP (most significant part) plus +/// each lower part. An all-null array is constant regardless of the bits its lower parts +/// hold in null slots. #[derive(Debug)] pub(crate) struct DecimalBytePartsIsConstantKernel; @@ -34,7 +37,27 @@ impl DynAggregateKernel for DecimalBytePartsIsConstantKernel { return Ok(None); }; - let result = is_constant(array.msp(), ctx)?; + let result = is_constant_parts(array, ctx)?; Ok(Some(IsConstant::make_partial(batch, result, ctx)?)) } } + +fn is_constant_parts( + array: ArrayView<'_, DecimalByteParts>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !is_constant(array.msp(), ctx)? { + return Ok(false); + } + // Null slots hold undefined bits in the lower parts, so they cannot make a constant + // (all-null) array non-constant. + if array.array().all_invalid(ctx)? { + return Ok(true); + } + for part in array.lower_parts().iter() { + if !is_constant(part, ctx)? { + return Ok(false); + } + } + Ok(true) +} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs index 5e8d28e3526..cb71ba7880c 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/kernel.rs @@ -1,9 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_array::ArrayVTable; -use vortex_array::arrays::Dict; -use vortex_array::arrays::dict::TakeExecuteAdaptor; use vortex_array::optimizer::kernels::ArrayKernelsExt; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; @@ -19,9 +16,4 @@ pub(crate) fn initialize(session: &VortexSession) { DecimalByteParts, CompareExecuteAdaptor(DecimalByteParts), ); - kernels.register_execute_parent_kernel( - Dict.id(), - DecimalByteParts, - TakeExecuteAdaptor(DecimalByteParts), - ); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mask.rs index e7dc95af84f..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..f9848e1b2e7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/mod.rs @@ -7,6 +7,7 @@ mod filter; pub(crate) mod is_constant; pub(crate) mod kernel; mod mask; +mod slice; mod take; #[cfg(test)] @@ -19,10 +20,36 @@ mod tests { use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array; use vortex_array::compute::conformance::consistency::test_array_consistency; use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::i256; + use vortex_array::validity::Validity; use vortex_buffer::buffer; use crate::DecimalByteParts; use crate::DecimalBytePartsArray; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_of; + use crate::decimal_byte_parts::testing::i256_parts; + + /// Values needing more than 64 bits, so the encoding carries lower parts. + fn wide_i128() -> Vec { + vec![ + 1 << 70, + -(1 << 70), + (1 << 64) - 1, + 0, + 99_999_999_999_999_999_999_999_999_999_999_999_999, + ] + } + + fn wide_i256() -> Vec { + vec![ + i256_of(1, 0), + i256_of(-1, 0), + i256_of(0, u128::MAX), + i256_of(1 << 64, 7), + i256_of(0, 0), + ] + } #[rstest] // Basic decimal byte parts arrays @@ -70,6 +97,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 +121,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..5b07af47252 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,103 @@ 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::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use crate::DecimalByteParts; + use crate::decimal_byte_parts::testing::encode; + use crate::decimal_byte_parts::testing::i256_of; + + /// Taking pushes down into the parts during optimization, with no execution context in + /// play: `ArrayRef::take` wraps the array in a `Dict` and optimizes, and the reduce rule + /// must rewrite that into a `DecimalByteParts` of taken parts. + #[test] + fn take_pushes_down_without_executing() -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + + let decimal = DecimalArray::new( + Buffer::from(vec![1i128 << 70, 2, 3]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + let indices = buffer![0u64, 2].into_array(); + let taken = encode(&decimal)?.into_array().take(indices)?; + + assert!( + taken.is::(), + "expected the take to reduce into the encoding, got {}", + taken.encoding_id() + ); + Ok(()) + } + + /// Taking with nullable indices must still round-trip the wide values, including the + /// null row, on arrays that carry lower parts. + #[rstest] + #[case::one_lower_part(DecimalArray::new( + Buffer::from(vec![1i128 << 70, 2, 3]), + DecimalDType::new(38, 2), + Validity::NonNullable, + ))] + #[case::three_lower_parts(DecimalArray::new( + Buffer::from(vec![i256_of(1, 1 << 70), i256_of(0, 2), i256_of(0, 3)]), + DecimalDType::new(76, 2), + Validity::NonNullable, + ))] + fn take_with_nullable_indices(#[case] decimal: DecimalArray) -> VortexResult<()> { + let session = array_session(); + crate::initialize(&session); + let mut ctx = session.create_execution_ctx(); + + let indices = PrimitiveArray::from_option_iter([Some(0u64), None, Some(2u64)]).into_array(); + let expected = decimal + .clone() + .into_array() + .take(indices.clone())? + .execute::(&mut ctx)?; + + let taken = encode(&decimal)?.into_array().take(indices)?; + let actual = taken.execute::(&mut ctx)?; + + assert_arrays_eq!(expected, actual, &mut ctx); + Ok(()) } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs deleted file mode 100644 index 1e561b149fe..00000000000 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/mod.rs +++ /dev/null @@ -1,324 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Splitting decimal values into 64-bit parts and reassembling them. -//! -//! A `DecimalByteParts` array stores each value as a signed most significant part (MSP) -//! followed by `k` unsigned 64-bit lower parts ordered most significant first. The encoded -//! value is -//! -//! ```text -//! msp * 2^(64k) + Σ_{i, -} - -impl DecimalParts { - /// Construct decimal parts from an MSP with no lower parts. - fn from_msp(values: Buffer, validity: Validity) -> Self { - Self { - msp: PrimitiveArray::new(values, validity).into_array(), - lower_parts: Vec::new(), - } - } - - fn new( - msp: Buffer, - lower_parts: impl IntoIterator>, - validity: Validity, - ) -> Self { - Self { - msp: PrimitiveArray::new(msp, validity).into_array(), - lower_parts: lower_parts - .into_iter() - .map(|part| PrimitiveArray::new(part, Validity::NonNullable).into_array()) - .collect(), - } - } -} - -/// 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. -/// -/// # 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()?; - Ok(match decimal.values_type() { - DecimalType::I8 => DecimalParts::from_msp(decimal.buffer::(), validity), - DecimalType::I16 => DecimalParts::from_msp(decimal.buffer::(), validity), - DecimalType::I32 => DecimalParts::from_msp(decimal.buffer::(), validity), - DecimalType::I64 => DecimalParts::from_msp(decimal.buffer::(), validity), - DecimalType::I128 => { - let mask = validity.execute_mask(decimal.len(), ctx)?; - let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i128_to_parts); - DecimalParts::new(msp, lower, validity) - } - DecimalType::I256 => { - let mask = validity.execute_mask(decimal.len(), ctx)?; - let (msp, lower) = split_wide(&decimal.buffer::(), &mask, i256_to_parts); - DecimalParts::new(msp, lower, validity) - } - }) -} - -/// 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. -fn split_wide( - values: &Buffer, - 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)); - - // Zero out all parts if all null - if validity.all_false() { - msp.push_n(0, len); - for part in &mut lower { - part.push_n(0, len); - } - return (msp.freeze(), lower.map(BufferMut::freeze)); - } - - // Allocate without zeroing, then initialize every part of each row together. - let msp_out = &mut msp.spare_capacity_mut()[..len]; - let mut lower_out = lower - .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!("AllFalse case addressed above"), - } - - // 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] -const fn i128_to_parts(value: i128) -> (i64, [u64; 1]) { - #[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] -const fn i256_to_parts(value: i256) -> (i64, [u64; MAX_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], - ) -} - -/// Reassemble primitive arrays that constitute decimal byte parts into a canonical decimal array. -/// -/// The MSP must be signed. There must be between zero and three (inclusive) `u64` lower parts, ordered -/// most significant first. The lower parts must be non-nullable. Every input array must have the same length. -/// -/// With no lower parts, the MSP buffer is reused as the decimal values. One lower part -/// assembles into `i128`. Two or three lower parts assemble into `i256`. -/// -/// # Errors -/// -/// Returns an error if the parts do not describe a valid decimal, or if the MSP's validity -/// cannot be derived. -pub fn assemble_decimal( - msp: &PrimitiveArray, - lower_parts: &[PrimitiveArray], - decimal_dtype: DecimalDType, -) -> VortexResult { - let validity = msp.validity()?; - vortex_ensure!(msp.dtype().as_ptype().is_signed_int()); - - if lower_parts.is_empty() { - return Ok(match_each_signed_integer_ptype!(msp.ptype(), |P| { - // SAFETY: the buffer is typed by the array's own ptype, the decimal dtype is the - // array's, and the validity is taken from the same array. - unsafe { DecimalArray::new_unchecked(msp.to_buffer::

(), decimal_dtype, validity) } - })); - } - - let len = msp.len(); - let lower: Vec<&[u64]> = lower_parts - .iter() - .map(|part| { - vortex_ensure!( - part.dtype() == &LOWER_PART_DTYPE, - "lower part must be non-nullable u64" - ); - let part = part.as_slice::(); - vortex_ensure!( - part.len() == len, - "lower part has len {}, expected {len}", - part.len() - ); - Ok(part) - }) - .collect::>()?; - - Ok(match lower.as_slice() { - [first] => DecimalArray::new( - assemble_wide::(msp, [first]), - decimal_dtype, - validity, - ), - [first, second] => DecimalArray::new( - assemble_wide::(msp, [first, second]), - decimal_dtype, - validity, - ), - [first, second, third] => DecimalArray::new( - assemble_wide::(msp, [first, second, third]), - decimal_dtype, - validity, - ), - _ => vortex_bail!( - "at most {MAX_LOWER_PARTS} lower parts are supported, got {}", - lower.len() - ), - }) -} - -/// Reassemble a signed MSP and `K` unsigned lower parts into wide integers. -/// -/// Each row starts with the MSP sign-extended to `T`. Appending a lower word shifts the -/// accumulated value left by 64 bits and fills the low bits with that word. Lower parts -/// are appended most significant first. -/// -/// The callers select `i128` for one lower part and `i256` for two or three. Since `K` -/// is constant, the compiler can unroll the loop that appends the lower words. -fn assemble_wide(msp: &PrimitiveArray, lower: [&[u64]; K]) -> Buffer -where - T: NativeDecimalType + Shl + BitOr, -{ - let mut out = BufferMut::::with_capacity(msp.len()); - match_each_signed_integer_ptype!(msp.ptype(), |P| { - out.extend_trusted(msp.as_slice::

().iter().enumerate().map(|(row, value)| { - #[allow( - clippy::useless_conversion, - reason = "the widening to i64 is a no-op only for the i64 arm of the ptype match" - )] - let mut value = T::from(i64::from(*value)).vortex_expect("MSP fits in the output type"); - for part in lower { - value = (value << LOWER_PART_BITS) - | T::from(part[row]).vortex_expect("lower word fits in the output type"); - } - value - })); - }); - out.freeze() -} - -#[cfg(test)] -mod tests; diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs deleted file mode 100644 index 3e3de06c44e..00000000000 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/limbs/tests.rs +++ /dev/null @@ -1,226 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use rstest::rstest; -use vortex_array::VortexSessionExecute; -use vortex_array::array_session; -use vortex_array::arrays::DecimalArray; -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 super::*; - -#[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 expected = PrimitiveArray::new( - decimal - .validity()? - .execute_mask(len, &mut ctx)? - .iter() - .map(|valid| if valid { u64::MAX } else { 0 }) - .collect::>(), - Validity::NonNullable, - ); - let parts = split_decimal(&decimal, &mut ctx)?; - for lower in parts.lower_parts { - assert_arrays_eq!(expected.clone(), lower, &mut ctx); - } - assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); - Ok(()) -} - -fn round_trip(decimal: DecimalArray) -> VortexResult { - let mut ctx = array_session().create_execution_ctx(); - let parts = split_decimal(&decimal, &mut ctx)?; - let msp = parts.msp.execute::(&mut ctx)?; - let lower = parts - .lower_parts - .into_iter() - .map(|part| part.execute::(&mut ctx)) - .collect::>>()?; - assemble_decimal(&msp, &lower, decimal.decimal_dtype()) -} - -#[rstest] -#[case::zero(0)] -#[case::one(1)] -#[case::minus_one(-1)] -#[case::limb_boundary(1i128 << 64)] -#[case::just_below_limb_boundary((1i128 << 64) - 1)] -#[case::negative_limb_boundary(-(1i128 << 64))] -#[case::max(i128::MAX)] -#[case::min(i128::MIN)] -fn test_split_assemble_i128(#[case] value: i128) -> VortexResult<()> { - let decimal = DecimalArray::new( - Buffer::from(vec![value]), - DecimalDType::new(38, 2), - Validity::NonNullable, - ); - let round_tripped = round_trip(decimal)?; - assert_eq!(round_tripped.buffer::().as_slice(), &[value]); - Ok(()) -} - -#[rstest] -#[case::zero(i256::ZERO)] -#[case::one(i256::ONE)] -#[case::minus_one(i256::ZERO - i256::ONE)] -#[case::max(i256::MAX)] -#[case::min(i256::MIN)] -#[case::word_1(i256::from_parts(1u128 << 64, 0))] -#[case::word_2(i256::from_parts(0, 1))] -#[case::word_3(i256::from_parts(0, 1i128 << 64))] -#[case::mixed(i256::from_parts(u128::MAX, -3))] -fn test_split_assemble_i256(#[case] value: i256) -> VortexResult<()> { - let decimal = DecimalArray::new( - Buffer::from(vec![value]), - DecimalDType::new(76, 2), - Validity::NonNullable, - ); - let round_tripped = round_trip(decimal)?; - assert_eq!(round_tripped.buffer::().as_slice(), &[value]); - Ok(()) -} - -#[rstest] -fn test_split_narrow_decimal_has_no_lower_parts( - #[values(Validity::NonNullable, Validity::AllInvalid, Validity::from_iter([true, false, true]))] - validity: Validity, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = DecimalArray::new(buffer![1i32, 2, 3], DecimalDType::new(2, 0), validity); - let parts = split_decimal(&decimal, &mut ctx)?; - assert!(parts.lower_parts.is_empty()); - assert_eq!(parts.msp.dtype().as_ptype(), PType::I32); - let msp = parts.msp.execute::(&mut ctx)?; - assert_eq!( - msp.as_slice::().as_ptr(), - decimal.buffer::().as_ptr() - ); - assert_arrays_eq!(decimal.clone(), round_trip(decimal)?, &mut ctx); - Ok(()) -} - -#[test] -fn test_split_i256_part_count_and_types() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = DecimalArray::new( - Buffer::from(vec![i256::from_i128(i128::MAX), i256::MIN]), - DecimalDType::new(76, 0), - Validity::NonNullable, - ); - let parts = split_decimal(&decimal, &mut ctx)?; - assert_eq!(parts.lower_parts.len(), MAX_LOWER_PARTS); - assert_eq!(parts.msp.dtype().as_ptype(), PType::I64); - for part in &parts.lower_parts { - assert_eq!(part.dtype(), &LOWER_PART_DTYPE); - } - Ok(()) -} - -#[rstest] -fn test_split_i256_part_order( - #[values(Validity::NonNullable, Validity::from_iter([true, false, true]))] validity: Validity, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let decimal = DecimalArray::new( - buffer![ - i256::from_parts((2u128 << 64) | 3, (1i128 << 64) | 4), - i256::ZERO, - i256::from_parts((6u128 << 64) | 7, (-2i128 << 64) | 5), - ], - DecimalDType::new(76, 0), - validity.clone(), - ); - let parts = split_decimal(&decimal, &mut ctx)?; - assert_arrays_eq!( - PrimitiveArray::new(buffer![1i64, 0, -2], validity), - parts.msp, - &mut ctx - ); - assert_eq!(parts.lower_parts.len(), 3); - for (part, expected) in parts.lower_parts.into_iter().zip([ - buffer![4u64, 0, 5], - buffer![2u64, 0, 6], - buffer![3u64, 0, 7], - ]) { - assert_arrays_eq!( - PrimitiveArray::new(expected, Validity::NonNullable), - part, - &mut ctx - ); - } - Ok(()) -} - -#[rstest] -fn test_assemble_rejects_mismatched_lower_lengths( - #[values(1, 2, 3)] lower_count: usize, - #[values(0, 1, 3)] lower_len: usize, -) { - let msp = PrimitiveArray::new(buffer![0i64; 2], Validity::NonNullable); - let mut lower = vec![PrimitiveArray::new(buffer![0u64; 2], Validity::NonNullable); lower_count]; - lower[lower_count - 1] = PrimitiveArray::new(buffer![0u64; lower_len], Validity::NonNullable); - let dtype = DecimalDType::new(if lower_count == 1 { 38 } else { 76 }, 0); - assert!(assemble_decimal(&msp, &lower, dtype).is_err()); -} - -#[rstest] -fn test_assemble_i256_part_order_and_sign_extension( - #[values(false, true)] narrow_msp: bool, - #[values(2, 3)] lower_count: usize, -) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let msp = if narrow_msp { - PrimitiveArray::new(buffer![3i8, -3], Validity::NonNullable) - } else { - PrimitiveArray::new(buffer![3i64, -3], Validity::NonNullable) - }; - let lower = - [4u64, 1, 2].map(|word| PrimitiveArray::new(buffer![word; 2], Validity::NonNullable)); - let dtype = DecimalDType::new(76, 0); - let actual = assemble_decimal(&msp, &lower[3 - lower_count..], dtype)?; - let low = (1u128 << 64) | 2; - let expected = if lower_count == 2 { - buffer![i256::from_parts(low, 3), i256::from_parts(low, -3)] - } else { - buffer![ - i256::from_parts(low, (3i128 << 64) | 4), - i256::from_parts(low, (-3i128 << 64) | 4), - ] - }; - assert_arrays_eq!( - DecimalArray::new(expected, dtype, Validity::NonNullable), - actual, - &mut ctx - ); - Ok(()) -} diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index 951a893a785..919b7bb44a3 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -1,385 +1,59 @@ // 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; +#[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 split::DecimalParts; +pub use split::dbp_encode; +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/prop_tests.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs new file mode 100644 index 00000000000..5630d1706a8 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs @@ -0,0 +1,183 @@ +// 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; +use super::testing::encode; + +/// 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 round_tripped = canonicalize( + encode(&decimal).vortex_expect("encode").into_array(), + &mut ctx, + ); + + assert_eq!(round_tripped.values_type(), decimal.values_type()); + assert_arrays_eq!(decimal, round_tripped, &mut ctx); +} + +/// Decoding an encoded array and encoding it again must not change the values it decodes to. +/// +/// Starting from the encoded side reaches part counts `split_decimal` never produces, so this +/// covers layouts the property above cannot generate. It compares decoded values rather than +/// the arrays themselves because re-encoding normalizes the part count: splitting an `i256` +/// always yields three lower parts, whatever the original array carried. +#[hegel::test] +fn encoded_survives_decode_then_encode(tc: TestCase) { + let (array, _len) = draw_encoded(&tc); + let mut ctx = ctx(); + + let decoded = canonicalize(array.into_array(), &mut ctx); + let re_decoded = canonicalize( + encode(&decoded).vortex_expect("encode").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..a3aa1b4afdc --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -0,0 +1,310 @@ +// 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; + +/// Create a [`DecimalBytePartsArray`] from a [`DecimalArray`] by splitting it into parts. +/// +/// # Errors +/// +/// Returns an error if the decimal cannot be split. +pub fn dbp_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 { + DecimalByteParts::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 [`super::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..2dfe2a55b3c --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -0,0 +1,47 @@ +// 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 vortex_error::VortexResult; + +use super::DecimalBytePartsArray; +use super::dbp_encode; + +/// Encode a canonical decimal array as byte parts, splitting wide values into lower parts. +pub(crate) fn encode(decimal: &DecimalArray) -> VortexResult { + dbp_encode(decimal, &mut array_session().create_execution_ctx()) +} + +/// An `i128`-backed decimal array, encoded as byte parts with one lower part. +pub(crate) fn i128_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { + encode(&DecimalArray::new( + Buffer::from(values), + DecimalDType::new(38, 2), + validity, + )) + .vortex_expect("valid decimal byte parts") +} + +/// An `i256`-backed decimal array, encoded as byte parts with three lower parts. +pub(crate) fn i256_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { + encode(&DecimalArray::new( + Buffer::from(values), + DecimalDType::new(76, 2), + validity, + )) + .vortex_expect("valid decimal byte parts") +} + +/// Build an `i256` from a signed high `i128` and unsigned low `u128`. +pub(super) fn i256_of(high: i128, low: u128) -> i256 { + i256::from_parts(low, high) +} diff --git a/vortex-array/src/dtype/bigint/mod.rs b/vortex-array/src/dtype/bigint/mod.rs index 47195526b1f..b03bdef43c2 100644 --- a/vortex-array/src/dtype/bigint/mod.rs +++ b/vortex-array/src/dtype/bigint/mod.rs @@ -130,6 +130,20 @@ impl From for i256 { } } +impl From for i256 { + #[inline] + fn from(value: i64) -> Self { + Self::from_i128(value.into()) + } +} + +impl From for i256 { + #[inline] + fn from(value: u64) -> Self { + Self::from_i128(value.into()) + } +} + impl Display for i256 { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.0) diff --git a/vortex-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 21e440a7436..5d3f610afc7 100644 --- a/vortex-btrblocks/src/trace_tests.rs +++ b/vortex-btrblocks/src/trace_tests.rs @@ -403,7 +403,7 @@ fn trace_scan_filter_on_compressed_table() -> VortexResult<()> { optimize root=vortex.filter(i16, len=43) session=false reduce_parent static:FilterReduceAdaptor(Dict) slot=0 parent=vortex.filter(i16, len=43) child=vortex.dict(i16, len=4096) -> vortex.dict(i16, len=43) done output=vortex.dict(i16, len=43) - reduce_parent static:DecimalBytePartsFilterPushDownRule slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) + reduce_parent static:FilterReduceAdaptor(DecimalByteParts) slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) done output=vortex.decimal_byte_parts(decimal(15,2), len=43) optimize root=vortex.filter(vortex.date[days](i32), len=43) session=false optimize root=vortex.filter(i32, len=43) session=false @@ -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(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=64) + done output=vortex.decimal_byte_parts(decimal(15,2), len=64) optimize root=vortex.dict(vortex.date[days](i32), len=64) session=false reduce_parent static:TakeReduceAdaptor(Extension) slot=1 parent=vortex.dict(vortex.date[days](i32), len=64) child=vortex.ext(vortex.date[days](i32), len=4096) -> vortex.ext(vortex.date[days](i32), len=64) done output=vortex.ext(vortex.date[days](i32), len=64) diff --git a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs index 3475f26a175..a54df06fb4c 100644 --- a/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs +++ b/vortex-cuda/src/kernel/encodings/decimal_byte_parts.rs @@ -39,6 +39,13 @@ impl CudaExecute for DecimalBytePartsExecutor { .dtype() .as_decimal_opt() .vortex_expect("DecimalBytePartsArray dtype must be decimal"); + + // Reassembling lower parts into wide decimals is not implemented on the GPU; the MSP + // alone is not the value. + if !array.lower_parts().is_empty() { + vortex_bail!("DecimalBytePartsArray with lower parts is not supported on GPU") + } + let msp = array.msp().clone(); let PrimitiveDataParts { buffer, From ad35b58a341c3d0dc19884408ac6f6829d8d3631 Mon Sep 17 00:00:00 2001 From: Matthew Katz <87445739+mhk197@users.noreply.github.com> Date: Wed, 16 Sep 2026 10:04:50 -0400 Subject: [PATCH 03/12] Add `ArrayPlugin` to enable serde for new `DecimalBytePartsArray` format (#9810) Add serialization for `DecimalBytePartsArray` with lower parts while preserving the frozen v1 file format. `DecimalBytePartsPlugin` owns both formats and selects the serialized ID from the array's layout: | Array layout | Serialized ID | | --- | --- | | Signed most significant part (MSP) only | `vortex.decimal_byte_parts` | | Signed MSP plus one to three unsigned lower parts | `vortex.decimal_byte_parts_v2` | Both formats deserialize into the same in-memory array representation. Each lower part represents a 64-bit window, but its storage dtype may be narrowed to `u8`, `u16`, or `u32` when its values fit. - Record the MSP's physical integer type, lower-part count, and each lower part's physical integer type in metadata. Deserialization validates these types and restores each child with its recorded dtype. - Keep the frozen format restricted to zero lower parts and require at least one for v2. Frozen metadata remains byte-identical. - Make DBP's VTable `serialize` and `deserialize` methods return errors directing callers to `DecimalBytePartsPlugin`. Keep DBP metadata and serde helpers in `plugin.rs` so the dedicated plugin owns both formats. - Add a separate v2 compatibility fixture covering positive and negative wide `i128` values and nullable `i256` values. The frozen fixture is unchanged. **Breaking:** directly registering `DecimalByteParts`, or calling its VTable serde methods, now errors for both serialization and deserialization, including the frozen v1 format. Replace: ```rust session.arrays().register(DecimalByteParts); ``` with: ```rust session.arrays().register(DecimalBytePartsPlugin); ``` `vortex_decimal_byte_parts::initialize(&session)` already registers the plugin. Existing v1 files remain readable through it; no file migration is required. --------- Signed-off-by: Matt Katz --- docs/specs/editions.md | 4 +- .../src/decimal_byte_parts/array.rs | 205 +---------- .../src/decimal_byte_parts/compute/filter.rs | 12 +- .../src/decimal_byte_parts/compute/mod.rs | 11 +- .../src/decimal_byte_parts/compute/take.rs | 13 +- .../src/decimal_byte_parts/mod.rs | 6 +- .../src/decimal_byte_parts/plugin/mod.rs | 98 ++++++ .../src/decimal_byte_parts/plugin/tests.rs | 324 ++++++++++++++++++ .../src/decimal_byte_parts/plugin/v1.rs | 84 +++++ .../src/decimal_byte_parts/plugin/v2.rs | 114 ++++++ .../src/decimal_byte_parts/prop_tests.rs | 13 +- .../src/decimal_byte_parts/split.rs | 33 +- .../src/decimal_byte_parts/testing.rs | 31 +- encodings/decimal-byte-parts/src/lib.rs | 4 +- vortex-btrblocks/src/trace_tests.rs | 10 +- .../golden__compact__decimal_prices.snap | 2 +- .../golden__regular__decimal_prices.snap | 2 +- .../encodings/decimal_byte_parts_v2.rs | 105 ++++++ .../arrays/synthetic/encodings/mod.rs | 2 + 19 files changed, 809 insertions(+), 264 deletions(-) create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/mod.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/tests.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v1.rs create mode 100644 encodings/decimal-byte-parts/src/decimal_byte_parts/plugin/v2.rs create mode 100644 vortex-test/compat-gen/src/fixtures/arrays/synthetic/encodings/decimal_byte_parts_v2.rs 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/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs index bcdb3579f53..7059fd09838 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -7,14 +7,12 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; -use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; use vortex_array::ArrayParts; use vortex_array::ArrayRef; -use vortex_array::ArraySlots; use vortex_array::ArrayView; use vortex_array::EqMode; use vortex_array::ExecutionCtx; @@ -24,7 +22,6 @@ use vortex_array::array_slots; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::PType; use vortex_array::scalar::DecimalValue; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -37,92 +34,18 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_session::VortexSession; -use vortex_session::registry::CachedId; -use super::LOWER_PART_DTYPE; use super::MAX_LOWER_PARTS; use super::assemble::assemble_decimal; use super::assemble::assemble_wide_decimal_value; +use super::decimal_byte_parts_v2_id; use super::rules::PARENT_RULES; /// A [`DecimalByteParts`]-encoded Vortex array. pub type DecimalBytePartsArray = Array; -#[derive(Clone, prost::Message)] -pub struct DecimalBytesPartsMetadata { - #[prost(enumeration = "PType", tag = "1")] - zeroth_child_ptype: i32, - #[prost(uint32, tag = "2")] - lower_part_count: u32, -} - -impl DecimalBytesPartsMetadata { - fn from_array(array: ArrayView<'_, DecimalByteParts>) -> VortexResult { - Ok(Self { - zeroth_child_ptype: PType::try_from(array.msp().dtype())? as i32, - lower_part_count: u32::try_from(array.lower_parts().len()) - .map_err(|_| vortex_err!("lower part count exceeds u32"))?, - }) - } - - fn into_array_parts( - self, - dtype: &DType, - len: usize, - children: &dyn ArrayChildren, - ) -> VortexResult> { - vortex_ensure!( - dtype.as_decimal_opt().is_some(), - "decoding decimal but given non decimal dtype {dtype}" - ); - - let encoded_dtype = DType::Primitive(self.zeroth_child_ptype(), dtype.nullability()); - - let lower_part_count = self.lower_part_count()?; - vortex_ensure!( - children.len() == DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, - "expected {} children, got {}", - DecimalBytePartsSlots::FIXED_COUNT + lower_part_count, - children.len() - ); - - let msp = children.get(DecimalBytePartsSlots::MSP, &encoded_dtype, len)?; - - let mut slots = ArraySlots::with_capacity(children.len()); - slots.push(Some(msp)); - for idx in 0..lower_part_count { - slots.push(Some(children.get( - DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx, - &LOWER_PART_DTYPE, - len, - )?)); - } - - Ok( - ArrayParts::new(DecimalByteParts, dtype.clone(), len, DecimalBytePartsData) - .with_slots(slots), - ) - } - - /// The number of lower parts encoded in this array. - /// - /// # Errors - /// - /// Returns an error if the count exceeds [`MAX_LOWER_PARTS`]. - fn lower_part_count(&self) -> VortexResult { - let count = usize::try_from(self.lower_part_count) - .map_err(|_| vortex_err!("lower part count {} out of range", self.lower_part_count))?; - vortex_ensure!( - count <= MAX_LOWER_PARTS, - "at most {MAX_LOWER_PARTS} lower parts are supported, got {count}" - ); - Ok(count) - } -} - /// This array encodes decimals by splitting them between 1-4 columns of primitive typed children. /// /// The most significant part (MSP) stores the most significant decimal bits. It is signed and is @@ -222,7 +145,7 @@ impl DecimalByteParts { /// /// 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 - /// canonical decimal array. + /// decimal array. /// /// # Errors /// @@ -274,8 +197,7 @@ impl VTable for DecimalByteParts { type ValidityVTable = ValidityVTableFromChild; fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("vortex.decimal_byte_parts"); - *ID + decimal_byte_parts_v2_id() } fn validate( @@ -331,33 +253,22 @@ impl VTable for DecimalByteParts { } fn serialize( - array: ArrayView<'_, Self>, + _array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - vortex_ensure!( - array.lower_parts().is_empty(), - "serializing DecimalByteParts with lower parts is not supported" - ); - Ok(Some( - DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), - )) + vortex_bail!("DecimalByteParts serialization requires DecimalBytePartsPlugin") } fn deserialize( &self, - dtype: &DType, - len: usize, - metadata: &[u8], + _dtype: &DType, + _len: usize, + _metadata: &[u8], _buffers: &[BufferHandle], - children: &dyn ArrayChildren, + _children: &dyn ArrayChildren, _session: &VortexSession, ) -> VortexResult> { - let metadata = DecimalBytesPartsMetadata::decode(metadata)?; - vortex_ensure!( - metadata.lower_part_count()? == 0, - "vortex.decimal_byte_parts must not carry lower parts" - ); - metadata.into_array_parts(dtype, len, children) + vortex_bail!("DecimalByteParts deserialization requires DecimalBytePartsPlugin") } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { @@ -510,18 +421,15 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; use vortex_array::validity::Validity; - use vortex_array::vtable::VTable; use vortex_buffer::buffer; use vortex_error::VortexResult; use super::DecimalByteParts; - use super::DecimalBytePartsArray; use super::DecimalBytePartsArraySlotsExt; use super::DecimalBytePartsData; use crate::decimal_byte_parts::LOWER_PART_DTYPE; use crate::decimal_byte_parts::MAX_LOWER_PARTS; use crate::decimal_byte_parts::testing::i128_parts; - use crate::decimal_byte_parts::testing::i256_of; use crate::decimal_byte_parts::testing::i256_parts; #[test] @@ -563,62 +471,6 @@ mod tests { ); } - /// The largest unscaled value a `Decimal(38, _)` can hold: `10^38 - 1`. - const MAX_PRECISION_38: i128 = 99_999_999_999_999_999_999_999_999_999_999_999_999; - - /// The largest unscaled value a `Decimal(76, _)` can hold: `10^76 - 1`. - fn max_precision_76() -> i256 { - i256::from_i128(10).wrapping_pow(76) - i256::ONE - } - - /// Values that exercise every 64-bit window of an `i128`, both signs, and the boundaries - /// where a lower part carries into the MSP. - fn wide_i128_values() -> Vec { - vec![ - 0, - 1, - -1, - (1 << 64) - 1, - 1 << 64, - -(1 << 64), - -((1 << 64) + 1), - MAX_PRECISION_38, - -MAX_PRECISION_38, - 1 << 100, - ] - } - - /// Values that exercise every 64-bit window of an `i256`. - fn wide_i256_values() -> Vec { - vec![ - i256::ZERO, - i256::ONE, - i256::ZERO - i256::ONE, - i256_of(0, u128::MAX), - i256_of(1, 0), - i256_of(-1, 0), - i256_of(-1, u128::MAX - 1), - i256_of(1 << 64, 12345), - max_precision_76(), - i256::ZERO - max_precision_76(), - ] - } - - #[rstest] - #[case::i128_non_nullable(i128_parts(wide_i128_values(), Validity::NonNullable))] - #[case::i256_non_nullable(i256_parts(wide_i256_values(), Validity::NonNullable))] - fn test_canonical_decimal_round_trips( - #[case] array: DecimalBytePartsArray, - ) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let canonical = array - .clone() - .into_array() - .execute::(&mut ctx)?; - assert_arrays_eq!(array, canonical, &mut ctx); - Ok(()) - } - #[test] fn test_lower_part_layout_i128() -> VortexResult<()> { let array = i128_parts(vec![(3i128 << 64) | 7], Validity::NonNullable); @@ -639,7 +491,7 @@ mod tests { #[test] fn test_lower_part_layout_i256() -> VortexResult<()> { let array = i256_parts( - vec![i256_of((5i128 << 64) | 6, (7u128 << 64) | 8)], + vec![i256::from_parts((7u128 << 64) | 8, (5i128 << 64) | 6)], Validity::NonNullable, ); assert_eq!(array.lower_parts().len(), MAX_LOWER_PARTS); @@ -654,27 +506,6 @@ mod tests { Ok(()) } - #[rstest] - #[case::i128(i128_parts(wide_i128_values(), Validity::AllValid))] - #[case::i256(i256_parts(wide_i256_values(), Validity::AllValid))] - fn test_scalar_at_matches_canonical(#[case] array: DecimalBytePartsArray) -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let canonical = array - .clone() - .into_array() - .execute::(&mut ctx)? - .into_array(); - let array = array.into_array(); - for idx in 0..array.len() { - assert_eq!( - array.execute_scalar(idx, &mut ctx)?, - canonical.execute_scalar(idx, &mut ctx)?, - "scalar mismatch at index {idx}" - ); - } - Ok(()) - } - #[rstest] fn test_scalar_at_matches_canonical_for_each_part_count( #[values(false, true)] narrow_msp: bool, @@ -860,7 +691,7 @@ mod tests { let canonical = i128_array.into_array().execute::(&mut ctx)?; assert_eq!(canonical.values_type(), DecimalType::I128); - let i256_array = i256_parts(vec![i256_of(1 << 100, 0)], Validity::NonNullable); + let 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); @@ -885,7 +716,10 @@ mod tests { )?; let canonical = array.into_array().execute::(&mut ctx)?; assert_eq!(canonical.values_type(), DecimalType::I256); - assert_eq!(canonical.buffer::().as_slice(), &[i256_of(1, 9)]); + assert_eq!( + canonical.buffer::().as_slice(), + &[i256::from_parts(9, 1)] + ); Ok(()) } @@ -912,11 +746,4 @@ mod tests { assert_arrays_eq!(array, canonical.into_array(), &mut ctx); Ok(()) } - #[test] - fn test_frozen_serializer_rejects_lower_parts() -> VortexResult<()> { - let session = array_session(); - let array = i128_parts(vec![1i128 << 70], Validity::NonNullable); - assert!(VTable::serialize(array.as_view(), &session).is_err()); - Ok(()) - } } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/compute/filter.rs index 49c4021dd18..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 @@ -27,12 +27,12 @@ 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_of; use crate::decimal_byte_parts::testing::i256_parts; #[test] @@ -72,11 +72,11 @@ mod test { let array = i256_parts( vec![ - i256_of(1, 0), - i256_of(-1, 5), - i256_of(0, u128::MAX), - i256_of(1 << 64, 7), - i256_of(0, 0), + 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]), ); 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 f9848e1b2e7..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 @@ -27,7 +27,6 @@ mod tests { use crate::DecimalByteParts; use crate::DecimalBytePartsArray; use crate::decimal_byte_parts::testing::i128_parts; - use crate::decimal_byte_parts::testing::i256_of; use crate::decimal_byte_parts::testing::i256_parts; /// Values needing more than 64 bits, so the encoding carries lower parts. @@ -43,11 +42,11 @@ mod tests { fn wide_i256() -> Vec { vec![ - i256_of(1, 0), - i256_of(-1, 0), - i256_of(0, u128::MAX), - i256_of(1 << 64, 7), - i256_of(0, 0), + 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), ] } 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 5b07af47252..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 @@ -38,14 +38,13 @@ mod tests { 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; - use crate::decimal_byte_parts::testing::encode; - use crate::decimal_byte_parts::testing::i256_of; /// Taking pushes down into the parts during optimization, with no execution context in /// play: `ArrayRef::take` wraps the array in a `Dict` and optimizes, and the reduce rule @@ -61,7 +60,9 @@ mod tests { Validity::NonNullable, ); let indices = buffer![0u64, 2].into_array(); - let taken = encode(&decimal)?.into_array().take(indices)?; + let taken = DecimalByteParts::encode(&decimal, &mut session.create_execution_ctx())? + .into_array() + .take(indices)?; assert!( taken.is::(), @@ -80,7 +81,7 @@ mod tests { Validity::NonNullable, ))] #[case::three_lower_parts(DecimalArray::new( - Buffer::from(vec![i256_of(1, 1 << 70), i256_of(0, 2), i256_of(0, 3)]), + 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, ))] @@ -96,7 +97,9 @@ mod tests { .take(indices.clone())? .execute::(&mut ctx)?; - let taken = encode(&decimal)?.into_array().take(indices)?; + let taken = DecimalByteParts::encode(&decimal, &mut ctx)? + .into_array() + .take(indices)?; let actual = taken.execute::(&mut ctx)?; assert_arrays_eq!(expected, actual, &mut ctx); 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 919b7bb44a3..65bb8222f7a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -22,6 +22,7 @@ use vortex_array::dtype::PType; mod array; mod assemble; pub(crate) mod compute; +mod plugin; #[cfg(test)] mod prop_tests; mod rules; @@ -30,8 +31,11 @@ mod split; mod testing; 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::dbp_encode; pub use split::split_decimal; #[doc(hidden)] 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 index 5630d1706a8..b1625485ee7 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/prop_tests.rs @@ -21,7 +21,6 @@ use vortex_error::VortexExpect; use super::DecimalByteParts; use super::DecimalBytePartsArray; -use super::testing::encode; /// Largest magnitude a `Decimal(38, _)` can hold: 38 nines. const MAX_I128: i128 = 10i128.pow(38) - 1; @@ -139,10 +138,8 @@ fn decoded_survives_encode_then_decode(tc: TestCase) { let decimal = draw_decimal(&tc); let mut ctx = ctx(); - let round_tripped = canonicalize( - encode(&decimal).vortex_expect("encode").into_array(), - &mut 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); @@ -160,10 +157,8 @@ fn encoded_survives_decode_then_encode(tc: TestCase) { let mut ctx = ctx(); let decoded = canonicalize(array.into_array(), &mut ctx); - let re_decoded = canonicalize( - encode(&decoded).vortex_expect("encode").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); } diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs index a3aa1b4afdc..8d8818924e9 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -27,22 +27,23 @@ use super::LOWER_PART_BITS; use super::MAX_I128_LOWER_PARTS; use super::MAX_I256_LOWER_PARTS; -/// Create a [`DecimalBytePartsArray`] from a [`DecimalArray`] by splitting it into parts. -/// -/// # Errors -/// -/// Returns an error if the decimal cannot be split. -pub fn dbp_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 { - DecimalByteParts::new_unchecked(parts.msp, parts.lower_parts, decimal.decimal_dtype()) - }) +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. diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs index 2dfe2a55b3c..af270bad407 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -11,37 +11,24 @@ use vortex_array::dtype::i256; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_error::VortexExpect; -use vortex_error::VortexResult; +use super::DecimalByteParts; use super::DecimalBytePartsArray; -use super::dbp_encode; - -/// Encode a canonical decimal array as byte parts, splitting wide values into lower parts. -pub(crate) fn encode(decimal: &DecimalArray) -> VortexResult { - dbp_encode(decimal, &mut array_session().create_execution_ctx()) -} /// An `i128`-backed decimal array, encoded as byte parts with one lower part. pub(crate) fn i128_parts(values: Vec, validity: Validity) -> DecimalBytePartsArray { - encode(&DecimalArray::new( - Buffer::from(values), - DecimalDType::new(38, 2), - validity, - )) + 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 { - encode(&DecimalArray::new( - Buffer::from(values), - DecimalDType::new(76, 2), - validity, - )) + DecimalByteParts::encode( + &DecimalArray::new(Buffer::from(values), DecimalDType::new(76, 2), validity), + &mut array_session().create_execution_ctx(), + ) .vortex_expect("valid decimal byte parts") } - -/// Build an `i256` from a signed high `i128` and unsigned low `u128`. -pub(super) fn i256_of(high: i128, low: u128) -> i256 { - i256::from_parts(low, high) -} diff --git a/encodings/decimal-byte-parts/src/lib.rs b/encodings/decimal-byte-parts/src/lib.rs index 36a53c3a614..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-btrblocks/src/trace_tests.rs b/vortex-btrblocks/src/trace_tests.rs index 5d3f610afc7..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:FilterReduceAdaptor(DecimalByteParts) slot=0 parent=vortex.filter(decimal(15,2), len=43) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=43) - done output=vortex.decimal_byte_parts(decimal(15,2), len=43) + 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) @@ -440,8 +440,8 @@ fn trace_scan_take_on_compressed_table() -> VortexResult<()> { insta::assert_snapshot!(optimized.trace.to_string(), @" optimize root=vortex.dict({l_quantity=decimal(15,2), l_shipdate=vortex.date[days](i32), l_shipmode=utf8}, len=64) session=false optimize root=vortex.dict(decimal(15,2), len=64) session=false - reduce_parent static:TakeReduceAdaptor(DecimalByteParts) slot=1 parent=vortex.dict(decimal(15,2), len=64) child=vortex.decimal_byte_parts(decimal(15,2), len=4096) -> vortex.decimal_byte_parts(decimal(15,2), len=64) - done output=vortex.decimal_byte_parts(decimal(15,2), len=64) + 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/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-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), From c5ea5c54a801c24592230a5086ab9b8311fa2702 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 16 Sep 2026 12:02:35 -0400 Subject: [PATCH 04/12] Declare the frozen v1 wire ID from the decimal compression scheme Signed-off-by: Matt Katz --- vortex-btrblocks/src/schemes/decimal.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) 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. From 35adc87189635c951022538af21bde4d4d557e3e Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 16 Sep 2026 13:51:18 -0400 Subject: [PATCH 05/12] Do not link the private lower-part limit from public docs Signed-off-by: Matt Katz --- encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs | 4 ++-- encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs index 7059fd09838..c1826ae46d4 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/array.rs @@ -82,7 +82,7 @@ impl DecimalBytePartsData { /// # 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`] + /// 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, @@ -169,7 +169,7 @@ impl DecimalByteParts { /// # 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 + /// 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, diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs index 8d8818924e9..f33101e4a3a 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/split.rs @@ -65,7 +65,7 @@ impl DecimalParts { } /// Construct decimal parts arrays from the buffers constituting a wide decimal (`i128` or `i256`). - /// Wide decimals have an `i64` MSP and up to [`super::MAX_LOWER_PARTS`] `u64` lower parts. + /// Wide decimals have an `i64` MSP and up to `MAX_LOWER_PARTS` `u64` lower parts. fn from_wide( msp: Buffer, lower_parts: impl IntoIterator>, From 7db92a09f91ce2204d425b9fc0cd6122ec1c61d5 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Fri, 4 Sep 2026 11:15:20 -0400 Subject: [PATCH 06/12] Tell the compressor which serialized IDs the writer may emit CascadingCompressor carries the snapshot of serialized IDs the writer may emit, filled by the file writer from the enabled editions through BtrBlocksCompressorBuilder::allow_serialized_ids. A scheme whose encoding has more than one wire format picks its compression mode from it with allows_serialized_id, the newest permitted one; without a restriction every ID is allowed. No scheme consults the set yet. This is the mechanism docs/specs/editions.md describes under compression with replacement encodings (#9779). Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 36 +++++++++- vortex-compressor/src/compressor/mod.rs | 34 ++++++++++ vortex-compressor/src/compressor/tests.rs | 81 +++++++++++++++++++++++ vortex-file/src/writer.rs | 3 +- 4 files changed, 152 insertions(+), 2 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 341569f408c..e2caa67798b 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -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,6 +109,7 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + allowed_serialized_ids: None, } } @@ -210,15 +213,33 @@ impl BtrBlocksCompressorBuilder { self } + /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier + /// call. A scheme whose encoding has several wire formats picks its compression mode from this + /// set: the newest permitted one. + /// + /// The file writer passes the serialized IDs its enabled editions permit. + pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { + self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(allowed).copied().collect(), + None => allowed.clone(), + }); + self + } + /// Builds the configured [`BtrBlocksCompressor`]. 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_array::arrays::Bool; use vortex_fastlanes::FoR; use super::*; @@ -283,6 +304,19 @@ mod tests { } } + /// Every serialized ID is allowed until the writer narrows the set to its editions. + #[test] + fn allowed_serialized_ids_reach_the_compressor() { + let default = BtrBlocksCompressorBuilder::default().build(); + assert!(default.0.allows_serialized_id(Bool.id())); + + let narrowed = BtrBlocksCompressorBuilder::default() + .allow_serialized_ids(&HashSet::from([FoR.id()])) + .build(); + assert!(narrowed.0.allows_serialized_id(FoR.id())); + assert!(!narrowed.0.allows_serialized_id(Bool.id())); + } + #[test] fn cuda_compatible_uses_fsst_for_strings() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 219b67e2519..774513c66a7 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,9 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -46,6 +49,10 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// The serialized IDs the output may use, or `None` for no restriction. See + /// [`allows_serialized_id`](Self::allows_serialized_id). + allowed_serialized_ids: Option>, } impl CascadingCompressor { @@ -63,9 +70,36 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + allowed_serialized_ids: None, } } + /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier + /// call. + /// + /// The file writer passes the serialized IDs its enabled editions permit. A scheme whose + /// encoding has several wire formats picks its compression mode from this set, the newest + /// permitted one, before estimating or compressing. + pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { + self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + Some(existing) => existing.intersection(&allowed).copied().collect(), + None => allowed, + }); + self + } + + /// Returns whether the writer may emit the serialized ID `id`. + /// + /// Schemes whose encoding has several wire formats consult this to pick their compression + /// mode. Without a restriction every ID is allowed, so the newest mode is chosen. The + /// serializer still emits the oldest wire form the resulting array fits, and the + /// serialization context validates that ID. + pub fn allows_serialized_id(&self, id: ArrayId) -> bool { + self.allowed_serialized_ids + .as_ref() + .is_none_or(|allowed| allowed.contains(&id)) + } + /// Returns whether the compressor was configured with `scheme`. pub fn has_scheme(&self, scheme: SchemeId) -> bool { self.schemes diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 3a2a6281047..5f50c62b6c6 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -9,11 +9,14 @@ 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::arrays::Bool; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -26,6 +29,7 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -96,6 +100,48 @@ impl Scheme for DirectRatioScheme { } } +/// What the last `FormatRecordingScheme::compress` call saw for `allows_serialized_id`. +static SEEN_FORMAT: Mutex> = Mutex::new(None); + +/// Stands in for a scheme whose encoding has several wire formats: it asks the compressor whether +/// the newer one is allowed and records the answer. +#[derive(Debug)] +struct FormatRecordingScheme; + +impl Scheme for FormatRecordingScheme { + fn scheme_name(&self) -> &'static str { + "test.format_recording" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + *SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id())); + Ok(data.array().clone()) + } +} + #[derive(Debug)] struct ImmediateAlwaysUseScheme; @@ -845,3 +891,38 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } + +#[test] +fn allowed_serialized_ids_default_to_everything_and_intersect() { + let compressor = compressor(); + assert!(compressor.allows_serialized_id(Constant.id())); + assert!(compressor.allows_serialized_id(Bool.id())); + + let restricted = + compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); + assert!(restricted.allows_serialized_id(Constant.id())); + assert!(!restricted.allows_serialized_id(Bool.id())); + + let narrowed = + restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); + assert!(narrowed.allows_serialized_id(Primitive.id())); + assert!(!narrowed.allows_serialized_id(Constant.id())); + assert!(!narrowed.allows_serialized_id(Bool.id())); +} + +/// A scheme sees the restriction through the compressor it is handed: everything is allowed until +/// the writer narrows the set to its editions. +#[test] +fn schemes_see_the_allowed_serialized_ids() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0..4096i32).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let unrestricted = CascadingCompressor::new(vec![&FormatRecordingScheme]); + unrestricted.compress(&array, &mut exec_ctx)?; + assert_eq!(*SEEN_FORMAT.lock(), Some(true)); + + let restricted = unrestricted.with_allowed_serialized_ids(HashSet::from([Primitive.id()])); + restricted.compress(&array, &mut exec_ctx)?; + assert_eq!(*SEEN_FORMAT.lock(), Some(false)); + Ok(()) +} diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 6c0d01dfa05..c9efe05274b 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -254,7 +254,8 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_serialized_ids), + .retain_allowed_encodings(&allowed_serialized_ids) + .allow_serialized_ids(&allowed_serialized_ids), ) .build(), }; From 8f52553c46f4189aca56920e468db8732277017b Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 10:47:58 -0400 Subject: [PATCH 07/12] Enable a scheme when any of its serialized IDs is permitted Scheme::produced_encodings now names the serialized IDs a scheme may write its output under, oldest first. BtrBlocksCompressorBuilder::allow_serialized_ids replaces retain_allowed_encodings: it keeps a scheme when at least one of those IDs is permitted and hands the set to the compressor, so the writer makes one call from the serialized IDs its editions permit instead of mapping them back to in-memory encodings, which could not tell two wire formats of one encoding apart. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 97 +++++++++++++++++++++++------ vortex-compressor/src/scheme/mod.rs | 14 ++--- vortex-file/src/writer.rs | 1 - 3 files changed, 85 insertions(+), 27 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index e2caa67798b..ddecd4faa8b 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -203,26 +203,22 @@ impl BtrBlocksCompressorBuilder { 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))); - self - } - - /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier - /// call. A scheme whose encoding has several wire formats picks its compression mode from this - /// set: the newest permitted one. + /// A scheme stays when at least one of its [produced IDs](Scheme::produced_encodings) is + /// permitted, and the compressor is handed the set so a scheme whose encoding has several + /// wire formats picks its compression mode from it: the newest permitted one. /// /// The file writer passes the serialized IDs its enabled editions permit. pub fn allow_serialized_ids(mut self, allowed: &HashSet) -> Self { - self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + let allowed: HashSet = match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(allowed).copied().collect(), None => allowed.clone(), - }); + }; + self.schemes + .retain(|s| s.produced_encodings().iter().any(|id| allowed.contains(id))); + self.allowed_serialized_ids = Some(allowed); self } @@ -238,11 +234,20 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; use vortex_array::VTable; use vortex_array::arrays::Bool; + use vortex_array::arrays::Primitive; + use vortex_compressor::scheme::CompressionEstimate; + use vortex_compressor::scheme::EstimateVerdict; + use vortex_error::VortexResult; use vortex_fastlanes::FoR; use super::*; + use crate::ArrayAndStats; + use crate::CompressorContext; #[test] fn empty_starts_with_no_schemes() { @@ -257,26 +262,80 @@ mod tests { } #[test] - fn retain_allowed_encodings_filters_schemes() { + fn allow_serialized_ids_filters_schemes() { let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&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()); + let none = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&HashSet::new()); assert!(none.schemes.is_empty()); } #[test] - fn retaining_all_declared_outputs_keeps_every_scheme() { + fn allowing_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); + let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } + /// Stands in for a scheme whose encoding has two wire formats. + #[derive(Debug)] + struct TwoFormatScheme; + + impl Scheme for TwoFormatScheme { + fn scheme_name(&self) -> &'static str { + "test.two_formats" + } + + fn matches(&self, _canonical: &Canonical) -> bool { + false + } + + fn produced_encodings(&self) -> Vec { + vec![FoR.id(), Bool.id()] + } + + 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 { + unreachable!("test helper never matches") + } + } + + /// A scheme with several wire formats stays while any of them is permitted; which one it + /// produces is decided when compressing. + #[test] + fn any_permitted_format_keeps_the_scheme() { + static TWO_FORMATS: TwoFormatScheme = TwoFormatScheme; + + let newer_only = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&TWO_FORMATS) + .allow_serialized_ids(&HashSet::from([Bool.id()])); + assert_eq!(newer_only.schemes.len(), 1); + + let neither = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&TWO_FORMATS) + .allow_serialized_ids(&HashSet::from([Primitive.id()])); + assert!(neither.schemes.is_empty()); + } + #[test] fn cuda_compatible_excludes_alprd() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..fa42231d422 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,15 +124,15 @@ 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 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. + /// An encoding with several wire formats lists every one of them, oldest first. The writer + /// keeps the scheme while any of them is permitted, and the scheme picks the newest + /// permitted one as its compression mode. fn produced_encodings(&self) -> Vec; /// Returns the stats generation options this scheme requires. The compressor merges all diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index c9efe05274b..e6bb9be4ff8 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -254,7 +254,6 @@ impl VortexWriteOptions { None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_serialized_ids) .allow_serialized_ids(&allowed_serialized_ids), ) .build(), From 5a90682006ca52be9548696304c12470c8ecee40 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 10:49:53 -0400 Subject: [PATCH 08/12] Read the permitted serialized IDs from the CompressorContext The compressor seeds each root CompressorContext with its permitted serialized IDs and every descent inherits them, so a scheme asks compress_ctx.allows_serialized_id both while estimating and while compressing and picks the same mode in both. The per-compressor accessor goes; allowed_serialized_ids remains for inspection. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 8 +- vortex-compressor/src/compressor/cascade.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 38 +++---- vortex-compressor/src/compressor/tests.rs | 110 ++++++++++++++------ vortex-compressor/src/scheme/ctx.rs | 27 ++++- 5 files changed, 127 insertions(+), 58 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index ddecd4faa8b..fcf2bb3934b 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -367,13 +367,15 @@ mod tests { #[test] fn allowed_serialized_ids_reach_the_compressor() { let default = BtrBlocksCompressorBuilder::default().build(); - assert!(default.0.allows_serialized_id(Bool.id())); + assert!(default.0.allowed_serialized_ids().is_none()); let narrowed = BtrBlocksCompressorBuilder::default() .allow_serialized_ids(&HashSet::from([FoR.id()])) .build(); - assert!(narrowed.0.allows_serialized_id(FoR.id())); - assert!(!narrowed.0.allows_serialized_id(Bool.id())); + assert_eq!( + narrowed.0.allowed_serialized_ids(), + Some(&HashSet::from([FoR.id()])) + ); } #[test] diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 86d45d2c0d9..dd98f4ea3c6 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()); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 774513c66a7..159842f2595 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,11 +9,14 @@ mod sample; mod select; mod structural; +use std::sync::Arc; + use vortex_array::ArrayId; 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; @@ -50,9 +53,9 @@ pub struct CascadingCompressor { /// list offsets). root_exclusions: Vec, - /// The serialized IDs the output may use, or `None` for no restriction. See - /// [`allows_serialized_id`](Self::allows_serialized_id). - allowed_serialized_ids: Option>, + /// The serialized IDs the writer may emit, or `None` for no restriction. Seeds every root + /// [`CompressorContext`], where schemes read it. + allowed_serialized_ids: Option>>, } impl CascadingCompressor { @@ -77,27 +80,26 @@ impl CascadingCompressor { /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier /// call. /// - /// The file writer passes the serialized IDs its enabled editions permit. A scheme whose - /// encoding has several wire formats picks its compression mode from this set, the newest - /// permitted one, before estimating or compressing. + /// The file writer passes the serialized IDs its enabled editions permit. Schemes read the + /// set through [`CompressorContext::allows_serialized_id`], so a scheme whose encoding has + /// several wire formats picks the newest permitted one as its mode, while estimating and + /// while compressing alike. pub fn with_allowed_serialized_ids(mut self, allowed: HashSet) -> Self { - self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() { + self.allowed_serialized_ids = Some(Arc::new(match self.allowed_serialized_ids.take() { Some(existing) => existing.intersection(&allowed).copied().collect(), None => allowed, - }); + })); self } - /// Returns whether the writer may emit the serialized ID `id`. - /// - /// Schemes whose encoding has several wire formats consult this to pick their compression - /// mode. Without a restriction every ID is allowed, so the newest mode is chosen. The - /// serializer still emits the oldest wire form the resulting array fits, and the - /// serialization context validates that ID. - pub fn allows_serialized_id(&self, id: ArrayId) -> bool { - self.allowed_serialized_ids - .as_ref() - .is_none_or(|allowed| allowed.contains(&id)) + /// The serialized IDs the writer may emit, or `None` when unrestricted. + pub fn allowed_serialized_ids(&self) -> Option<&HashSet> { + self.allowed_serialized_ids.as_deref() + } + + /// The context a compress call starts from. + pub(crate) fn root_context(&self) -> CompressorContext { + CompressorContext::new(self.allowed_serialized_ids.clone()) } /// Returns whether the compressor was configured with `scheme`. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 5f50c62b6c6..ef388c7d373 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -132,12 +132,12 @@ impl Scheme for FormatRecordingScheme { fn compress( &self, - compressor: &CascadingCompressor, + _compressor: &CascadingCompressor, data: &ArrayAndStats, - _compress_ctx: CompressorContext, + compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - *SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id())); + *SEEN_FORMAT.lock() = Some(compress_ctx.allows_serialized_id(Constant.id())); Ok(data.array().clone()) } } @@ -420,8 +420,12 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -438,8 +442,12 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -456,8 +464,12 @@ fn callback_skip_is_ignored() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -474,8 +486,12 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -492,8 +508,12 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -510,8 +530,12 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -528,8 +552,12 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(winner.is_none()); Ok(()) @@ -630,8 +658,12 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -651,7 +683,7 @@ fn threshold_reflects_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -672,7 +704,7 @@ fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner // `None`) because the zero-byte sample is never stored as the best. @@ -691,7 +723,7 @@ fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert_eq!(observed, Some(None)); @@ -711,7 +743,7 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -732,8 +764,12 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + let winner = compressor.choose_best_scheme( + &schemes, + &data, + CompressorContext::new(None), + &mut exec_ctx, + )?; assert!(matches!( winner, @@ -780,7 +816,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // A context with default stats_options (count_distinct_values = false) and // marked as a sample so the function skips the sampling step and compresses // the array directly. - let ctx = CompressorContext::new().with_sampling(); + let ctx = CompressorContext::new(None).with_sampling(); // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" @@ -895,19 +931,27 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { #[test] fn allowed_serialized_ids_default_to_everything_and_intersect() { let compressor = compressor(); - assert!(compressor.allows_serialized_id(Constant.id())); - assert!(compressor.allows_serialized_id(Bool.id())); + let root = compressor.root_context(); + assert!(root.allows_serialized_id(Constant.id())); + assert!(root.allows_serialized_id(Bool.id())); let restricted = compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); - assert!(restricted.allows_serialized_id(Constant.id())); - assert!(!restricted.allows_serialized_id(Bool.id())); + let root = restricted.root_context(); + assert!(root.allows_serialized_id(Constant.id())); + assert!(!root.allows_serialized_id(Bool.id())); let narrowed = restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); - assert!(narrowed.allows_serialized_id(Primitive.id())); - assert!(!narrowed.allows_serialized_id(Constant.id())); - assert!(!narrowed.allows_serialized_id(Bool.id())); + let root = narrowed.root_context(); + assert!(root.allows_serialized_id(Primitive.id())); + assert!(!root.allows_serialized_id(Constant.id())); + assert!(!root.allows_serialized_id(Bool.id())); + + // Descending keeps the set. + let child = root.descend_with_scheme(IntDictScheme.id(), 0); + assert!(child.allows_serialized_id(Primitive.id())); + assert!(!child.allows_serialized_id(Constant.id())); } /// A scheme sees the restriction through the compressor it is handed: everything is allowed until diff --git a/vortex-compressor/src/scheme/ctx.rs b/vortex-compressor/src/scheme/ctx.rs index 4eed7538daa..83685031c33 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,8 +4,11 @@ //! Compression context for recursive compression. use std::fmt; +use std::sync::Arc; +use vortex_array::ArrayId; use vortex_error::VortexExpect; +use vortex_utils::aliases::hash_set::HashSet; use crate::compressor::ROOT_SCHEME_ID; use crate::scheme::SchemeId; @@ -38,18 +41,24 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, + + /// The serialized IDs the writer may emit, or `None` for no restriction. Shared by every + /// context of one compress call, so cloning at each descent is a pointer bump. + allowed_serialized_ids: Option>>, } impl CompressorContext { - /// Creates a new `CompressorContext`. + /// Creates a new root `CompressorContext` for a compressor that may emit the given serialized + /// IDs, or any ID when `None`. /// /// This should **only** be created by the compressor. - pub(crate) fn new() -> Self { + pub(crate) fn new(allowed_serialized_ids: Option>>) -> Self { Self { is_sample: false, allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), + allowed_serialized_ids, } } } @@ -57,7 +66,7 @@ impl CompressorContext { #[cfg(test)] impl Default for CompressorContext { fn default() -> Self { - Self::new() + Self::new(None) } } @@ -67,6 +76,18 @@ impl CompressorContext { self.is_sample } + /// Returns whether the writer may emit the serialized ID `id`. + /// + /// A scheme whose encoding has several wire formats picks its compression mode from this, + /// the newest permitted one, and the same answer is available while estimating and while + /// compressing. Without a restriction every ID is allowed. The serializer still emits the + /// oldest wire form the resulting array fits, and the serialization context validates it. + pub fn allows_serialized_id(&self, id: ArrayId) -> bool { + self.allowed_serialized_ids + .as_ref() + .is_none_or(|allowed| allowed.contains(&id)) + } + /// Returns the merged stats generation options for this compression site. pub fn merged_stats_options(&self) -> GenerateStatsOptions { self.merged_stats_options From b652b40211e80df972f73afe03b8454b3eadae5a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 14:25:55 -0400 Subject: [PATCH 09/12] add predecessor schemes and make produced_encodings refer to serialized ids Signed-off-by: Matt Katz --- .../src/{builder.rs => builder/mod.rs} | 217 +++----------- vortex-btrblocks/src/builder/tests.rs | 216 ++++++++++++++ .../src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- .../src/schemes/string/zstd_buffers.rs | 2 +- vortex-compressor/src/compressor/cascade.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 98 +++++-- vortex-compressor/src/compressor/select.rs | 21 +- vortex-compressor/src/compressor/tests.rs | 171 ++--------- .../src/compressor/version_tests.rs | 266 ++++++++++++++++++ vortex-compressor/src/scheme/ctx.rs | 27 +- vortex-compressor/src/scheme/exclusion.rs | 6 +- vortex-compressor/src/scheme/mod.rs | 22 +- vortex-file/src/tests.rs | 35 +++ 15 files changed, 686 insertions(+), 403 deletions(-) rename vortex-btrblocks/src/{builder.rs => builder/mod.rs} (58%) create mode 100644 vortex-btrblocks/src/builder/tests.rs create mode 100644 vortex-compressor/src/compressor/version_tests.rs diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder/mod.rs similarity index 58% rename from vortex-btrblocks/src/builder.rs rename to vortex-btrblocks/src/builder/mod.rs index fcf2bb3934b..9744e4b18e3 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. @@ -116,7 +116,7 @@ impl BtrBlocksCompressorBuilder { /// 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 /// @@ -197,32 +197,53 @@ 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 } /// Restricts compression to the serialized IDs in `allowed`, intersecting with any earlier /// call. /// - /// A scheme stays when at least one of its [produced IDs](Scheme::produced_encodings) is - /// permitted, and the compressor is handed the set so a scheme whose encoding has several - /// wire formats picks its compression mode from it: the newest permitted one. - /// - /// The file writer passes the serialized IDs its enabled editions permit. + /// At build time, each scheme is replaced by the newest version in its predecessor chain + /// whose [`required_serialized_ids`](Scheme::required_serialized_ids) 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.schemes - .retain(|s| s.produced_encodings().iter().any(|id| allowed.contains(id))); 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 { let compressor = CascadingCompressor::new(self.schemes); BtrBlocksCompressor(match self.allowed_serialized_ids { @@ -233,178 +254,4 @@ impl BtrBlocksCompressorBuilder { } #[cfg(test)] -mod tests { - use vortex_array::ArrayRef; - use vortex_array::Canonical; - use vortex_array::ExecutionCtx; - use vortex_array::VTable; - use vortex_array::arrays::Bool; - use vortex_array::arrays::Primitive; - use vortex_compressor::scheme::CompressionEstimate; - use vortex_compressor::scheme::EstimateVerdict; - use vortex_error::VortexResult; - use vortex_fastlanes::FoR; - - use super::*; - use crate::ArrayAndStats; - use crate::CompressorContext; - - #[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 allow_serialized_ids_filters_schemes() { - let allowed: HashSet = [FoR.id()].into_iter().collect(); - let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); - assert_eq!(builder.schemes.len(), 1); - assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); - - let none = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&HashSet::new()); - assert!(none.schemes.is_empty()); - } - - #[test] - fn allowing_all_declared_outputs_keeps_every_scheme() { - let allowed: HashSet = ALL_SCHEMES - .iter() - .flat_map(|scheme| scheme.produced_encodings()) - .collect(); - let builder = BtrBlocksCompressorBuilder::default().allow_serialized_ids(&allowed); - assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); - } - - /// Stands in for a scheme whose encoding has two wire formats. - #[derive(Debug)] - struct TwoFormatScheme; - - impl Scheme for TwoFormatScheme { - fn scheme_name(&self) -> &'static str { - "test.two_formats" - } - - fn matches(&self, _canonical: &Canonical) -> bool { - false - } - - fn produced_encodings(&self) -> Vec { - vec![FoR.id(), Bool.id()] - } - - 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 { - unreachable!("test helper never matches") - } - } - - /// A scheme with several wire formats stays while any of them is permitted; which one it - /// produces is decided when compressing. - #[test] - fn any_permitted_format_keeps_the_scheme() { - static TWO_FORMATS: TwoFormatScheme = TwoFormatScheme; - - let newer_only = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&TWO_FORMATS) - .allow_serialized_ids(&HashSet::from([Bool.id()])); - assert_eq!(newer_only.schemes.len(), 1); - - let neither = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&TWO_FORMATS) - .allow_serialized_ids(&HashSet::from([Primitive.id()])); - assert!(neither.schemes.is_empty()); - } - - #[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" - ); - } - } - - /// Every serialized ID is allowed until the writer narrows the set to its editions. - #[test] - fn allowed_serialized_ids_reach_the_compressor() { - let default = BtrBlocksCompressorBuilder::default().build(); - assert!(default.0.allowed_serialized_ids().is_none()); - - let narrowed = BtrBlocksCompressorBuilder::default() - .allow_serialized_ids(&HashSet::from([FoR.id()])) - .build(); - assert_eq!( - narrowed.0.allowed_serialized_ids(), - Some(&HashSet::from([FoR.id()])) - ); - } - - #[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..1b736ba895b --- /dev/null +++ b/vortex-btrblocks/src/builder/tests.rs @@ -0,0 +1,216 @@ +// 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(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(integer::FoRScheme.id())); + assert!(!compressor.has_scheme(NewFoRScheme.id())); +} diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 3f06d65b061..5204e4e2478 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 46b2f1e302e..86b69bba47e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index a1bc8643775..dddaa349de1 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index cf691c70fcb..98d5feee2d9 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn required_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index dd98f4ea3c6..ecfd3c2c542 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -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 159842f2595..03733aca5e1 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,9 +9,8 @@ mod sample; mod select; mod structural; -use std::sync::Arc; - use vortex_array::ArrayId; +use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; use crate::builtins::IntDictScheme; @@ -53,16 +52,38 @@ pub struct CascadingCompressor { /// list offsets). root_exclusions: Vec, - /// The serialized IDs the writer may emit, or `None` for no restriction. Seeds every root - /// [`CompressorContext`], where schemes read it. - allowed_serialized_ids: Option>>, + /// 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`](Self::has_scheme) 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 { @@ -73,40 +94,68 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + scheme_aliases, allowed_serialized_ids: None, } } - /// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier - /// call. + /// Selects the newest eligible version of each scheme, intersecting with any earlier call. /// - /// The file writer passes the serialized IDs its enabled editions permit. Schemes read the - /// set through [`CompressorContext::allows_serialized_id`], so a scheme whose encoding has - /// several wire formats picks the newest permitted one as its mode, while estimating and - /// while compressing alike. + /// A version is eligible only when all of its [`Scheme::required_serialized_ids`] 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 { - self.allowed_serialized_ids = Some(Arc::new(match self.allowed_serialized_ids.take() { + 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 serialized IDs the writer may emit, or `None` when unrestricted. - pub fn allowed_serialized_ids(&self) -> Option<&HashSet> { - self.allowed_serialized_ids.as_deref() - } - /// The context a compress call starts from. pub(crate) fn root_context(&self) -> CompressorContext { - CompressorContext::new(self.allowed_serialized_ids.clone()) + CompressorContext::new() } - /// Returns whether the compressor was configured with `scheme`. + /// Returns whether a version of `scheme` is enabled. + /// + /// Any ID in a registered predecessor chain refers to the selected version, including when + /// the selected version is older or newer than the specified ID. pub fn has_scheme(&self, scheme: SchemeId) -> bool { - self.schemes - .iter() - .any(|candidate| candidate.id() == scheme) + self.scheme_aliases.contains_key(&scheme) + } + + fn resolve_scheme_id(&self, scheme: SchemeId) -> SchemeId { + self.scheme_aliases.get(&scheme).copied().unwrap_or(scheme) } } @@ -114,3 +163,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/tests.rs b/vortex-compressor/src/compressor/tests.rs index ef388c7d373..3a2a6281047 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -9,14 +9,11 @@ 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::arrays::Bool; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; -use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -29,7 +26,6 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_session::VortexSession; -use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -100,48 +96,6 @@ impl Scheme for DirectRatioScheme { } } -/// What the last `FormatRecordingScheme::compress` call saw for `allows_serialized_id`. -static SEEN_FORMAT: Mutex> = Mutex::new(None); - -/// Stands in for a scheme whose encoding has several wire formats: it asks the compressor whether -/// the newer one is allowed and records the answer. -#[derive(Debug)] -struct FormatRecordingScheme; - -impl Scheme for FormatRecordingScheme { - fn scheme_name(&self) -> &'static str { - "test.format_recording" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn produced_encodings(&self) -> Vec { - Vec::new() - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - *SEEN_FORMAT.lock() = Some(compress_ctx.allows_serialized_id(Constant.id())); - Ok(data.array().clone()) - } -} - #[derive(Debug)] struct ImmediateAlwaysUseScheme; @@ -420,12 +374,8 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -442,12 +392,8 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -464,12 +410,8 @@ fn callback_skip_is_ignored() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -486,12 +428,8 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -508,12 +446,8 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -530,12 +464,8 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -552,12 +482,8 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(winner.is_none()); Ok(()) @@ -658,12 +584,8 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -683,7 +605,7 @@ fn threshold_reflects_pass_one_best() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -704,7 +626,7 @@ fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner // `None`) because the zero-byte sample is never stored as the best. @@ -723,7 +645,7 @@ fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert_eq!(observed, Some(None)); @@ -743,7 +665,7 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?; + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; let observed = *OBSERVED_THRESHOLD.lock(); assert!(matches!( @@ -764,12 +686,8 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( let data = estimate_test_data(); let mut exec_ctx = SESSION.create_execution_ctx(); - let winner = compressor.choose_best_scheme( - &schemes, - &data, - CompressorContext::new(None), - &mut exec_ctx, - )?; + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; assert!(matches!( winner, @@ -816,7 +734,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // A context with default stats_options (count_distinct_values = false) and // marked as a sample so the function skips the sampling step and compresses // the array directly. - let ctx = CompressorContext::new(None).with_sampling(); + let ctx = CompressorContext::new().with_sampling(); // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" @@ -927,46 +845,3 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } - -#[test] -fn allowed_serialized_ids_default_to_everything_and_intersect() { - let compressor = compressor(); - let root = compressor.root_context(); - assert!(root.allows_serialized_id(Constant.id())); - assert!(root.allows_serialized_id(Bool.id())); - - let restricted = - compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()])); - let root = restricted.root_context(); - assert!(root.allows_serialized_id(Constant.id())); - assert!(!root.allows_serialized_id(Bool.id())); - - let narrowed = - restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()])); - let root = narrowed.root_context(); - assert!(root.allows_serialized_id(Primitive.id())); - assert!(!root.allows_serialized_id(Constant.id())); - assert!(!root.allows_serialized_id(Bool.id())); - - // Descending keeps the set. - let child = root.descend_with_scheme(IntDictScheme.id(), 0); - assert!(child.allows_serialized_id(Primitive.id())); - assert!(!child.allows_serialized_id(Constant.id())); -} - -/// A scheme sees the restriction through the compressor it is handed: everything is allowed until -/// the writer narrows the set to its editions. -#[test] -fn schemes_see_the_allowed_serialized_ids() -> VortexResult<()> { - let array = PrimitiveArray::from_iter(0..4096i32).into_array(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let unrestricted = CascadingCompressor::new(vec![&FormatRecordingScheme]); - unrestricted.compress(&array, &mut exec_ctx)?; - assert_eq!(*SEEN_FORMAT.lock(), Some(true)); - - let restricted = unrestricted.with_allowed_serialized_ids(HashSet::from([Primitive.id()])); - restricted.compress(&array, &mut exec_ctx)?; - assert_eq!(*SEEN_FORMAT.lock(), Some(false)); - Ok(()) -} diff --git a/vortex-compressor/src/compressor/version_tests.rs b/vortex-compressor/src/compressor/version_tests.rs new file mode 100644 index 00000000000..197583331cf --- /dev/null +++ b/vortex-compressor/src/compressor/version_tests.rs @@ -0,0 +1,266 @@ +// 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(version.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(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(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 83685031c33..0b9d8e3d4b8 100644 --- a/vortex-compressor/src/scheme/ctx.rs +++ b/vortex-compressor/src/scheme/ctx.rs @@ -4,11 +4,8 @@ //! Compression context for recursive compression. use std::fmt; -use std::sync::Arc; -use vortex_array::ArrayId; use vortex_error::VortexExpect; -use vortex_utils::aliases::hash_set::HashSet; use crate::compressor::ROOT_SCHEME_ID; use crate::scheme::SchemeId; @@ -41,24 +38,18 @@ pub struct CompressorContext { /// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions /// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions cascade_history: Vec<(SchemeId, usize)>, - - /// The serialized IDs the writer may emit, or `None` for no restriction. Shared by every - /// context of one compress call, so cloning at each descent is a pointer bump. - allowed_serialized_ids: Option>>, } impl CompressorContext { - /// Creates a new root `CompressorContext` for a compressor that may emit the given serialized - /// IDs, or any ID when `None`. + /// Creates a new root `CompressorContext`. /// /// This should **only** be created by the compressor. - pub(crate) fn new(allowed_serialized_ids: Option>>) -> Self { + pub(crate) fn new() -> Self { Self { is_sample: false, allowed_cascading: MAX_CASCADE, merged_stats_options: GenerateStatsOptions::default(), cascade_history: Vec::new(), - allowed_serialized_ids, } } } @@ -66,7 +57,7 @@ impl CompressorContext { #[cfg(test)] impl Default for CompressorContext { fn default() -> Self { - Self::new(None) + Self::new() } } @@ -76,18 +67,6 @@ impl CompressorContext { self.is_sample } - /// Returns whether the writer may emit the serialized ID `id`. - /// - /// A scheme whose encoding has several wire formats picks its compression mode from this, - /// the newest permitted one, and the same answer is available while estimating and while - /// compressing. Without a restriction every ID is allowed. The serializer still emits the - /// oldest wire form the resulting array fits, and the serialization context validates it. - pub fn allows_serialized_id(&self, id: ArrayId) -> bool { - self.allowed_serialized_ids - .as_ref() - .is_none_or(|allowed| allowed.contains(&id)) - } - /// Returns the merged stats generation options for this compression site. pub fn merged_stats_options(&self) -> GenerateStatsOptions { self.merged_stats_options 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 fa42231d422..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 may write its output under. + /// The serialized IDs this scheme may write its output under. Every ID must be permitted before + /// this scheme can be selected. /// /// 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. /// - /// An encoding with several wire formats lists every one of them, oldest first. The writer - /// keeps the scheme while any of them is permitted, and the scheme picks the newest - /// permitted one as its compression mode. + /// 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-file/src/tests.rs b/vortex-file/src/tests.rs index 640de874d2b..f740aacb318 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -1729,6 +1729,41 @@ async fn test_encoding_registered_after_write_options() -> 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(); From b43711ae833b7c8081f030b4676fbf95bace618c Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 8 Sep 2026 14:40:19 -0400 Subject: [PATCH 10/12] Fix optional scheme trait methods and documentation links Restore produced_encodings in feature-gated schemes and fix stale trait links. Document scheme ID resolution for Clippy. Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder/mod.rs | 2 +- vortex-btrblocks/src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- vortex-btrblocks/src/schemes/string/zstd_buffers.rs | 2 +- vortex-compressor/src/compressor/mod.rs | 3 ++- 6 files changed, 7 insertions(+), 6 deletions(-) diff --git a/vortex-btrblocks/src/builder/mod.rs b/vortex-btrblocks/src/builder/mod.rs index 9744e4b18e3..7143038ea88 100644 --- a/vortex-btrblocks/src/builder/mod.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -227,7 +227,7 @@ impl BtrBlocksCompressorBuilder { /// call. /// /// At build time, each scheme is replaced by the newest version in its predecessor chain - /// whose [`required_serialized_ids`](Scheme::required_serialized_ids) are all permitted. + /// 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 { diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 5204e4e2478..3f06d65b061 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 86b69bba47e..46b2f1e302e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index dddaa349de1..a1bc8643775 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index 98d5feee2d9..cf691c70fcb 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn required_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 03733aca5e1..464768e8b36 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -101,7 +101,7 @@ impl CascadingCompressor { /// Selects the newest eligible version of each scheme, intersecting with any earlier call. /// - /// A version is eligible only when all of its [`Scheme::required_serialized_ids`] are allowed. + /// 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 { @@ -154,6 +154,7 @@ impl CascadingCompressor { self.scheme_aliases.contains_key(&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) } From 3a0d848bead22dbd03c19241732d916b6cfec5b3 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 9 Sep 2026 17:52:01 -0400 Subject: [PATCH 11/12] fix Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder/tests.rs | 7 +- vortex-btrblocks/tests/scheme_versions.rs | 103 ++++++++++++++++++ vortex-compressor/src/compressor/mod.rs | 17 ++- .../src/compressor/version_tests.rs | 10 +- 4 files changed, 127 insertions(+), 10 deletions(-) create mode 100644 vortex-btrblocks/tests/scheme_versions.rs diff --git a/vortex-btrblocks/src/builder/tests.rs b/vortex-btrblocks/src/builder/tests.rs index 1b736ba895b..4310b9416c8 100644 --- a/vortex-btrblocks/src/builder/tests.rs +++ b/vortex-btrblocks/src/builder/tests.rs @@ -200,7 +200,8 @@ fn restrictions_select_predecessors_of_schemes_added_later() { .with_new_scheme(&NewFoRScheme) .build(); assert!(compressor.has_scheme(integer::FoRScheme.id())); - assert!(compressor.has_scheme(NewFoRScheme.id())); + assert!(compressor.has_scheme_family(NewFoRScheme.id())); + assert!(!compressor.has_scheme(NewFoRScheme.id())); } #[rstest] @@ -211,6 +212,6 @@ fn excluding_any_version_removes_the_chain(#[case] excluded: SchemeId) { .with_new_scheme(&NewFoRScheme) .exclude_schemes([excluded]) .build(); - assert!(!compressor.has_scheme(integer::FoRScheme.id())); - assert!(!compressor.has_scheme(NewFoRScheme.id())); + assert!(!compressor.has_scheme_family(integer::FoRScheme.id())); + assert!(!compressor.has_scheme_family(NewFoRScheme.id())); } diff --git a/vortex-btrblocks/tests/scheme_versions.rs b/vortex-btrblocks/tests/scheme_versions.rs new file mode 100644 index 00000000000..ff4b67961c4 --- /dev/null +++ b/vortex-btrblocks/tests/scheme_versions.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![cfg(feature = "unstable_encodings")] + +#[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-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 464768e8b36..3f88debbc24 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -63,7 +63,7 @@ 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`](Self::has_scheme) checks. + /// 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 @@ -146,14 +146,21 @@ impl CascadingCompressor { CompressorContext::new() } - /// Returns whether a version of `scheme` is enabled. + /// Returns whether any version in the scheme's family is enabled. /// - /// Any ID in a registered predecessor chain refers to the selected version, including when - /// the selected version is older or newer than the specified ID. - pub fn has_scheme(&self, scheme: SchemeId) -> bool { + /// 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.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) diff --git a/vortex-compressor/src/compressor/version_tests.rs b/vortex-compressor/src/compressor/version_tests.rs index 197583331cf..1fc6dfbac99 100644 --- a/vortex-compressor/src/compressor/version_tests.rs +++ b/vortex-compressor/src/compressor/version_tests.rs @@ -152,8 +152,13 @@ fn newest_eligible_version_is_selected_before_estimation() -> VortexResult<()> { )?; assert_eq!(winner.map(|(scheme, _)| scheme.id()), Some(expected)); for version in [&V1, &V2, &V3] { - assert!(compressor.has_scheme(version.id())); + assert!(compressor.has_scheme_family(version.id())); + assert_eq!( + compressor.has_scheme(version.id()), + version.id() == expected + ); } + assert!(!compressor.has_scheme(OTHER.id())); } Ok(()) } @@ -164,6 +169,7 @@ fn no_eligible_version_removes_the_entire_chain() { 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())); } } @@ -193,7 +199,7 @@ fn successive_restrictions_keep_aliases_and_intersect_wire_ids() { let compressor = compressor.with_allowed_serialized_ids(HashSet::from([*V2_ID, *AUX_ID])); assert!(compressor.schemes.is_empty()); - assert!(!compressor.has_scheme(V3.id())); + assert!(!compressor.has_scheme_family(V3.id())); } static PUSH_OLD: TestScheme = TestScheme { From 29fe244c9037d011e7308b5f6564f6d0fb771a6a Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 16 Sep 2026 11:47:29 -0400 Subject: [PATCH 12/12] Adapt to develop: rename remaining scheme-filter callers and drop the unstable_encodings gate Signed-off-by: Matt Katz --- vortex-bench/src/lib.rs | 2 +- vortex-btrblocks/src/builder/mod.rs | 2 +- vortex-btrblocks/src/canonical_compressor.rs | 2 +- vortex-btrblocks/tests/golden.rs | 4 ++-- vortex-btrblocks/tests/scheme_versions.rs | 2 -- vortex-cuda/ffi/src/lib.rs | 2 +- vortex-cuda/gpu-scan-cli/src/main.rs | 2 +- vortex-python/src/io.rs | 2 +- vortex-tui/src/convert.rs | 2 +- 9 files changed, 9 insertions(+), 11 deletions(-) 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/mod.rs b/vortex-btrblocks/src/builder/mod.rs index 7143038ea88..3699964e702 100644 --- a/vortex-btrblocks/src/builder/mod.rs +++ b/vortex-btrblocks/src/builder/mod.rs @@ -161,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. 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/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 index ff4b67961c4..9c68725ae0a 100644 --- a/vortex-btrblocks/tests/scheme_versions.rs +++ b/vortex-btrblocks/tests/scheme_versions.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -#![cfg(feature = "unstable_encodings")] - #[cfg(test)] mod tests { use rstest::rstest; 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 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