diff --git a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs index b9babd52aaa..a49c9c87f17 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -5,11 +5,13 @@ use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hasher; +use prost::Message as _; use vortex_array::Array; use vortex_array::ArrayParts; use vortex_array::ArrayView; pub(crate) mod compute; mod limbs; +mod plugin; mod rules; #[cfg(test)] pub(crate) mod testing; @@ -22,7 +24,8 @@ pub mod _benchmarking { pub use super::limbs::assemble_decimal; } -use prost::Message as _; +pub use plugin::DecimalBytePartsPlugin; +pub use plugin::decimal_byte_parts_v2_id; use vortex_array::ArrayEq; use vortex_array::ArrayHash; use vortex_array::ArrayId; @@ -177,7 +180,9 @@ impl DecimalByteParts { lower_parts: Vec, decimal_dtype: DecimalDType, ) -> VortexResult { - // Lower parts are supported in memory; the frozen serializer still rejects them. + // Building lower parts in memory is never gated — reading a file requires it. What is + // gated is the serialized form: an array carrying lower parts serializes under the + // `vortex.decimal_byte_parts_v2` format ID, which only editions that contain it may write. let len = msp.len(); let dtype = DType::Decimal(decimal_dtype, msp.dtype().nullability()); let slots = DecimalBytePartsSlots { msp, lower_parts }.into_slots(); @@ -256,7 +261,7 @@ impl VTable for DecimalByteParts { ) -> VortexResult>> { vortex_ensure!( array.lower_parts().is_empty(), - "serializing DecimalByteParts with lower parts is not supported" + "serializing DecimalByteParts with lower parts requires DecimalBytePartsPlugin" ); Ok(Some( DecimalBytesPartsMetadata::from_array(array)?.encode_to_vec(), @@ -506,6 +511,8 @@ mod tests { use crate::decimal_byte_parts::testing::i128_parts; use crate::decimal_byte_parts::testing::i256_of; use crate::decimal_byte_parts::testing::i256_parts; + use crate::decimal_byte_parts::testing::wide_i128_values; + use crate::decimal_byte_parts::testing::wide_i256_values; #[test] fn test_scalar_at_decimal_parts() { @@ -546,47 +553,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))] @@ -817,11 +783,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/plugin.rs b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin.rs new file mode 100644 index 00000000000..e7fb773e049 --- /dev/null +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/plugin.rs @@ -0,0 +1,351 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serialization of decimal byte parts under the frozen and v2 format IDs. + +use prost::Message as _; +use vortex_array::Array; +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::VTable; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use super::DecimalByteParts; +use super::DecimalBytePartsArraySlotsExt; +use super::DecimalBytesPartsMetadata; + +/// The `vortex.decimal_byte_parts_v2` serialized format ID, for `DecimalBytePartsArray`s carrying +/// lower parts. +/// +/// The `vortex.decimal_byte_parts` Id corresponds to the previous version of the `DecimalBytePartsArray`, +/// which does not support lower parts. Both IDs deserialize back into the same `DecimalBytePartsArray`. +pub fn decimal_byte_parts_v2_id() -> ArrayId { + static ID: CachedId = CachedId::new("vortex.decimal_byte_parts_v2"); + *ID +} + +/// The [`ArrayPlugin`] for [`DecimalByteParts`], owning both of its serialized formats. +/// +/// An array without lower parts serializes under the frozen `vortex.decimal_byte_parts` ID, +/// byte-identical to files written before lower parts existed. An array carrying lower parts +/// serializes under [`decimal_byte_parts_v2_id`]. Reading holds each ID to its own contract: +/// the frozen ID carries no lower parts and the v2 ID carries at least one. +/// +/// Register this plugin, or call [`crate::initialize`], to enable both formats. Registering +/// [`DecimalByteParts`] directly only supports the frozen format. +#[derive(Clone, Debug)] +pub struct DecimalBytePartsPlugin; + +impl ArrayPlugin for DecimalBytePartsPlugin { + fn id(&self) -> ArrayId { + VTable::id(&DecimalByteParts) + } + + fn serialized_ids(&self) -> Vec { + vec![VTable::id(&DecimalByteParts), 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_id = if view.lower_parts().is_empty() { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + Ok(Some(ArraySerialization::from_array( + serialized_id, + array, + DecimalBytesPartsMetadata::from_array(view)?.encode_to_vec(), + ))) + } + + fn deserialize( + &self, + parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + let metadata = DecimalBytesPartsMetadata::decode(parts.metadata)?; + let lower_part_count = metadata.lower_part_count()?; + if parts.serialized_id == decimal_byte_parts_v2_id() { + vortex_ensure!( + lower_part_count > 0, + "{} must carry at least one lower part", + parts.serialized_id + ); + } else { + vortex_ensure!( + parts.serialized_id == VTable::id(&DecimalByteParts), + "DecimalByteParts plugin does not recognize serialized ID {}", + parts.serialized_id + ); + vortex_ensure!( + lower_part_count == 0, + "{} must not carry lower parts, got {lower_part_count}", + parts.serialized_id + ); + } + Ok(Array::try_from_parts(metadata.into_array_parts( + parts.dtype, + parts.len, + parts.children, + )?)? + .into_array()) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayContext; + use vortex_array::ArrayVTable; + use vortex_array::VortexSessionExecute; + use vortex_array::array_session; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::Primitive; + 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::decimal_byte_parts::testing::encode; + use crate::decimal_byte_parts::testing::i128_parts; + use crate::decimal_byte_parts::testing::i256_parts; + use crate::decimal_byte_parts::testing::wide_i128_values; + use crate::decimal_byte_parts::testing::wide_i256_values; + + #[rstest] + #[case::no_lower_parts(DecimalByteParts::try_new( + buffer![1i32, 2, 3].into_array(), DecimalDType::new(9, 2), + ))] + #[case::one_lower_part(Ok(i128_parts(wide_i128_values(), Validity::NonNullable)))] + #[case::three_lower_parts(Ok(i256_parts(wide_i256_values(), Validity::NonNullable)))] + #[case::nullable_three_lower_parts(Ok(i256_parts( + wide_i256_values(), + Validity::from_iter([true, false, true, true, true, false, true, true, true, true]), + )))] + #[case::wider_i64_storage(encode(&DecimalArray::new( + buffer![-99i64, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, + )))] + #[case::wider_i128_storage(encode(&DecimalArray::new( + buffer![-99i128, 0, 99], DecimalDType::new(2, 0), Validity::NonNullable, + )))] + #[case::wider_i256_storage(encode(&DecimalArray::new( + buffer![i256::from_i128(-99), i256::ZERO, i256::from_i128(99)], + DecimalDType::new(2, 0), Validity::NonNullable, + )))] + #[case::redundant_two_lower_parts(DecimalByteParts::try_new_with_lower_parts( + buffer![0i64; 3].into_array(), + vec![buffer![0u64; 3].into_array(), lower_part()], + DecimalDType::new(38, 2), + ))] + #[case::redundant_three_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), + ))] + fn test_serde_round_trip( + #[case] array: VortexResult, + ) -> VortexResult<()> { + let session = session(); + let array = array?; + let lower_part_count = array.lower_parts().len(); + let array = array.into_array(); + let dtype = array.dtype().clone(); + let len = array.len(); + + let expected_id = if lower_part_count == 0 { + VTable::id(&DecimalByteParts) + } 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() + .len(), + lower_part_count, + "lower parts must survive serde" + ); + + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) + } + + #[rstest] + #[case::missing_lower_part(1, vec![msp()])] + #[case::extra_lower_part(0, vec![msp(), lower_part()])] + #[case::too_many_lower_parts( + 4, vec![msp(), lower_part(), lower_part(), lower_part(), lower_part()], + )] + fn test_deserialize_rejects_child_count_mismatch( + #[case] lower_part_count: u32, + #[case] children: Vec, + ) { + let serialized_id = if lower_part_count == 0 { + VTable::id(&DecimalByteParts) + } else { + decimal_byte_parts_v2_id() + }; + assert!(plugin_deserialize_with(serialized_id, lower_part_count, children).is_err()); + } + + fn plugin_deserialize_with( + serialized_id: ArrayId, + lower_part_count: u32, + children: Vec, + ) -> VortexResult { + let metadata = DecimalBytesPartsMetadata { + zeroth_child_ptype: PType::I64 as i32, + lower_part_count, + } + .encode_to_vec(); + let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); + DecimalBytePartsPlugin.deserialize( + ArrayDeserialization::new(serialized_id, &dtype, 3, &metadata, &[], &children), + &array_session(), + ) + } + + /// Each serialized ID keeps its own contract: the frozen ID never carries lower parts, and + /// the v2 ID is never written without them. + #[rstest] + #[case::frozen_without_lower_parts(VTable::id(&DecimalByteParts), 0, vec![msp()], true)] + #[case::frozen_with_lower_parts( + VTable::id(&DecimalByteParts), + 1, + vec![msp(), lower_part()], + false + )] + #[case::v2_with_lower_parts(decimal_byte_parts_v2_id(), 1, vec![msp(), lower_part()], true)] + #[case::v2_without_lower_parts(decimal_byte_parts_v2_id(), 0, vec![msp()], false)] + fn plugin_holds_each_id_to_its_contract( + #[case] serialized_id: ArrayId, + #[case] lower_part_count: u32, + #[case] children: Vec, + #[case] accepted: bool, + ) { + let result = plugin_deserialize_with(serialized_id, lower_part_count, children); + assert_eq!(result.is_ok(), accepted, "{serialized_id}: {result:?}"); + } + + fn msp() -> ArrayRef { + buffer![1i64, 2, 3].into_array() + } + + fn lower_part() -> ArrayRef { + buffer![1u64, 2, 3].into_array() + } + + fn session() -> VortexSession { + let session = array_session(); + crate::initialize(&session); + session + } + + #[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( + [ + ArrayVTable::id(&DecimalByteParts), + 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( + [ + ArrayVTable::id(&DecimalByteParts), + 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_wide_serialization() -> VortexResult<()> { + let session = array_session(); + session.arrays().register(DecimalByteParts); + let array = DecimalByteParts::try_new_with_lower_parts( + msp(), + vec![lower_part()], + DecimalDType::new(38, 2), + )? + .into_array(); + assert!(session.array_serialize(&array).is_err()); + 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 index d2ce68f3700..950c963f7a2 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/testing.rs @@ -51,3 +51,44 @@ pub(crate) fn i256_parts(values: Vec, validity: Validity) -> DecimalBytePa pub(crate) fn i256_of(high: i128, low: u128) -> i256 { i256::from_parts(low, high) } + +/// 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. +pub(crate) 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`. +pub(crate) 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(), + ] +} 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-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),