From 78415c2b7eb304cc975ea8de8a6305738534b147 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Thu, 17 Sep 2026 11:47:30 -0400 Subject: [PATCH] Compress wide decimals into lower parts when the writer permits the v2 format Signed-off-by: Matt Katz --- Cargo.lock | 1 + .../src/decimal_byte_parts/mod.rs | 2 +- vortex-btrblocks/src/schemes/decimal.rs | 89 ------ vortex-btrblocks/src/schemes/decimal/mod.rs | 116 ++++++++ vortex-btrblocks/src/schemes/decimal/tests.rs | 257 ++++++++++++++++++ vortex-file/Cargo.toml | 1 + vortex-file/src/tests.rs | 156 +++++++++++ 7 files changed, 532 insertions(+), 90 deletions(-) delete mode 100644 vortex-btrblocks/src/schemes/decimal.rs create mode 100644 vortex-btrblocks/src/schemes/decimal/mod.rs create mode 100644 vortex-btrblocks/src/schemes/decimal/tests.rs diff --git a/Cargo.lock b/Cargo.lock index ca9cd626f1b..9d40c25d641 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11198,6 +11198,7 @@ dependencies = [ "object_store", "parking_lot", "pin-project-lite", + "rand 0.10.2", "rstest", "tokio", "tracing", 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 65bb8222f7a..733360589f4 100644 --- a/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs +++ b/encodings/decimal-byte-parts/src/decimal_byte_parts/mod.rs @@ -54,7 +54,7 @@ const MAX_I128_LOWER_PARTS: usize = 1; const MAX_I256_LOWER_PARTS: usize = 3; /// The maximum number of 64-bit lower parts an encoded decimal can carry. -const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS; +pub const MAX_LOWER_PARTS: usize = MAX_I256_LOWER_PARTS; /// Number of bits stored in each lower part. const LOWER_PART_BITS: usize = 64; diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs deleted file mode 100644 index f77a77d8c50..00000000000 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ /dev/null @@ -1,89 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Decimal compression scheme using byte-part decomposition. - -use vortex_array::ArrayId; -use vortex_array::ArrayRef; -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::decimal_byte_parts_v1_id; -use vortex_error::VortexResult; - -use crate::ArrayAndStats; -use crate::CascadingCompressor; -use crate::CompressorContext; -use crate::Scheme; -use crate::SchemeExt; - -/// 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`. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct DecimalScheme; - -impl Scheme for DecimalScheme { - fn scheme_name(&self) -> &'static str { - "vortex.decimal.byte_parts" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches!(canonical, Canonical::Decimal(_)) - } - - 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()] - } - - /// Children: primitive=0. - fn num_children(&self) -> usize { - 1 - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - // Decimal compression is almost always beneficial (narrowing + primitive compression). - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) - } - - fn compress( - &self, - compressor: &CascadingCompressor, - data: &ArrayAndStats, - compress_ctx: CompressorContext, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - // TODO(joe): add support splitting i128/256 buffers into chunks of primitive values - // for compression. 2 for i128 and 4 for i256. - let decimal = data.array().clone().execute::(exec_ctx)?; - let decimal = narrowed_decimal(decimal); - let validity = decimal.validity()?; - let prim = match decimal.values_type() { - DecimalType::I8 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I16 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I32 => PrimitiveArray::new(decimal.buffer::(), validity), - DecimalType::I64 => PrimitiveArray::new(decimal.buffer::(), validity), - _ => return Ok(decimal.into_array()), - }; - - let compressed = - compressor.compress_child(&prim.into_array(), &compress_ctx, self.id(), 0, exec_ctx)?; - - DecimalByteParts::try_new(compressed, decimal.decimal_dtype()).map(|d| d.into_array()) - } -} diff --git a/vortex-btrblocks/src/schemes/decimal/mod.rs b/vortex-btrblocks/src/schemes/decimal/mod.rs new file mode 100644 index 00000000000..87d577f2a75 --- /dev/null +++ b/vortex-btrblocks/src/schemes/decimal/mod.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal compression via byte-part decomposition. + +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::decimal::narrowed_decimal; +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::MAX_LOWER_PARTS; +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; +use crate::CascadingCompressor; +use crate::CompressorContext; +use crate::Scheme; +use crate::SchemeExt; + +/// Compression scheme for decimal arrays via byte-part decomposition. +/// +/// Narrows the decimal to the smallest integer type and splits it into a signed most significant +/// part plus up to three unsigned 64-bit lower parts, each compressed as its own child. Values that +/// fit one signed part produce a single-part array under the frozen `vortex.decimal_byte_parts` +/// format. Wider values need lower parts, and so the `vortex.decimal_byte_parts.v2` format. They +/// are split only when the writer may emit that format, and stay canonical otherwise. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct DecimalScheme; + +impl Scheme for DecimalScheme { + fn scheme_name(&self) -> &'static str { + "vortex.decimal.byte_parts" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches!(canonical, Canonical::Decimal(_)) + } + + fn produced_encodings(&self) -> Vec { + // Single-part arrays are always possible. The v2 format is written only when the + // compression context permits it, see `compress`. + vec![decimal_byte_parts_v1_id()] + } + + /// Children: msp=0, then up to [`MAX_LOWER_PARTS`] lower parts. + fn num_children(&self) -> usize { + DecimalBytePartsSlots::FIXED_COUNT + MAX_LOWER_PARTS + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + // Decimal compression is almost always beneficial (narrowing + primitive compression). + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + compressor: &CascadingCompressor, + data: &ArrayAndStats, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let decimal = data.array().clone().execute::(exec_ctx)?; + let decimal = narrowed_decimal(decimal); + let parts = split_decimal(&decimal, exec_ctx)?; + + // Lower parts need the v2 format. Leave wide values canonical when the writer may not + // emit it, so a frozen-format file never carries an array it cannot serialize. + if !parts.lower_parts.is_empty() + && !compress_ctx.allows_serialized_id(&decimal_byte_parts_v2_id()) + { + return Ok(decimal.into_array()); + } + + 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_with_lower_parts(msp, lower_parts, decimal.decimal_dtype()) + .map(IntoArray::into_array) + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-btrblocks/src/schemes/decimal/tests.rs b/vortex-btrblocks/src/schemes/decimal/tests.rs new file mode 100644 index 00000000000..39f0d39dd00 --- /dev/null +++ b/vortex-btrblocks/src/schemes/decimal/tests.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::iter; +use std::sync::LazyLock; + +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; +use rstest::rstest; +use vortex_array::ArrayId; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DecimalArray; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::i256; +use vortex_array::session::ArraySessionExt; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +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_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; + +use super::DecimalScheme; +use crate::BtrBlocksCompressor; +use crate::BtrBlocksCompressorBuilder; +use crate::SchemeExt; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session(); + vortex_decimal_byte_parts::initialize(&session); + session +}); + +/// Number of values per array: large enough for cascaded integer schemes to use sampling. +const N: usize = 16_384; + +fn ten_pow(exp: u32) -> i256 { + i256::from_i128(10).wrapping_pow(exp) +} + +/// Deterministic 24-bit noise, so the low part of each value is neither constant nor a +/// sequence — the realistic shape for a wide decimal column with a large fixed magnitude. +fn noise(seed: u64) -> impl Iterator { + let mut rng = StdRng::seed_from_u64(seed); + iter::repeat_with(move || i128::from(rng.random::() >> 8)) +} + +/// `i128`-backed values that need more than 64 bits, so the encoding must carry one lower +/// part. +fn wide_i128_array(validity: Validity) -> DecimalArray { + let base = 10i128.pow(25); + let values: Buffer = noise(7).take(N).map(|delta| base + delta).collect(); + DecimalArray::new(values, DecimalDType::new(38, 2), validity) +} + +/// `i256`-backed values that need more than 128 bits, so the encoding must carry three +/// lower parts. +fn wide_i256_array(validity: Validity) -> DecimalArray { + let base = ten_pow(40); + let values: Buffer = noise(11) + .take(N) + .map(|delta| base + i256::from_i128(delta)) + .collect(); + DecimalArray::new(values, DecimalDType::new(76, 2), validity) +} + +/// Compress with no restriction on serialized IDs, so wide values may split. +fn compress(array: &ArrayRef) -> VortexResult { + BtrBlocksCompressor::default().compress(array, &mut SESSION.create_execution_ctx()) +} + +/// Compress as a writer whose editions permit the frozen byte-parts format but not v2. +fn compress_v1_only(array: &ArrayRef) -> VortexResult { + let v1_only = HashSet::from([decimal_byte_parts_v1_id()]); + BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&v1_only) + .build() + .compress(array, &mut SESSION.create_execution_ctx()) +} + +fn byte_parts(array: &ArrayRef) -> &ArrayRef { + assert!( + array.is::(), + "expected DecimalByteParts, got {}", + array.encoding_id() + ); + array +} + +fn lower_part_count(array: &ArrayRef) -> usize { + byte_parts(array) + .as_opt::() + .vortex_expect("byte parts array") + .lower_parts() + .len() +} + +/// When the writer may emit the v2 format, values too wide for a single signed part split into +/// lower parts: one for `i128` storage, three for `i256`. +#[rstest] +#[case::i128(wide_i128_array(Validity::NonNullable).into_array(), 1)] +#[case::i128_nullable(wide_i128_array(Validity::from_iter((0..N).map(|i| i % 3 != 0))).into_array(), 1)] +#[case::i256(wide_i256_array(Validity::NonNullable).into_array(), 3)] +#[case::i256_nullable(wide_i256_array(Validity::from_iter((0..N).map(|i| i % 5 != 0))).into_array(), 3)] +fn test_wide_decimals_split_when_v2_is_permitted( + #[case] array: ArrayRef, + #[case] expected_lower_parts: usize, + #[values(false, true)] explicit_ids: bool, +) -> VortexResult<()> { + let mut builder = BtrBlocksCompressorBuilder::default(); + if explicit_ids { + builder = builder.retain_allowed_encodings(&HashSet::from([ + decimal_byte_parts_v1_id(), + decimal_byte_parts_v2_id(), + ])); + } + let compressed = builder + .build() + .compress(&array, &mut SESSION.create_execution_ctx())?; + assert_eq!(lower_part_count(&compressed), expected_lower_parts); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + + let serialization = SESSION + .array_serialize(&compressed)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialization.serialized_id, decimal_byte_parts_v2_id()); + Ok(()) +} + +/// A writer that may emit only the frozen format leaves wide values as the canonical decimal: +/// splitting them would need lower parts, for which no single-part form exists. +#[rstest] +#[case::i128(wide_i128_array(Validity::NonNullable).into_array())] +#[case::i128_nullable(wide_i128_array(Validity::from_iter((0..N).map(|i| i % 3 != 0))).into_array())] +#[case::i256(wide_i256_array(Validity::NonNullable).into_array())] +#[case::i256_nullable(wide_i256_array(Validity::from_iter((0..N).map(|i| i % 5 != 0))).into_array())] +fn test_wide_decimals_stay_canonical_without_v2(#[case] array: ArrayRef) -> VortexResult<()> { + let compressed = compress_v1_only(&array)?; + + assert!( + compressed.as_opt::().is_none(), + "expected the wide decimal to be left canonical, got {}", + compressed.encoding_id() + ); + assert_eq!(compressed.dtype(), array.dtype()); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[test] +fn test_i256_decimal_round_trips_extreme_values() -> VortexResult<()> { + // Every 64-bit window exercised, including the sign boundary of the most significant + // part. Bounded by the precision so the values are legal `Decimal(76, 0)` scalars. + let max = ten_pow(76) - i256::ONE; + let values: Buffer = (0..N) + .map(|i| match i % 8 { + 0 => i256::ZERO, + 1 => i256::ONE, + 2 => i256::ZERO - i256::ONE, + 3 => i256::from_parts(u128::MAX, 0), + 4 => i256::from_parts(0, 1), + 5 => i256::from_parts(0, -1), + 6 => max, + _ => i256::ZERO - max, + }) + .collect(); + let array = + DecimalArray::new(values, DecimalDType::new(76, 0), Validity::NonNullable).into_array(); + + let compressed = compress(&array)?; + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +#[rstest] +fn test_narrow_decimal_has_no_lower_parts( + #[values(false, true)] v1_only: bool, +) -> VortexResult<()> { + // Values that fit 64 bits are narrowed rather than split, even when the declared + // precision needs an i256. + let values: Buffer = (0..N as i128).map(|i| i256::from_i128(i * 3)).collect(); + let array = + DecimalArray::new(values, DecimalDType::new(76, 2), Validity::NonNullable).into_array(); + + let compressed = if v1_only { + compress_v1_only(&array)? + } else { + compress(&array)? + }; + assert_eq!(lower_part_count(&compressed), 0); + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + + // Narrow values keep the frozen format even when v2 is permitted. + let serialization = SESSION + .array_serialize(&compressed)? + .vortex_expect("byte parts arrays are serializable"); + assert_eq!(serialization.serialized_id, decimal_byte_parts_v1_id()); + Ok(()) +} + +#[rstest] +fn test_narrow_precision_with_wide_null_slot( + #[values(false, true)] v1_only: bool, +) -> VortexResult<()> { + let array = DecimalArray::new( + buffer![1i64, i64::MAX, 3], + DecimalDType::new(2, 0), + Validity::from_iter([true, false, true]), + ) + .into_array(); + let compressed = if v1_only { + compress_v1_only(&array)? + } else { + compress(&array)? + }; + assert_arrays_eq!(array, compressed, &mut SESSION.create_execution_ctx()); + Ok(()) +} + +/// The frozen format is the one the scheme always writes, so it must be permitted for the scheme +/// to run at all. The v2 format is optional. +#[rstest] +#[case::neither(vec![], false)] +#[case::v1(vec![decimal_byte_parts_v1_id()], true)] +#[case::v2_without_v1(vec![decimal_byte_parts_v2_id()], false)] +#[case::both(vec![decimal_byte_parts_v1_id(), decimal_byte_parts_v2_id()], true)] +fn test_decimal_scheme_needs_the_frozen_format(#[case] allowed: Vec, #[case] kept: bool) { + let compressor = BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed.into_iter().collect()) + .build(); + assert_eq!(compressor.has_scheme(DecimalScheme.id()), kept); +} + +#[test] +fn test_canonical_of_compressed_wide_decimal_keeps_storage_width() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + + let array = wide_i128_array(Validity::NonNullable).into_array(); + let canonical = compress(&array)?.execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I128); + + let array = wide_i256_array(Validity::NonNullable).into_array(); + let canonical = compress(&array)?.execute::(&mut ctx)?; + assert_eq!(canonical.values_type(), DecimalType::I256); + Ok(()) +} diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index db947430ed8..f50bc276d7e 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -65,6 +65,7 @@ vortex-build = { workspace = true } [dev-dependencies] allocator-api2 = { workspace = true } divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f740aacb318..0b3acdb1f1a 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -11,6 +11,9 @@ use flatbuffers::FlatBufferBuilder; use futures::StreamExt; use futures::TryStreamExt; use futures::pin_mut; +use rand::RngExt; +use rand::SeedableRng as _; +use rand::rngs::StdRng; use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::IntoArray; @@ -38,6 +41,7 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; +use vortex_array::dtype::i256; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::and; @@ -72,9 +76,16 @@ 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_decimal_byte_parts::split_decimal; +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_error::vortex_err; use vortex_io::session::RuntimeSession; use vortex_layout::DynLayout; use vortex_layout::LayoutStrategy; @@ -2879,3 +2890,148 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { assert_eq!(result.len(), 1); Ok(()) } + +/// End-to-end check that decimals wider than 64 bits survive a write/read round trip. +/// +/// The test session permits both byte-parts wire formats, so wide values can be split and compressed. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_wide_decimal_round_trips_through_a_file() -> VortexResult<()> { + const N: usize = 16_384; + + /// Deterministic 24-bit noise, so the low bits of each value are neither constant nor a + /// sequence. + fn noise(seed: u64) -> impl Iterator { + let mut rng = StdRng::seed_from_u64(seed); + iter::repeat_with(move || i128::from(rng.random::() >> 8)) + } + + // Values that need more than 64 bits, so `i128` storage cannot be narrowed away. + let decimal_38 = DecimalArray::new( + noise(7) + .take(N) + .map(|delta| 10i128.pow(25) + delta) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array(); + + // Values that need more than 128 bits, so `i256` storage cannot be narrowed away. + let base = i256::from_i128(10).wrapping_pow(40); + let decimal_76 = DecimalArray::new( + noise(11) + .take(N) + .map(|delta| base + i256::from_i128(delta)) + .collect::>(), + DecimalDType::new(76, 4), + Validity::from_iter((0..N).map(|i| i % 9 != 0)), + ) + .into_array(); + + let st = StructArray::from_fields(&[ + ("decimal_38", decimal_38), + ("decimal_76_nullable", decimal_76), + ])? + .into_array(); + let dtype = st.dtype().clone(); + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, st.clone().to_array_stream()) + .await?; + + let chunks: Vec<_> = SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .try_collect() + .await?; + let read = ChunkedArray::try_new(chunks, dtype)?.into_array(); + + let mut ctx = SESSION.create_execution_ctx(); + assert_eq!(read.len(), N); + assert_arrays_eq!(st, read, &mut ctx); + + Ok(()) +} + +/// The default writer splits wide decimals only when its enabled editions permit the v2 +/// format, including when its input already carries lower parts. +#[rstest] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_default_writer_selects_decimal_scheme_version( + #[values(false, true)] v2_enabled: bool, +) -> VortexResult<()> { + let session = array_session() + .with::() + .with::(); + crate::register_default_encodings(&session); + if v2_enabled { + crate::enable_all_registered_array_encodings(&session); + } else { + for declaration in EDITION_DECLARATIONS { + session + .register_edition(declaration) + .map_err(|error| vortex_err!("{error}"))?; + } + session + .enable_edition(CORE_2026_08_3) + .map_err(|error| vortex_err!("{error}"))?; + } + + let decimal = DecimalArray::new( + (0..64i128) + .map(|i| (1i128 << 70) + i) + .collect::>(), + DecimalDType::new(38, 2), + Validity::NonNullable, + ); + + // Building the encoded array is allowed; only getting it into a file is restricted. + let parts = split_decimal(&decimal, &mut session.create_execution_ctx())?; + assert_eq!(parts.lower_parts.len(), 1, "expected a wide split"); + let encoded = DecimalByteParts::try_new_with_lower_parts( + parts.msp, + parts.lower_parts, + decimal.decimal_dtype(), + )? + .into_array(); + + let st = StructArray::from_fields(&[("wide", encoded)])?.into_array(); + let mut buf = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buf, st.clone().to_array_stream()) + .await?; + + let chunks: Vec<_> = session + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .try_collect() + .await?; + + let lower_part_counts: Vec = chunks + .iter() + .flat_map(|chunk| chunk.depth_first_traversal()) + .filter_map(|node| { + node.as_opt::() + .map(|array| array.lower_parts().len()) + }) + .collect(); + assert_eq!(lower_part_counts.is_empty(), !v2_enabled); + assert!( + lower_part_counts.iter().all(|count| *count == 1), + "expected one lower part per array, got {lower_part_counts:?}" + ); + + let mut ctx = session.create_execution_ctx(); + let read = ChunkedArray::try_new(chunks, st.dtype().clone())?.into_array(); + assert_arrays_eq!(st, read, &mut ctx); + Ok(()) +}