diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 2d203323fef..5ef6a1f6ed7 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -3,6 +3,8 @@ //! Builder for configuring `BtrBlocksCompressor` instances. +use vortex_array::ArrayId; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; use vortex_session::VortexSession; use vortex_utils::aliases::hash_set::HashSet; @@ -85,6 +87,16 @@ impl CompressionMode { } excluded } + + /// Returns the serialized IDs [`build`](BtrBlocksCompressorBuilder::build) disallows in this + /// mode, on top of the session's restrictions. + fn excluded_encodings(self) -> Vec { + match self { + Self::All | Self::Default | Self::Compact => Vec::new(), + // Multi-part DecimalByteParts arrays have no CUDA decode kernel. + Self::Cuda => vec![decimal_byte_parts_v2_id()], + } + } } /// Builder for creating configured [`BtrBlocksCompressor`] instances. @@ -218,15 +230,14 @@ impl BtrBlocksCompressorBuilder { fn allowed_schemes(&self) -> Vec<&'static dyn Scheme> { let excluded: HashSet = self.mode.excluded_schemes().into_iter().collect(); + let excluded_ids: HashSet = self.mode.excluded_encodings().into_iter().collect(); + let allowed = |id: &ArrayId| self.allowed.is_allowed(id) && !excluded_ids.contains(id); self.schemes .iter() .copied() .filter(|s| !excluded.contains(&s.id())) - .filter(|s| { - s.produced_encodings() - .iter() - .all(|id| self.allowed.is_allowed(id)) - }) + .map(|s| s.refine(&allowed).unwrap_or(s)) + .filter(|s| s.produced_encodings().iter().all(&allowed)) .collect() } } diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index f77a77d8c50..ba2a5e26546 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -9,13 +9,15 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; use vortex_array::dtype::DecimalType; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::EstimateVerdict; use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsSlots; use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_decimal_byte_parts::split_decimal; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -24,12 +26,54 @@ use crate::CompressorContext; use crate::Scheme; use crate::SchemeExt; +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum DecimalSchemeMode { + V1, + V2, +} + +/// The v1 decimal scheme, which the default [`CompressionSession`](crate::CompressionSession) +/// registers. +pub(crate) static DECIMAL_V1: DecimalScheme = DecimalScheme::v1(); +static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); + /// Compression scheme for decimal arrays via byte-part decomposition. /// -/// Narrows the decimal to the smallest integer type, compresses the underlying primitive, and wraps -/// the result in a `DecimalBytePartsArray`. +/// Narrows the decimal to the smallest integer type and compresses its byte parts independently. +/// The v1 mode leaves values wider than `i64` canonical; v2 splits them into a signed most +/// significant part and up to three unsigned lower parts. Single-part arrays serialize as v1 +/// in either mode, while arrays with lower parts serialize as v2. +/// +/// The default uses v1. [`refine`](Scheme::refine) picks v2 when both decimal IDs are allowed +/// and v1 otherwise. #[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct DecimalScheme; +pub struct DecimalScheme { + mode: DecimalSchemeMode, +} + +impl DecimalScheme { + /// Creates a decimal scheme configured for v1, disallowing splitting of wide decimals. + /// + /// Values that remain wider than `i64` after narrowing stay canonical. + pub const fn v1() -> Self { + Self { + mode: DecimalSchemeMode::V1, + } + } + + /// Creates a decimal scheme configured for v2, allowing splitting of wide decimals. + pub const fn v2() -> Self { + Self { + mode: DecimalSchemeMode::V2, + } + } +} + +impl Default for DecimalScheme { + fn default() -> Self { + Self::v1() + } +} impl Scheme for DecimalScheme { fn scheme_name(&self) -> &'static str { @@ -41,14 +85,28 @@ impl Scheme for DecimalScheme { } fn produced_encodings(&self) -> Vec { - // 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()] + match self.mode { + DecimalSchemeMode::V1 => vec![decimal_byte_parts_v1_id()], + DecimalSchemeMode::V2 => { + vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()] + } + } + } + + fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> Option<&'static dyn Scheme> { + if allowed(&decimal_byte_parts_v1_id()) && allowed(&decimal_byte_parts_v2_id()) { + Some(&DECIMAL_V2) + } else { + Some(&DECIMAL_V1) + } } - /// Children: primitive=0. + /// Children: msp=0, then up to three lower parts in v2 mode. fn num_children(&self) -> usize { - 1 + match self.mode { + DecimalSchemeMode::V1 => 1, + DecimalSchemeMode::V2 => 4, + } } fn expected_compression_ratio( @@ -68,22 +126,73 @@ impl Scheme for DecimalScheme { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - // TODO(joe): add support splitting i128/256 buffers into chunks of primitive values - // for compression. 2 for i128 and 4 for i256. let decimal = data.array().clone().execute::(exec_ctx)?; let decimal = narrowed_decimal(decimal); - let validity = decimal.validity()?; - let prim = match decimal.values_type() { - DecimalType::I8 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I16 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I32 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I64 => PrimitiveArray::new(decimal.buffer::(), validity), - _ => return Ok(decimal.into_array()), - }; + if self.mode == DecimalSchemeMode::V1 + && matches!(decimal.values_type(), DecimalType::I128 | DecimalType::I256) + { + return Ok(decimal.into_array()); + } - let compressed = - compressor.compress_child(&prim.into_array(), &compress_ctx, self.id(), 0, exec_ctx)?; + let parts = split_decimal(&decimal, exec_ctx)?; + let msp = compressor.compress_child( + &parts.msp, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::MSP, + exec_ctx, + )?; + let lower_parts = parts + .lower_parts + .iter() + .enumerate() + .map(|(idx, part)| { + compressor.compress_child( + part, + &compress_ctx, + self.id(), + DecimalBytePartsSlots::LOWER_PARTS_OFFSET + idx, + exec_ctx, + ) + }) + .collect::>>()?; - DecimalByteParts::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array()) + DecimalByteParts::try_new_with_lower_parts(msp, lower_parts, decimal.decimal_dtype()) + .map(IntoArray::into_array) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayId; + use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; + use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; + + use super::DECIMAL_V1; + use super::DECIMAL_V2; + use crate::Scheme; + use crate::SchemeExt; + + /// Both variants refine to v2 exactly when both decimal IDs are allowed. + #[rstest] + #[case::neither(false, false, false)] + #[case::v1(true, false, false)] + #[case::v2_only(false, true, false)] + #[case::both(true, true, true)] + fn refine_picks_v2_only_when_both_ids_are_allowed( + #[case] allow_v1: bool, + #[case] allow_v2: bool, + #[case] expect_v2: bool, + #[values(&DECIMAL_V1, &DECIMAL_V2)] scheme: &'static dyn Scheme, + ) { + let allowed = |id: &ArrayId| { + (allow_v1 && *id == decimal_byte_parts_v1_id()) + || (allow_v2 && *id == decimal_byte_parts_v2_id()) + }; + let refined = scheme.refine(&allowed).unwrap_or(scheme); + assert_eq!(refined.id(), scheme.id()); + let expected: &dyn Scheme = if expect_v2 { &DECIMAL_V2 } else { &DECIMAL_V1 }; + assert_eq!(refined.produced_encodings(), expected.produced_encodings()); } } diff --git a/vortex-btrblocks/src/session.rs b/vortex-btrblocks/src/session.rs index 67dd1f278bc..87093daa95b 100644 --- a/vortex-btrblocks/src/session.rs +++ b/vortex-btrblocks/src/session.rs @@ -87,8 +87,8 @@ impl Default for CompressionSession { &binary::ZstdScheme, #[cfg(feature = "zstd")] &binary::ZstdBuffersScheme, - // Decimal schemes. - &decimal::DecimalScheme, + // Decimal schemes. `refine` switches to v2 where both decimal IDs are allowed. + &decimal::DECIMAL_V1, // Temporal schemes. &temporal::TemporalScheme, ], diff --git a/vortex-btrblocks/tests/decimal_config.rs b/vortex-btrblocks/tests/decimal_config.rs new file mode 100644 index 00000000000..4431fd04f95 --- /dev/null +++ b/vortex-btrblocks/tests/decimal_config.rs @@ -0,0 +1,229 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal scheme refinement by allowed serialized IDs, and compression of wide decimal parts. + +#![cfg(test)] + +use std::sync::LazyLock; + +use rstest::rstest; +use vortex_array::ArrayContext; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Decimal; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +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_btrblocks::BtrBlocksCompressorBuilder; +use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::schemes::decimal::DecimalScheme; +use vortex_btrblocks::schemes::integer::BitPackingScheme; +use vortex_btrblocks::schemes::integer::FoRScheme; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBufferMut; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use vortex_decimal_byte_parts::decimal_byte_parts_v1_id; +use vortex_decimal_byte_parts::decimal_byte_parts_v2_id; +use vortex_edition::EDITION_DECLARATIONS; +use vortex_edition::EDITION_FAMILIES; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +static DECIMAL_V2: DecimalScheme = DecimalScheme::v2(); + +/// Registers the decimal and fastlanes encodings, and enables no editions. +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + session +}); + +/// Like [`SESSION`], with the latest core edition enabled: it allows decimal v1 but not v2. +static CORE_SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session().with::(); + for family in EDITION_FAMILIES { + session + .editions() + .declare_family(family) + .expect("first-party edition family"); + } + for declaration in EDITION_DECLARATIONS { + session + .register_edition(declaration) + .expect("first-party edition"); + } + session + .enable_edition(CORE_2026_08_3) + .expect("core edition is registered"); + vortex_decimal_byte_parts::initialize(&session); + vortex_fastlanes::initialize(&session); + session +}); + +/// Which decimal output a builder is expected to produce. +#[derive(Debug, Clone, Copy)] +enum Expected { + /// No decimal scheme survives, so decimals stay canonical. + Canonical, + /// Single-part arrays only: wide decimals stay canonical. + V1, + /// Wide decimals split into multi-part arrays. + V2, +} + +fn decimal_array(wide: bool) -> ArrayRef { + let base = if wide { 1i128 << 70 } else { 0 }; + DecimalArray::new( + (0..128i128).map(|i| base + i).collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array() +} + +fn assert_decimal_output( + builder: BtrBlocksCompressorBuilder, + wide: bool, + expected_id: Option, +) -> VortexResult<()> { + let array = decimal_array(wide); + let mut ctx = SESSION.create_execution_ctx(); + let compressed = builder.build().compress(&array, &mut ctx)?; + if let Some(expected_id) = expected_id { + assert!(compressed.is::()); + let serialized = SESSION + .array_serialize(&compressed)? + .ok_or_else(|| vortex_err!("expected serializable decimal byte parts"))?; + assert_eq!(serialized.serialized_id, expected_id); + } else { + assert!(compressed.is::()); + } + assert_arrays_eq!(array, compressed, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::no_editions(|| BtrBlocksCompressorBuilder::from_session(&SESSION), Expected::Canonical)] +#[case::core_edition(|| BtrBlocksCompressorBuilder::from_session(&CORE_SESSION), Expected::V1)] +#[case::editions_disabled( + || BtrBlocksCompressorBuilder::from_session(&SESSION).disable_editions(), + Expected::V2 +)] +#[case::unrestricted( + || BtrBlocksCompressorBuilder::from_session(&SESSION).unrestricted(), + Expected::V2 +)] +#[case::cuda( + || BtrBlocksCompressorBuilder::from_session(&SESSION).unrestricted().only_cuda_compatible(), + Expected::V1 +)] +fn decimal_scheme_follows_allowed_ids( + #[case] builder: fn() -> BtrBlocksCompressorBuilder, + #[case] expected: Expected, + #[values(false, true)] wide: bool, +) -> VortexResult<()> { + let expected_id = match expected { + Expected::Canonical => None, + Expected::V1 | Expected::V2 if !wide => Some(decimal_byte_parts_v1_id()), + Expected::V1 => None, + Expected::V2 => Some(decimal_byte_parts_v2_id()), + }; + assert_decimal_output(builder(), wide, expected_id) +} + +#[test] +fn refinement_does_not_restore_excluded_decimal() -> VortexResult<()> { + let builder = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .exclude_schemes([DECIMAL_V2.id()]); + assert_decimal_output(builder, false, None) +} + +#[rstest] +#[case::i128(false, 1)] +#[case::i256(true, 3)] +fn wide_decimal_parts_roundtrip( + #[case] use_i256: bool, + #[case] lower_part_count: usize, + #[values(false, true)] negative: bool, + #[values(false, true)] nullable: bool, + #[values(false, true)] compress_children: bool, +) -> VortexResult<()> { + let validity = if nullable { + Validity::from_iter((0..2048).map(|i| i % 7 != 0)) + } else { + Validity::NonNullable + }; + let array = if use_i256 { + let values = (1..=2048u32) + .map(|i| { + let msp = if negative { + -i128::from(i) + } else { + i128::from(i) + }; + i256::from_parts( + (u128::from(i) << 64) | u128::from(i * 131 + 17), + (msp << 64) | i128::from(i * 3 + 1), + ) + }) + .collect::>(); + DecimalArray::new(values, DecimalDType::new(76, 2), validity) + } else { + let values = (1..=2048i128) + .map(|i| { + let msp = if negative { -i } else { i }; + (msp << 70) + i * 131 + 17 + }) + .collect::>(); + DecimalArray::new(values, DecimalDType::new(38, 2), validity) + } + .into_array(); + let mut builder = BtrBlocksCompressorBuilder::empty().with_new_scheme(&DECIMAL_V2); + if compress_children { + builder = builder + .with_new_scheme(&FoRScheme) + .with_new_scheme(&BitPackingScheme); + } + let mut ctx = SESSION.create_execution_ctx(); + let compressed = builder.build().compress(&array, &mut ctx)?; + // Every part varies, so without child compression splitting alone cannot save space. + assert_eq!(compressed.is::(), compress_children); + if compress_children { + let parts = compressed + .as_opt::() + .ok_or_else(|| vortex_err!("expected decimal byte parts"))?; + assert_eq!(parts.lower_parts().len(), lower_part_count); + assert!(!parts.msp().is_canonical()); + assert!(parts.lower_parts().iter().all(|part| !part.is_canonical())); + } + + let array_ctx = ArrayContext::empty(); + let mut bytes = ByteBufferMut::empty(); + for buffer in compressed.serialize(&array_ctx, &SESSION, &SerializeOptions::default())? { + bytes.extend_from_slice(buffer.as_ref()); + } + let decoded = SerializedArray::try_from(bytes.freeze())?.decode( + array.dtype(), + array.len(), + &ReadContext::new(array_ctx.to_ids()), + &SESSION, + )?; + assert_arrays_eq!(array, decoded, &mut ctx); + Ok(()) +} diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 0ba1c90202a..59fe48fc066 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -135,6 +135,15 @@ pub trait Scheme: Debug + Send + Sync { /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. fn produced_encodings(&self) -> Vec; + /// Returns the variant of this scheme to use given which serialized IDs are `allowed`. + /// + /// `None` keeps this scheme. A variant must share this scheme's [`SchemeId`]. Every ID the + /// variant declares in [`produced_encodings`](Self::produced_encodings) must still be allowed + /// for it to be used. + fn refine(&self, _allowed: &dyn Fn(&ArrayId) -> bool) -> 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 81ade17fd1e..6f74e583b7c 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -38,6 +38,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; +use vortex_array::dtype::i256; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::and; @@ -72,7 +73,12 @@ use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; +use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_decimal_byte_parts::DecimalBytePartsArraySlotsExt; +use vortex_edition::EDITION_DECLARATIONS; use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::declarations::core::CORE_2026_08_3; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_io::session::RuntimeSession; @@ -173,6 +179,90 @@ async fn test_read_simple() { assert_eq!(row_count, 8); } +/// Wide decimals split into multi-part arrays only when the writer allows the v2 format. The +/// default writer allows every registered encoding once editions are disabled; a strategy built +/// from the session keeps the enabled editions' restrictions regardless. +#[rstest] +#[case::default_writer(false, false)] +#[case::custom_layout(true, false)] +#[case::explicit_compressor(true, true)] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn decimal_writer_uses_default_or_session_restrictions( + #[values(false, true)] use_i256: bool, + #[case] custom_strategy: bool, + #[case] explicit_compressor: bool, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::() + .with::(); + crate::register_default_encodings(&session); + for declaration in EDITION_DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(CORE_2026_08_3)?; + + let array = if use_i256 { + DecimalArray::new( + (0..1024u128) + .map(|i| i256::from_parts(i * 17, 1i128 << 70)) + .collect::>(), + DecimalDType::new(76, 2), + Validity::NonNullable, + ) + } else { + DecimalArray::new( + (0..1024i128) + .map(|i| (1i128 << 70) + i * 17) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + } + .into_array(); + let strategy = crate::strategy::WriteStrategyBuilder::from_session(&session) + .with_row_block_size(256) + .with_data_block_target_bytes(None); + let strategy = if explicit_compressor { + strategy.with_btrblocks_builder(BtrBlocksCompressorBuilder::from_session(&session)) + } else { + strategy + } + .build(); + + for disable_editions in [false, true] { + let mut options = session.write_options(); + if custom_strategy { + options = options.with_strategy(Arc::clone(&strategy)); + } + let options = if disable_editions { + options.disable_editions() + } else { + options + }; + let mut buffer = ByteBufferMut::empty(); + options + .write(&mut buffer, array.clone().to_array_stream()) + .await?; + let actual = session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + let uses_v2 = actual.depth_first_traversal().any(|array| { + array + .as_opt::() + .is_some_and(|parts| !parts.lower_parts().is_empty()) + }); + assert_eq!(uses_v2, disable_editions && !custom_strategy); + assert_arrays_eq!(array, actual, &mut session.create_execution_ctx()); + } + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_round_trip_many_types() {