From 936aa0a349cf4931203f698cf840ad26ec201305 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Mon, 2 Nov 2026 15:24:02 -0500 Subject: [PATCH 1/3] Support decimal values in Widen Keep the logical decimal dtype while storing values in a narrower decimal child. Execute into the logical dtype's canonical buffer width and serialize narrow leaves through the legacy decimal wire format without executing them. Signed-off-by: Connor Tsui --- scripts/codspeed_bench_counts.json | 3 +- vortex-array/Cargo.toml | 4 + vortex-array/benches/widen_decimal.rs | 337 ++++++++++++++++++ vortex-array/src/arrays/widen/array.rs | 46 ++- vortex-array/src/arrays/widen/encode.rs | 97 +++++ vortex-array/src/arrays/widen/mod.rs | 34 +- .../src/arrays/widen/tests/decimal.rs | 181 ++++++++++ .../arrays/widen/tests/decimal_fallback.rs | 175 +++++++++ vortex-array/src/arrays/widen/tests/mod.rs | 10 +- vortex-array/src/arrays/widen/tests/serde.rs | 240 +++++++++++++ vortex-array/src/arrays/widen/vtable/mod.rs | 69 +++- .../src/arrays/widen/vtable/plugin.rs | 86 +++++ vortex-array/src/session/mod.rs | 2 + vortex-ipc/src/messages/decoder.rs | 40 ++- vortex-layout/src/layouts/flat/writer.rs | 50 +++ 15 files changed, 1333 insertions(+), 41 deletions(-) create mode 100644 vortex-array/benches/widen_decimal.rs create mode 100644 vortex-array/src/arrays/widen/tests/decimal.rs create mode 100644 vortex-array/src/arrays/widen/tests/decimal_fallback.rs create mode 100644 vortex-array/src/arrays/widen/tests/serde.rs create mode 100644 vortex-array/src/arrays/widen/vtable/plugin.rs diff --git a/scripts/codspeed_bench_counts.json b/scripts/codspeed_bench_counts.json index 3c45eb9b7ec..00ced92f623 100644 --- a/scripts/codspeed_bench_counts.json +++ b/scripts/codspeed_bench_counts.json @@ -74,7 +74,8 @@ "varbinview_fill_null": 2, "varbinview_try_new": 6, "varbinview_zip": 2, - "widen": 46 + "widen": 46, + "widen_decimal": 42 }, "vortex-buffer": { "allocation": 35, diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 8560aa9e6f9..5a29e064667 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -341,6 +341,10 @@ harness = false name = "widen" harness = false +[[bench]] +name = "widen_decimal" +harness = false + [[bench]] name = "run_end_null_count" harness = false diff --git a/vortex-array/benches/widen_decimal.rs b/vortex-array/benches/widen_decimal.rs new file mode 100644 index 00000000000..699043a02fe --- /dev/null +++ b/vortex-array/benches/widen_decimal.rs @@ -0,0 +1,337 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compute on decimal `Widen` arrays and canonical arrays with the same values. +//! +//! The cases cover precision 18 over precision 9, precision 10 over precision 5, and precision 38 +//! over precision 18, all at scale 2. The precision 5 values fit in `i16`, but require `i32` to +//! represent their full decimal domain. Each iteration shares value buffers and builds fresh arrays +//! so cached statistics cannot remove the operation's work. `Widen` buffers are selected once with +//! [`WidenArray::encode_decimal`], outside the timed region. + +use std::fmt; +use std::sync::LazyLock; + +use divan::Bencher; +use mimalloc::MiMalloc; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::min_max; +use vortex_array::aggregate_fn::fns::sum; +use vortex_array::aggregate_fn::session::AggregateFnSession; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Decimal; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::WidenArray; +use vortex_array::arrays::widen::WidenArraySlotsExt; +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::memory::MemorySession; +use vortex_array::optimizer::kernels::KernelSession; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + LazyLock::force(&SESSION); + + // Keep process-wide execution initialization outside the first measured operation. + let case = Case { + input: Input::Widen, + ..CASES[0] + }; + case.values() + .array(case) + .execute::(&mut SESSION.create_execution_ctx()) + .vortex_expect("the benchmark input executes to canonical decimals"); + + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(|| { + VortexSession::empty() + .with::() + .with::() + .with::() +}); + +const LEN: usize = 32_768; +const SCALE: i8 = 2; + +#[derive(Clone, Copy, Debug)] +struct Case { + precision: u8, + child_precision: u8, + input: Input, +} + +const CASES: &[Case] = &[ + Case { + precision: 18, + child_precision: 9, + input: Input::Canonical, + }, + Case { + precision: 10, + child_precision: 5, + input: Input::Canonical, + }, + Case { + precision: 38, + child_precision: 18, + input: Input::Canonical, + }, +]; + +#[derive(Clone, Copy, Debug)] +enum Input { + Canonical, + Widen, +} + +fn cases() -> Vec { + CASES + .iter() + .flat_map(|&case| { + [ + case, + Case { + input: Input::Widen, + ..case + }, + ] + }) + .collect() +} + +impl fmt::Display for Case { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "p{}_childp{}_{}", + self.precision, + self.child_precision, + match self.input { + Input::Canonical => "canonical", + Input::Widen => "widen", + }, + ) + } +} + +impl Case { + fn decimal_dtype(self) -> DecimalDType { + DecimalDType::new(self.precision, SCALE) + } + + fn values(self) -> Values { + let (minimum, maximum) = match self.child_precision { + 5 => (10_000, 32_767), + 9 => (-999_999_999, 999_999_999), + 18 => (-999_999_999_999_999_999, 999_999_999_999_999_999), + _ => unreachable!("CASES contains only child precisions 5, 9, and 18"), + }; + let mut rng = StdRng::seed_from_u64(0); + + // Fixed extrema make the required child precision independent of the random sample. + let values = [minimum, maximum] + .into_iter() + .chain((2..LEN).map(|_| rng.random_range(minimum..=maximum))); + + let values = if self.precision <= 18 { + Values::I64(values.collect()) + } else { + Values::I128(values.map(i128::from).collect()) + }; + + if matches!(self.input, Input::Canonical) { + return values; + } + + let encoded = WidenArray::encode_decimal( + values.decimal_array(self.decimal_dtype()), + &mut SESSION.create_execution_ctx(), + ) + .vortex_expect("the benchmark values fit a narrower decimal precision"); + let widen = encoded.as_::(); + let child = widen.values().as_::(); + assert_eq!(child.precision(), self.child_precision); + assert_eq!( + child.values_type(), + child.decimal_dtype().canonical_values_type(), + ); + + match child.values_type() { + DecimalType::I32 => Values::I32(child.buffer::()), + DecimalType::I64 => Values::I64(child.buffer::()), + _ => unreachable!("the benchmark child precisions require i32 or i64"), + } + } + + fn constant(self, value: i64) -> ArrayRef { + let scalar = Scalar::decimal(value.into(), self.decimal_dtype(), Nullability::NonNullable); + ConstantArray::new(scalar, LEN).into_array() + } + + /// Returns a constant that splits the values, so `compare_constant` is not uniform. + /// + /// The precision 5 values lie in `10_000..=32_767`, so 0 would be below all of them. The other + /// ranges are symmetric about zero, so 0 splits them. + fn in_range(self) -> i64 { + if self.child_precision == 5 { 20_000 } else { 0 } + } + + fn out_of_range(self) -> i64 { + 10i64.pow(u32::from(self.child_precision)) + } + + fn cast_target(self) -> DType { + let precision = if self.precision == 38 { 76 } else { 38 }; + DType::Decimal( + DecimalDType::new(precision, SCALE), + Nullability::NonNullable, + ) + } +} + +enum Values { + I32(Buffer), + I64(Buffer), + I128(Buffer), +} + +impl Values { + fn decimal_array(&self, dtype: DecimalDType) -> DecimalArray { + match self { + Self::I32(values) => DecimalArray::new(values.clone(), dtype, Validity::NonNullable), + Self::I64(values) => DecimalArray::new(values.clone(), dtype, Validity::NonNullable), + Self::I128(values) => DecimalArray::new(values.clone(), dtype, Validity::NonNullable), + } + } + + fn array(&self, case: Case) -> ArrayRef { + match case.input { + Input::Canonical => self.decimal_array(case.decimal_dtype()).into_array(), + Input::Widen => { + let child_dtype = DecimalDType::new(case.child_precision, SCALE); + let child = self.decimal_array(child_dtype).into_array(); + let dtype = DType::Decimal(case.decimal_dtype(), Nullability::NonNullable); + WidenArray::try_new(child, dtype) + .vortex_expect("encode_decimal selected a strictly narrower child dtype") + .into_array() + } + } + } +} + +fn bench_case( + bencher: Bencher, + setup: impl Fn() -> T + Sync, + operation: impl Fn(T, &mut ExecutionCtx) -> O + Sync, +) { + bencher + .with_inputs(|| (setup(), SESSION.create_execution_ctx())) + .bench_values(|(operands, mut ctx)| operation(operands, &mut ctx)); +} + +fn execute_canonical(array: VortexResult, ctx: &mut ExecutionCtx) -> Canonical { + array + .vortex_expect("the operation is valid for the benchmark inputs") + .execute::(ctx) + .vortex_expect("the operation executes to canonical for the benchmark inputs") +} + +#[divan::bench(args = cases())] +fn compare_constant(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || (values.array(case), case.constant(case.in_range())), + |(array, constant), ctx| execute_canonical(array.binary(constant, Operator::Lt), ctx), + ); +} + +#[divan::bench(args = cases())] +fn compare_constant_out_of_range(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || (values.array(case), case.constant(case.out_of_range())), + |(array, constant), ctx| execute_canonical(array.binary(constant, Operator::Lt), ctx), + ); +} + +#[divan::bench(args = cases())] +fn sum(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| sum::sum(&array, ctx).vortex_expect("the benchmark decimal sum is valid"), + ); +} + +#[divan::bench(args = cases())] +fn min_max(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| { + min_max::min_max(&array, ctx, NumericalAggregateOpts::default()) + .vortex_expect("the benchmark decimal extrema are valid") + }, + ); +} + +#[divan::bench(args = cases())] +fn filter(bencher: Bencher, case: Case) { + let values = case.values(); + let mut rng = StdRng::seed_from_u64(1); + let selection: BitBuffer = (0..LEN).map(|_| rng.random_bool(0.5)).collect(); + bench_case( + bencher, + || (values.array(case), Mask::from_buffer(selection.clone())), + |(array, mask), ctx| execute_canonical(array.filter(mask), ctx), + ); +} + +#[divan::bench(args = cases())] +fn execute(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| execute_canonical(Ok(array), ctx), + ); +} + +#[divan::bench(args = cases())] +fn cast_widen(bencher: Bencher, case: Case) { + let values = case.values(); + bench_case( + bencher, + || values.array(case), + |array, ctx| execute_canonical(array.cast(case.cast_target()), ctx), + ); +} diff --git a/vortex-array/src/arrays/widen/array.rs b/vortex-array/src/arrays/widen/array.rs index ed08a50dc47..803fd8a0ce0 100644 --- a/vortex-array/src/arrays/widen/array.rs +++ b/vortex-array/src/arrays/widen/array.rs @@ -16,7 +16,7 @@ use crate::dtype::DType; /// The child slots of a [`WidenArray`]. #[array_slots(Widen)] pub struct WidenSlots { - /// The stored values, in an integer type strictly narrower than the logical dtype. + /// The stored values, with a canonical width strictly narrower than the logical dtype. #[slot(0)] pub values: ArrayRef, } @@ -25,9 +25,12 @@ impl WidenArray { /// Creates a `Widen` that presents `values` with the wider logical `dtype`, without reading /// any values. /// - /// `values` and `dtype` must both be integer dtypes with the same signedness and nullability, - /// and `values` must be strictly narrower. When `values` is itself a `Widen`, the new array - /// stores its values instead, so that `Widen` arrays do not nest. + /// `values` and `dtype` **must** both be integers with the same signedness, or decimals with + /// the same scale. They **must** have the same nullability, and the canonical width of `values` + /// **must** be strictly narrower. When `values` is itself a `Widen`, the new array stores its + /// values instead, so that `Widen` arrays do not nest. + /// + /// A decimal child **must** be a valid decimal array for its declared dtype. pub fn try_new(mut values: ArrayRef, dtype: DType) -> VortexResult { // NB: Check `values` before the flatten. The values of a `Widen` child are narrower than // the child, so the check in `try_from_parts` alone would accept a child as wide as `dtype` @@ -52,23 +55,34 @@ impl WidenArray { /// Checks that the dtypes of `values` and `logical` meet the requirements of /// [`WidenArray::try_new`]. pub(super) fn validate_dtypes(values: &DType, logical: &DType) -> VortexResult<()> { - let (values_ptype, logical_ptype) = match (values, logical) { + match (values, logical) { (DType::Primitive(values_ptype, _), DType::Primitive(logical_ptype, _)) if values_ptype.is_int() && logical_ptype.is_int() => { - (*values_ptype, *logical_ptype) + vortex_ensure!( + values_ptype.is_signed_int() == logical_ptype.is_signed_int(), + "Widen requires values with the signedness of {logical}, got {values}" + ); + vortex_ensure!( + values_ptype.byte_width() < logical_ptype.byte_width(), + "Widen requires values narrower than {logical}, got {values}" + ); + } + (DType::Decimal(values_decimal, _), DType::Decimal(logical_decimal, _)) => { + vortex_ensure!( + values_decimal.scale() == logical_decimal.scale(), + "Widen requires values with the scale of {logical}, got {values}" + ); + vortex_ensure!( + values_decimal.canonical_values_type() < logical_decimal.canonical_values_type(), + "Widen requires values narrower than {logical}, got {values}" + ); } - _ => vortex_bail!("Widen requires integer dtypes, got {values} values for {logical}"), - }; + _ => vortex_bail!( + "Widen requires integer dtypes or decimal dtypes for {logical}, got {values}" + ), + } - vortex_ensure!( - values_ptype.is_signed_int() == logical_ptype.is_signed_int(), - "Widen requires values with the signedness of {logical}, got {values}" - ); - vortex_ensure!( - values_ptype.byte_width() < logical_ptype.byte_width(), - "Widen requires values narrower than {logical}, got {values}" - ); vortex_ensure!( values.nullability() == logical.nullability(), "Widen requires values with the nullability of {logical}, got {values}" diff --git a/vortex-array/src/arrays/widen/encode.rs b/vortex-array/src/arrays/widen/encode.rs index 24d41c334fb..141fccde6d2 100644 --- a/vortex-array/src/arrays/widen/encode.rs +++ b/vortex-array/src/arrays/widen/encode.rs @@ -1,9 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Selection of narrower `Widen` children from valid values. +//! +//! Null payloads do not constrain the child dtype. The result keeps the logical dtype and uses +//! narrower storage only when that storage can hold every valid value. + use std::ops::RangeInclusive; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use crate::ArrayRef; @@ -11,12 +17,21 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::aggregate_fn::NumericalAggregateOpts; use crate::aggregate_fn::fns::min_max::min_max; +use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::copy_to_canonical_width; +use crate::arrays::decimal::valid_range; use crate::arrays::widen::compute::wrap_values; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; use crate::dtype::PType; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; use crate::scalar::Scalar; impl WidenArray { @@ -30,6 +45,51 @@ impl WidenArray { wrap_values(narrow_values(array, ctx)?.into_array(), dtype) } + /// Stores a decimal array at the smallest precision that its valid values and scale allow. + /// + /// The result keeps the input dtype and nullability. Empty and all-null arrays use the + /// precision required by the scale, or 1 for a non-positive scale. Returns the input unchanged + /// when the selected precision has the canonical width of the input dtype. + /// + /// # Errors + /// + /// Returns an error if a valid value does not fit the input precision. + pub fn encode_decimal(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { + let logical_decimal_dtype = array.decimal_dtype(); + let validity = array.validity()?; + let validity_mask = validity.execute_mask(array.len(), ctx)?; + let valid_value_range = match_each_decimal_value_type!(array.values_type(), |T| { + valid_range(array.buffer::().as_slice(), &validity_mask).map(|(min, max)| { + DecimalValue::from(min).as_i256()..=DecimalValue::from(max).as_i256() + }) + }); + + let child_decimal_dtype = smallest_decimal_dtype(valid_value_range, logical_decimal_dtype)?; + if child_decimal_dtype.canonical_values_type() + == logical_decimal_dtype.canonical_values_type() + { + return Ok(array.into_array()); + } + + let child_buffer_handle = match_each_decimal_value_type!(array.values_type(), |T| { + copy_to_canonical_width( + array.buffer::().as_slice(), + &validity_mask, + child_decimal_dtype, + )? + }); + let logical_dtype = array.dtype().clone(); + + let child = DecimalArray::new_handle( + child_buffer_handle, + child_decimal_dtype.canonical_values_type(), + child_decimal_dtype, + validity, + ); + + Ok(WidenArray::try_new(child.into_array(), logical_dtype)?.into_array()) + } + /// Returns `array` cast to the type that [`Self::encode`] chooses for its values, without a /// `Widen` around them. /// @@ -107,3 +167,40 @@ pub(super) fn integer_value(scalar: &Scalar) -> VortexResult { u64::try_from(scalar).map(i128::from) } } + +/// Returns the smallest decimal dtype that holds an unscaled valid-value range at the parent scale. +/// +/// The precision is at least 1 and at least a positive scale. An absent range represents an empty +/// or all-null array and needs only that minimum precision. The result never exceeds the parent +/// precision. Callers **must** exclude null payloads from the range, for example with +/// [`valid_range`]. +/// +/// # Errors +/// +/// Returns an error if the minimum exceeds the maximum or an endpoint exceeds the parent precision. +pub(crate) fn smallest_decimal_dtype( + range: Option>, + parent: DecimalDType, +) -> VortexResult { + let minimum_precision = parent.scale().max(1) as u8; + let Some(range) = range else { + // Empty and all-null arrays have no valid values to constrain precision. The scale still + // requires at least that many digits when positive, or precision 1 when non-positive. + return Ok(DecimalDType::new(minimum_precision, parent.scale())); + }; + + let (min, max) = range.into_inner(); + vortex_ensure!( + min <= max, + "Decimal range requires ordered extrema, got {min} to {max}" + ); + + for precision in minimum_precision..=parent.precision() { + let index = usize::from(precision); + if i256::MIN_BY_PRECISION[index] <= min && max <= i256::MAX_BY_PRECISION[index] { + return Ok(DecimalDType::new(precision, parent.scale())); + } + } + + vortex_bail!("Decimal range must fit {parent}, got {min} to {max}") +} diff --git a/vortex-array/src/arrays/widen/mod.rs b/vortex-array/src/arrays/widen/mod.rs index ac748864ad2..7007a309f8e 100644 --- a/vortex-array/src/arrays/widen/mod.rs +++ b/vortex-array/src/arrays/widen/mod.rs @@ -1,40 +1,50 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! An in-memory encoding that stores an integer array in a narrower integer type. +//! An in-memory encoding that stores integer and decimal arrays at a narrower canonical width. //! -//! A [`WidenArray`] keeps the logical dtype of an integer array, and stores its values and their -//! validity in a child with a strictly narrower integer type of the same signedness and -//! nullability. For example, an `i64` array whose values are all between 0 and 200 can store `i16` -//! values, but not `u8` values. [`WidenArray::encode`] picks the narrowest such type for a -//! [`PrimitiveArray`]. +//! A [`WidenArray`] keeps the logical dtype of an integer or decimal array, and stores its values +//! and their validity in a child with the same nullability and a strictly narrower canonical width. +//! +//! An integer child has a narrower integer type of the same signedness. For example, an `i64` array +//! whose values are all between 0 and 200 can store `i16` values, but not `u8` values. +//! [`WidenArray::encode`] picks the narrowest such type for a [`PrimitiveArray`]. +//! +//! A decimal child has the logical scale and a lower precision, and **must** be valid for its dtype. +//! [`WidenArray::encode_decimal`] picks the smallest precision that holds the valid values and the +//! scale. Decimal execution copies the child values into the logical canonical width. //! //! Narrowing is a physical choice, so it **must not** change a result. Execution casts the values -//! to the logical dtype, and an operation runs on the values only where it gives the canonical -//! result. A `Widen` never wraps another `Widen`. +//! to the logical dtype for integers, and an operation runs on the values only where it gives the +//! canonical result. A `Widen` never wraps another `Widen`. //! //! Take executes its gather instead of leaving a lazy [`Dict`] under a `Widen`, because code such //! as the DuckDB exporter and the dict layout writer looks for a `Dict` at the root of an array. //! -//! `Widen` has no wire form yet, so the default [`ArraySession`] does not register it, and an -//! array that holds a `Widen` cannot be serialized. +//! `Widen` has no wire form yet. The default [`ArraySession`] registers [`WidenPlugin`], which +//! writes a decimal `Widen` with a [`Decimal`](crate::arrays::Decimal) leaf using the legacy +//! `vortex.decimal` wire form. The leaf supplies the stored width and the containing format supplies +//! the parent dtype. Serialization keeps the parent statistics and never executes the array. +//! Integer arrays and non-decimal leaves require explicit preparation by the caller. //! //! [`ArraySession`]: crate::session::ArraySession //! [`Dict`]: crate::arrays::Dict //! [`PrimitiveArray`]: crate::arrays::PrimitiveArray +use vortex_session::VortexSession; + mod array; pub use array::WidenArraySlotsExt; pub use array::WidenSlots; pub use array::WidenSlotsView; -pub use vtable::WidenArray; pub(crate) mod compute; mod encode; mod vtable; -use vortex_session::VortexSession; pub use vtable::Widen; +pub use vtable::WidenArray; +pub use vtable::WidenPlugin; use crate::array::VTable; use crate::arrays::Dict; diff --git a/vortex-array/src/arrays/widen/tests/decimal.rs b/vortex-array/src/arrays/widen/tests/decimal.rs new file mode 100644 index 00000000000..556d2643c53 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal.rs @@ -0,0 +1,181 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Construction, execution, and encoding contracts for decimal `Widen` arrays. + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::ArrayRef; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::arrays::widen::encode::smallest_decimal_dtype; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::Nullability::NonNullable; +use crate::dtype::Nullability::Nullable; +use crate::dtype::PType; +use crate::dtype::i256; +use crate::validity::Validity; + +fn dtype(precision: u8, scale: i8, nullability: Nullability) -> DType { + DType::Decimal(DecimalDType::new(precision, scale), nullability) +} + +#[rstest] +#[case::same_width(dtype(10, 2, NonNullable), dtype(18, 2, NonNullable), "narrower")] +#[case::scale(dtype(9, 1, NonNullable), dtype(18, 2, NonNullable), "scale")] +#[case::higher_precision(dtype(18, 2, NonNullable), dtype(9, 2, NonNullable), "narrower")] +#[case::integer_child( + PType::I32.into(), + dtype(18, 2, NonNullable), + "integer dtypes or decimal dtypes", +)] +#[case::integer_parent( + dtype(9, 2, NonNullable), + PType::I64.into(), + "integer dtypes or decimal dtypes", +)] +fn test_rejects_invalid_decimal_dtypes( + #[case] child: DType, + #[case] parent: DType, + #[case] requirement: &str, +) { + let values = Canonical::empty(&child).into_array(); + + let error = WidenArray::try_new(values, parent) + .expect_err("the child and parent violate a Widen dtype requirement"); + assert!(error.to_string().contains(requirement), "{error}"); +} + +#[rstest] +#[case::canonical( + DecimalArray::from_iter([-999i32, 0, 999], DecimalDType::new(9, 2)), + [-999, 0, 999], +)] +#[case::wider_stored( + DecimalArray::from_iter([-999i128, 0, 999], DecimalDType::new(9, 2)), + [-999, 0, 999], +)] +#[case::null_garbage( + DecimalArray::new( + buffer![i128::MIN, 999, i128::MAX], + DecimalDType::new(9, 2), + Validity::from_iter([false, true, false]), + ), + [0, 999, 0], +)] +fn test_decimal_execute_widens_and_terminates( + #[case] child: DecimalArray, + #[case] expected: [i64; 3], +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let validity = child.validity()?; + let parent = dtype(18, 2, validity.nullability()); + let reference = DecimalArray::new( + expected.into_iter().collect::>(), + DecimalDType::new(18, 2), + validity, + ); + + let widened = WidenArray::try_new(child.into_array(), parent.clone())?.into_array(); + let executed = widened.execute::(&mut ctx)?.into_array(); + + assert_eq!(executed.dtype(), &parent); + let decimal = executed.as_::(); + assert_eq!(decimal.values_type(), DecimalType::I64); + assert_arrays_eq!(executed, reference, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::scale_raises_precision(vec![Some(1)], 7, Some(7))] +#[case::negative_boundary( vec![Some(-128), Some(99)], 0, Some(3))] +#[case::negative_scale( vec![Some(-99), Some(99)], -3, Some(2))] +#[case::all_null( vec![None, None], 0, Some(1))] +#[case::empty( vec![], 0, Some(1))] +#[case::empty_positive_scale( vec![], 12, None)] +#[case::empty_negative_scale( vec![], -4, Some(1))] +#[case::no_narrower_width( vec![Some(1_000_000_000)], 2, None)] +fn test_encode_decimal_best_precision( + #[case] values: Vec>, + #[case] scale: i8, + #[case] child_precision: Option, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = DecimalArray::from_option_iter(values, DecimalDType::new(18, scale)); + + let encoded = WidenArray::encode_decimal(input.clone(), &mut ctx)?; + + if let Some(precision) = child_precision { + let widened = encoded.as_::(); + let child = widened.values().as_::(); + assert_eq!(child.decimal_dtype(), DecimalDType::new(precision, scale)); + assert_eq!( + child.values_type(), + child.decimal_dtype().canonical_values_type() + ); + } else { + assert!(ArrayRef::ptr_eq(&encoded, &input.clone().into_array())); + } + assert_arrays_eq!(encoded, input, &mut ctx); + + Ok(()) +} + +#[test] +fn test_encode_decimal_ignores_null_payloads() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let input = DecimalArray::new( + buffer![i64::MIN, -128, i64::MAX], + DecimalDType::new(18, 0), + Validity::from_iter([false, true, false]), + ); + + let encoded = WidenArray::encode_decimal(input.clone(), &mut ctx)?; + + assert_eq!( + encoded.as_::().values().dtype(), + &dtype(3, 0, Nullable) + ); + assert_arrays_eq!(encoded, input, &mut ctx); + + Ok(()) +} + +#[test] +fn test_encode_decimal_rejects_value_outside_precision() { + let input = DecimalArray::from_iter([1_000_000_000_000_000_000i64], DecimalDType::new(18, 0)); + + let error = WidenArray::encode_decimal(input, &mut SESSION.create_execution_ctx()) + .expect_err("the valid value exceeds the input precision"); + assert!( + error.to_string().contains("Decimal range must fit"), + "{error}" + ); +} + +#[test] +fn test_smallest_decimal_dtype_i256_extrema() -> VortexResult<()> { + let parent = DecimalDType::new(76, 0); + let range = Some(i256::MIN_BY_PRECISION[76]..=i256::MAX_BY_PRECISION[76]); + assert_eq!(smallest_decimal_dtype(range, parent)?, parent); + assert!(smallest_decimal_dtype(Some(i256::MIN..=i256::MAX), parent).is_err()); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/decimal_fallback.rs b/vortex-array/src/arrays/widen/tests/decimal_fallback.rs new file mode 100644 index 00000000000..2df0d452f23 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_fallback.rs @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Conformance and canonical fallbacks for decimal widening. +//! +//! Arithmetic, zip, chunked and interleave execution preserve the logical decimal dtype and nulls. + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::SESSION; +use super::canonical; +use crate::ArrayRef; +use crate::Canonical; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::compute::conformance::binary_numeric::test_binary_numeric_array; +use crate::compute::conformance::cast::test_cast_conformance; +use crate::compute::conformance::consistency::test_array_consistency; +use crate::compute::conformance::filter::test_filter_conformance; +use crate::compute::conformance::mask::test_mask_conformance; +use crate::compute::conformance::take::test_take_conformance; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::scalar_fn::fns::operators::Operator; +use crate::validity::Validity; + +fn decimal_widen(validity: Validity) -> VortexResult { + let values = DecimalArray::new( + buffer![-125i16, 50, 50, 250, 900], + DecimalDType::new(4, 2), + validity, + ); + let dtype = DType::Decimal(DecimalDType::new(18, 2), values.dtype().nullability()); + + Ok(WidenArray::try_new(values.into_array(), dtype)?.into_array()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_conformance(#[case] validity: Validity) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(validity)?; + assert!(array.is::()); + + test_array_consistency(&array, &mut ctx); + test_cast_conformance(&array, &mut ctx); + test_filter_conformance(&array, &mut ctx); + test_take_conformance(&array, &mut ctx); + test_mask_conformance(&array, &mut ctx); + test_binary_numeric_array(&array, &mut ctx); + + Ok(()) +} + +#[rstest] +fn test_decimal_arithmetic_falls_back_to_canonical( + #[values(Operator::Add, Operator::Sub, Operator::Mul, Operator::Div)] operator: Operator, + #[values(false, true)] widen_on_right: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(Validity::from_iter([true, false, true, true, false]))?; + let other = DecimalArray::from_iter([25i64, 50, -100, 125, -200], DecimalDType::new(18, 2)) + .into_array(); + let (lhs, rhs) = if widen_on_right { + (other, array) + } else { + (array, other) + }; + let expected = canonical(&lhs, &mut ctx)? + .binary(canonical(&rhs, &mut ctx)?, operator)? + .execute::(&mut ctx)?; + + let actual = lhs.binary(rhs, operator)?.execute::(&mut ctx)?; + + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::true_branch(true, false)] +#[case::false_branch(false, true)] +#[case::both_branches(true, true)] +fn test_decimal_zip_falls_back_to_canonical( + #[case] widen_if_true: bool, + #[case] widen_if_false: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let if_true = decimal_widen(Validity::from_iter([true, false, true, true, false]))?; + let if_false = WidenArray::try_new( + DecimalArray::from_iter([500i16, -600, 700, -800, 900], DecimalDType::new(4, 2)) + .into_array(), + DType::Decimal(DecimalDType::new(18, 2), Nullability::NonNullable), + )? + .into_array(); + let canonical_true = canonical(&if_true, &mut ctx)?; + let canonical_false = canonical(&if_false, &mut ctx)?; + let mask = + BoolArray::from_iter([Some(true), Some(true), None, Some(false), Some(true)]).into_array(); + let expected = mask + .zip(canonical_true.clone(), canonical_false.clone())? + .execute::(&mut ctx)?; + let if_true = if widen_if_true { + if_true + } else { + canonical_true + }; + let if_false = if widen_if_false { + if_false + } else { + canonical_false + }; + + let actual = mask + .zip(if_true, if_false)? + .execute::(&mut ctx)?; + + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_decimal_chunked_execution_matches_canonical() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(Validity::from_iter([true, false, true, true, false]))?; + let expected = canonical(&array, &mut ctx)?; + let chunks = ChunkedArray::from_iter([ + array.slice(0..2)?, + expected.slice(2..3)?, + array.slice(3..5)?, + ]) + .into_array(); + + let actual = chunks.execute::(&mut ctx)?; + + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_decimal_interleave_falls_back_to_canonical() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(Validity::from_iter([true, false, true, true, false]))?; + let other = DecimalArray::from_iter([500i64, -600, 700, -800, 950], DecimalDType::new(18, 2)) + .into_array(); + let array_indices = buffer![0u8, 1, 0, 1, 0, 0, 1].into_array(); + let row_indices = buffer![4u32, 0, 1, 4, 0, 4, 1].into_array(); + let expected = array_indices + .interleave( + [canonical(&array, &mut ctx)?, other.clone()], + row_indices.clone(), + )? + .execute::(&mut ctx)?; + + let actual = array_indices + .interleave([array, other], row_indices)? + .execute::(&mut ctx)?; + + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/mod.rs b/vortex-array/src/arrays/widen/tests/mod.rs index 6a959e199d3..56c15dd4803 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -1,15 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Tests of the `Widen` encoding. +//! Tests of the `Widen` encoding over integer and decimal values. //! -//! The named fixtures cover signed and unsigned values, each kind of validity and a single row. -//! The `compute` tests run their operations over them. +//! The named integer fixtures cover signed and unsigned values, each kind of validity and a single +//! row, and the `compute` tests run their operations over them. The `decimal` modules build their +//! own decimal fixtures, and compare each result with the result on the canonical array. mod aggregates; mod compute; +mod decimal; +mod decimal_fallback; mod encoding; mod numeric; +mod serde; use std::sync::LazyLock; diff --git a/vortex-array/src/arrays/widen/tests/serde.rs b/vortex-array/src/arrays/widen/tests/serde.rs new file mode 100644 index 00000000000..44baaf43f0e --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/serde.rs @@ -0,0 +1,240 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serialization through existing wire forms, including nested arrays and stored statistics. +//! +//! Wire IDs and logical values are checked separately because a `Widen` can read back as another +//! encoding. + +use rstest::rstest; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use super::SESSION; +use crate::ArrayContext; +use crate::ArrayId; +use crate::ArrayRef; +use crate::ArrayVTable; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::BoolArray; +use crate::arrays::ChunkedArray; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::DictArray; +use crate::arrays::SliceArray; +use crate::arrays::StructArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::expr::stats::Precision; +use crate::expr::stats::Stat; +use crate::expr::stats::StatsProvider; +use crate::scalar::Scalar; +use crate::serde::SerializeOptions; +use crate::serde::SerializedArray; +use crate::session::ArraySessionExt; +use crate::validity::Validity; + +fn serialize_blob( + array: &ArrayRef, + ctx: &ArrayContext, + session: &VortexSession, +) -> VortexResult { + let mut blob = ByteBufferMut::empty(); + for buffer in array.serialize(ctx, session, &SerializeOptions::default())? { + blob.extend_from_slice(buffer.as_ref()); + } + + Ok(blob.freeze()) +} + +/// Returns a `Widen` of the logical dtype `logical` over `values`, with the nullability of +/// `values`. +fn widen_over(values: impl IntoArray, logical: DType) -> VortexResult { + let values = values.into_array(); + let logical = logical.with_nullability(values.dtype().nullability()); + + Ok(WidenArray::try_new(values, logical)?.into_array()) +} + +/// Serializes `array` with an unrestricted context and decodes it again. Returns the decoded +/// array and the IDs that the context interned. +fn serialize_round_trip( + array: &ArrayRef, + session: &VortexSession, +) -> VortexResult<(ArrayRef, Vec)> { + let ctx = ArrayContext::empty(); + let serialized = SerializedArray::try_from(serialize_blob(array, &ctx, session)?)?; + let ids = ctx.to_ids(); + let decoded = serialized.decode( + array.dtype(), + array.len(), + &ReadContext::new(ids.clone()), + session, + )?; + + Ok((decoded, ids)) +} + +#[test] +fn decimal_widen_writes_narrow_bytes() -> VortexResult<()> { + let values = [-999_999_999i32, 0, 999_999_999]; + let child = DecimalArray::from_iter(values, DecimalDType::new(9, 2)); + let child_ptr = child.buffer::().as_ptr(); + let logical_dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable); + let widen = widen_over(child, logical_dtype.clone())?; + let serialization = SESSION + .array_serialize(&widen)? + .vortex_expect("WidenPlugin serializes a decimal Widen over a Decimal leaf"); + assert_eq!(serialization.buffers[0].as_ptr(), child_ptr.cast::()); + + let ctx = ArrayContext::empty(); + let serialized = SerializedArray::try_from(serialize_blob(&widen, &ctx, &SESSION)?)?; + let read_ctx = ReadContext::new(ctx.to_ids()); + + assert_eq!( + read_ctx.resolve(serialized.encoding_id()), + Some(ArrayVTable::id(&Decimal)) + ); + assert_eq!(serialized.nbuffers(), 1); + let buffer = serialized.buffer(0)?; + assert_eq!(buffer.len(), 4 * widen.len()); + let expected_bytes = values + .into_iter() + .flat_map(i32::to_le_bytes) + .collect::>(); + assert_eq!(buffer.as_host().as_slice(), expected_bytes); + let decoded = serialized.decode(widen.dtype(), widen.len(), &read_ctx, &SESSION)?; + assert_eq!(decoded.dtype(), &logical_dtype); + assert_arrays_eq!(decoded, widen, &mut SESSION.create_execution_ctx()); + + Ok(()) +} + +#[rstest] +#[case::nullable(DecimalArray::from_option_iter( + [Some(-999_999_999i32), None, Some(42)], DecimalDType::new(9, 2), +), 38)] +#[case::all_null(DecimalArray::new( + buffer![i32::MIN, i32::MAX], DecimalDType::new(9, 2), Validity::AllInvalid, +), 38)] +#[case::empty(DecimalArray::from_iter(Vec::::new(), DecimalDType::new(9, 2)), 38)] +#[case::i128_child(DecimalArray::from_iter( + [i128::MIN / 2, 0, i128::MAX / 2], DecimalDType::new(38, 0), +), 76)] +fn decimal_widen_round_trip( + #[case] child: DecimalArray, + #[case] precision: u8, +) -> VortexResult<()> { + let scale = child.decimal_dtype().scale(); + let widen = widen_over( + child, + DType::Decimal( + DecimalDType::new(precision, scale), + Nullability::NonNullable, + ), + )?; + let (decoded, ids) = serialize_round_trip(&widen, &SESSION)?; + + assert!(!ids.contains(&ArrayVTable::id(&Widen))); + assert_arrays_eq!(decoded, widen, &mut SESSION.create_execution_ctx()); + + Ok(()) +} + +#[test] +fn decimal_widen_nested_round_trip() -> VortexResult<()> { + let child = DecimalArray::from_option_iter( + [Some(-999_999_999i32), None, Some(42)], + DecimalDType::new(9, 2), + ); + let widen = widen_over( + child, + DType::Decimal(DecimalDType::new(38, 2), Nullability::Nullable), + )?; + let nested = StructArray::from_fields(&[("decimal", widen.clone())])?.into_array(); + let chunked = ChunkedArray::try_new(vec![widen.clone(), widen.clone()], widen.dtype().clone())? + .into_array(); + + for array in [nested, chunked] { + let (decoded, ids) = serialize_round_trip(&array, &SESSION)?; + assert!(!ids.contains(&ArrayVTable::id(&Widen))); + assert_arrays_eq!(decoded, array, &mut SESSION.create_execution_ctx()); + } + + Ok(()) +} + +#[test] +fn serialization_keeps_parent_statistics() -> VortexResult<()> { + let child = DecimalArray::from_iter([-3i32, 5, 1], DecimalDType::new(9, 2)); + let dtype = DecimalDType::new(38, 2); + let widen = widen_over(child, DType::Decimal(dtype, Nullability::NonNullable))?; + let min = Scalar::decimal((-3i128).into(), dtype, Nullability::NonNullable); + let max = Scalar::decimal(5i128.into(), dtype, Nullability::NonNullable); + + for (stat, value) in [(Stat::Min, &min), (Stat::Max, &max)] { + let value = value + .value() + .cloned() + .vortex_expect("extrema are non-null decimal scalars"); + widen.statistics().set(stat, Precision::Exact(value)); + } + + let (decoded, _) = serialize_round_trip(&widen, &SESSION)?; + assert_eq!(decoded.statistics().get(Stat::Min), Precision::Exact(min)); + assert_eq!(decoded.statistics().get(Stat::Max), Precision::Exact(max)); + + Ok(()) +} + +#[test] +fn serialization_requires_a_ready_decimal_leaf() -> VortexResult<()> { + let child = DecimalArray::from_iter([1i32, 2, 3], DecimalDType::new(9, 2)).into_array(); + let dict = DictArray::try_new(buffer![0u8, 1, 2].into_array(), child)?; + let widen = widen_over( + dict, + DType::Decimal(DecimalDType::new(38, 2), Nullability::NonNullable), + )?; + let error = serialize_blob(&widen, &ArrayContext::empty(), &SESSION) + .expect_err("serialization requires a ready, serializable Decimal leaf"); + assert!(error.to_string().contains("requires a Decimal leaf")); + + let integer = widen_over(buffer![1i8, 2, 3], DType::from(PType::I64))?; + let error = serialize_blob(&integer, &ArrayContext::empty(), &SESSION) + .expect_err("WidenPlugin only serializes decimal arrays"); + assert!(error.to_string().contains("requires a decimal dtype")); + + Ok(()) +} + +#[test] +fn serialization_does_not_execute_validity() -> VortexResult<()> { + let validity = + SliceArray::new(BoolArray::from_iter([true, false, true]).into_array(), 0..2).into_array(); + let child = DecimalArray::new( + buffer![1i32, 2], + DecimalDType::new(9, 2), + Validity::Array(validity), + ); + let widen = widen_over( + child, + DType::Decimal(DecimalDType::new(38, 2), Nullability::Nullable), + )?; + let error = serialize_blob(&widen, &ArrayContext::empty(), &SESSION) + .expect_err("serialization requires a ready, serializable Decimal leaf"); + assert!(error.to_string().contains("vortex.slice")); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/vtable/mod.rs b/vortex-array/src/arrays/widen/vtable/mod.rs index 3af2178b523..6d8dec5ea98 100644 --- a/vortex-array/src/arrays/widen/vtable/mod.rs +++ b/vortex-array/src/arrays/widen/vtable/mod.rs @@ -1,7 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Vtable and wire adapter for the in-memory `Widen` encoding. +//! +//! Execution restores the logical width. The writer plugin uses an existing wire form because +//! `Widen` has no wire form of its own. + mod operations; + +mod plugin; +pub use plugin::WidenPlugin; + mod validity; use vortex_error::VortexResult; @@ -24,6 +33,10 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityVTableFromChild; use crate::array::with_empty_buffers; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::copy_to_canonical_width; use crate::arrays::widen::WidenArraySlotsExt; use crate::arrays::widen::WidenSlots; use crate::arrays::widen::array::validate_dtypes; @@ -31,13 +44,16 @@ use crate::arrays::widen::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::match_each_decimal_value_type; +use crate::require_child; use crate::serde::ArrayChildren; /// The encoding of a [`WidenArray`]. #[derive(Clone, Debug)] pub struct Widen; -/// A [`Widen`]-encoded integer array. +/// A [`Widen`]-encoded integer or decimal array. pub type WidenArray = Array; impl VTable for Widen { @@ -91,9 +107,7 @@ impl VTable for Widen { _array: ArrayView<'_, Self>, _session: &VortexSession, ) -> VortexResult>> { - // `ArraySession` does not register `Widen`, so `array_serialize` fails before it calls - // this method. Only a session that registers `Widen` itself can reach it. - vortex_panic!("TODO(connor)[Widen]: Widen has no wire form yet (#10733)") + Ok(None) } fn deserialize( @@ -112,10 +126,28 @@ impl VTable for Widen { WidenSlots::NAMES[idx].to_string() } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - // Yield a cast, not a canonical array, so that one kernel decodes and widens the values. - let widened = array.values().cast(array.dtype().clone())?; - Ok(ExecutionResult::done(widened)) + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + match array.dtype() { + DType::Decimal(decimal_dtype, _) => { + // The result must use the logical decimal dtype's canonical buffer width, + // independently of the child's stored width. + let decimal_dtype = *decimal_dtype; + let array = require_child!(array, array.values(), WidenSlots::VALUES => Decimal); + let widened = + copy_decimal_values(array.values().as_::(), decimal_dtype, ctx)?; + + Ok(ExecutionResult::done(widened.into_array())) + } + DType::Primitive(..) => { + // Integer casts let the child's kernel decode and widen its values in one pass. + let widened = array.values().cast(array.dtype().clone())?; + + Ok(ExecutionResult::done(widened)) + } + dtype => { + vortex_panic!("Widen validation requires an integer or decimal dtype, got {dtype}") + } + } } fn reduce(array: ArrayView<'_, Self>) -> VortexResult> { @@ -136,3 +168,24 @@ impl VTable for Widen { PARENT_RULES.evaluate(array, parent, child_idx) } } + +/// Copies the executed decimal `values` of a [`WidenArray`] into the canonical width of +/// `decimal_dtype`, the logical dtype of the `Widen`. +fn copy_decimal_values( + values: ArrayView<'_, Decimal>, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values.validity()?; + let mask = validity.execute_mask(values.len(), ctx)?; + let buffer_handle = match_each_decimal_value_type!(values.values_type(), |T| { + copy_to_canonical_width(values.buffer::().as_slice(), &mask, decimal_dtype)? + }); + + Ok(DecimalArray::new_handle( + buffer_handle, + decimal_dtype.canonical_values_type(), + decimal_dtype, + validity, + )) +} diff --git a/vortex-array/src/arrays/widen/vtable/plugin.rs b/vortex-array/src/arrays/widen/vtable/plugin.rs new file mode 100644 index 00000000000..bc2a92e42a4 --- /dev/null +++ b/vortex-array/src/arrays/widen/vtable/plugin.rs @@ -0,0 +1,86 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Serialization of decimal `Widen` arrays as narrow `vortex.decimal` nodes. +//! +//! The node stores the child's buffer width, while its containing format supplies the logical dtype +//! of the `Widen`. No `Widen` node or child decimal dtype is written. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use crate::ArrayDeserialization; +use crate::ArrayId; +use crate::ArrayPlugin; +use crate::ArrayRef; +use crate::ArraySerialization; +use crate::ArrayVTable; +use crate::arrays::Decimal; +use crate::arrays::Widen; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::session::ArraySessionExt; + +/// Writes a decimal [`Widen`] as one `vortex.decimal` node with its child's values buffer. +/// +/// The old decimal wire format allows a buffer narrower than the logical dtype's canonical width. +/// For example, a `Widen` of `Decimal(38, 2)` over a `Decimal(9, 2)` child writes the child's `i32` +/// buffer while keeping the logical dtype `Decimal(38, 2)`. +/// +/// The child **must** be a [`Decimal`] leaf. Its serializer supplies width metadata, buffers and +/// validity children through [`ArraySerialization`], which contains no dtype. The containing format +/// supplies the original `Widen` dtype and length, and array serialization keeps its statistics. +/// +/// This plugin registers only a writer. The decimal reader handles the resulting node. Serialization +/// executes neither values nor validity, so integer arrays, non-decimal children and validity that +/// cannot be serialized require preparation by the caller. +#[derive(Clone, Debug)] +pub struct WidenPlugin; + +impl ArrayPlugin for WidenPlugin { + fn id(&self) -> ArrayId { + ArrayVTable::id(&Widen) + } + + fn read_ids(&self) -> Vec { + vec![] + } + + fn write_ids(&self) -> Vec { + vec![ArrayVTable::id(&Decimal)] + } + + fn serialize( + &self, + array: &ArrayRef, + session: &VortexSession, + ) -> VortexResult> { + let widen = array.as_::(); + if !array.dtype().is_decimal() { + vortex_bail!( + "Widen serialization requires a decimal dtype, got {}", + array.dtype() + ); + } + + let values = widen.values(); + if !values.is::() { + vortex_bail!( + "Decimal Widen serialization requires a Decimal leaf, got {}", + values.encoding_id() + ); + } + + // Only the child's wire components are returned, not its dtype. The enclosing format and + // serialization tree still use the original Widen's logical dtype and statistics. + session.array_serialize(values) + } + + fn deserialize( + &self, + _parts: ArrayDeserialization<'_>, + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!("Widen has no wire form yet (#10733)") + } +} diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 9bc6a3d8c4c..96be71b0454 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -44,6 +44,7 @@ use crate::arrays::Union; use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; +use crate::arrays::widen::WidenPlugin; /// Registry of array encodings. pub type ArrayRegistry = ArcSwapMap; @@ -115,6 +116,7 @@ impl Default for ArraySession { this.register(PiecewiseSequence); this.register(RunEnd); this.register(VarBin); + this.register(WidenPlugin); this } diff --git a/vortex-ipc/src/messages/decoder.rs b/vortex-ipc/src/messages/decoder.rs index 73cdd235485..d6168ac558f 100644 --- a/vortex-ipc/src/messages/decoder.rs +++ b/vortex-ipc/src/messages/decoder.rs @@ -167,14 +167,23 @@ impl MessageDecoder { } #[cfg(test)] -mod test { +mod tests { use bytes::BytesMut; use vortex_array::ArrayRef; + use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::Widen; + use vortex_array::arrays::WidenArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_error::vortex_panic; use super::*; @@ -220,4 +229,33 @@ mod test { assert_eq!(array.nbuffers(), 1, "Array should have a single buffer"); write_and_read(&array.into_array()); } + + #[test] + fn array_ipc_writes_decimal_widen() -> VortexResult<()> { + let values = DecimalArray::new( + buffer![-125i32, 50, 900], + DecimalDType::new(9, 2), + Validity::from_iter([true, false, true]), + ); + let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::Nullable); + let array = WidenArray::try_new(values.into_array(), dtype)?.into_array(); + + let mut ipc_bytes = BytesMut::new(); + let mut encoder = MessageEncoder::new(SESSION.clone()); + for buf in encoder.encode(EncoderMessage::Array(&array))? { + ipc_bytes.extend_from_slice(buf.as_ref()); + } + + let (array_parts, ctx, row_count) = + match MessageDecoder::default().read_next(&mut ipc_bytes)? { + PollRead::Some(DecoderMessage::Array(array_parts)) => array_parts, + otherwise => vortex_panic!("Expected an array, got {:?}", otherwise), + }; + let actual = array_parts.decode(array.dtype(), row_count, &ctx, &SESSION)?; + + assert!(!ctx.ids().contains(&ArrayVTable::id(&Widen))); + assert_arrays_eq!(actual, array, &mut SESSION.create_execution_ctx()); + + Ok(()) + } } diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index 6cf043fdc34..1228c06f61c 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -176,16 +176,22 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::ArrayRef; + use vortex_array::ArrayVTable; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; + use vortex_array::arrays::Widen; + use vortex_array::arrays::WidenArray; use vortex_array::arrays::struct_::StructArrayExt; + use vortex_array::assert_arrays_eq; use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; @@ -197,6 +203,7 @@ mod tests { use vortex_buffer::BitBufferMut; use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::AllOr; @@ -404,4 +411,47 @@ mod tests { assert_eq!(field_b.as_slice::(), &[3, 4]); }) } + + #[test] + fn flat_writes_decimal_widen() -> VortexResult<()> { + let values = DecimalArray::new( + buffer![-125i32, 50, 900], + DecimalDType::new(9, 2), + Validity::from_iter([true, false, true]), + ); + let dtype = DType::Decimal(DecimalDType::new(38, 2), Nullability::Nullable); + let array = WidenArray::try_new(values.into_array(), dtype)?.into_array(); + + block_on(|handle| async { + let mut ctx_exec = SESSION.create_execution_ctx(); + let session = new_session().with_handle(handle); + let ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let layout = FlatLayoutStrategy::default() + .write_stream( + ctx.clone().into(), + Arc::::clone(&segments), + array.clone().to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; + let expr = root().bind(reader.dtype())?; + let result = reader + .projection_evaluation( + &(0..layout.row_count()), + &expr, + MaskFuture::new_true(usize::try_from(layout.row_count())?), + )? + .await?; + + assert!(!ctx.to_ids().contains(&ArrayVTable::id(&Widen))); + assert_arrays_eq!(result, array, &mut ctx_exec); + + Ok(()) + }) + } } From 95d6c868932c6f1c73d74729df2485337efc626e Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 3 Nov 2026 08:10:09 -0500 Subject: [PATCH 2/3] Keep decimal Widen benchmarks focused and small Signed-off-by: Connor Tsui --- vortex-array/benches/widen_decimal.rs | 118 +++++--------------------- 1 file changed, 19 insertions(+), 99 deletions(-) diff --git a/vortex-array/benches/widen_decimal.rs b/vortex-array/benches/widen_decimal.rs index 699043a02fe..2f7584a9b5c 100644 --- a/vortex-array/benches/widen_decimal.rs +++ b/vortex-array/benches/widen_decimal.rs @@ -3,11 +3,11 @@ //! Compute on decimal `Widen` arrays and canonical arrays with the same values. //! -//! The cases cover precision 18 over precision 9, precision 10 over precision 5, and precision 38 -//! over precision 18, all at scale 2. The precision 5 values fit in `i16`, but require `i32` to -//! represent their full decimal domain. Each iteration shares value buffers and builds fresh arrays -//! so cached statistics cannot remove the operation's work. `Widen` buffers are selected once with -//! [`WidenArray::encode_decimal`], outside the timed region. +//! Execution measures the widening copy from `i32` to `i64` and from `i64` to `i128`. Filtering +//! compares each `Widen` with its canonical input, since it can filter the narrow child first. +//! Canonical execution is omitted because it only returns the existing array. Each iteration shares +//! value buffers and builds fresh arrays so cached statistics cannot remove the operation's work. +//! [`WidenArray::encode_decimal`] selects the child buffers outside the timed region. use std::fmt; use std::sync::LazyLock; @@ -22,25 +22,18 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::NumericalAggregateOpts; -use vortex_array::aggregate_fn::fns::min_max; -use vortex_array::aggregate_fn::fns::sum; use vortex_array::aggregate_fn::session::AggregateFnSession; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Decimal; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::Widen; use vortex_array::arrays::WidenArray; use vortex_array::arrays::widen::WidenArraySlotsExt; -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::memory::MemorySession; use vortex_array::optimizer::kernels::KernelSession; -use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; @@ -75,7 +68,8 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::() }); -const LEN: usize = 32_768; +// Keep per-row work below 1 ms in CodSpeed CPU simulation, including later kernels. +const LEN: usize = 4_096; const SCALE: i8 = 2; #[derive(Clone, Copy, Debug)] @@ -91,11 +85,6 @@ const CASES: &[Case] = &[ child_precision: 9, input: Input::Canonical, }, - Case { - precision: 10, - child_precision: 5, - input: Input::Canonical, - }, Case { precision: 38, child_precision: 18, @@ -124,6 +113,16 @@ fn cases() -> Vec { .collect() } +fn widen_cases() -> Vec { + CASES + .iter() + .map(|&case| Case { + input: Input::Widen, + ..case + }) + .collect() +} + impl fmt::Display for Case { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!( @@ -146,10 +145,9 @@ impl Case { fn values(self) -> Values { let (minimum, maximum) = match self.child_precision { - 5 => (10_000, 32_767), 9 => (-999_999_999, 999_999_999), 18 => (-999_999_999_999_999_999, 999_999_999_999_999_999), - _ => unreachable!("CASES contains only child precisions 5, 9, and 18"), + _ => unreachable!("CASES contains only child precisions 9 and 18"), }; let mut rng = StdRng::seed_from_u64(0); @@ -187,31 +185,6 @@ impl Case { _ => unreachable!("the benchmark child precisions require i32 or i64"), } } - - fn constant(self, value: i64) -> ArrayRef { - let scalar = Scalar::decimal(value.into(), self.decimal_dtype(), Nullability::NonNullable); - ConstantArray::new(scalar, LEN).into_array() - } - - /// Returns a constant that splits the values, so `compare_constant` is not uniform. - /// - /// The precision 5 values lie in `10_000..=32_767`, so 0 would be below all of them. The other - /// ranges are symmetric about zero, so 0 splits them. - fn in_range(self) -> i64 { - if self.child_precision == 5 { 20_000 } else { 0 } - } - - fn out_of_range(self) -> i64 { - 10i64.pow(u32::from(self.child_precision)) - } - - fn cast_target(self) -> DType { - let precision = if self.precision == 38 { 76 } else { 38 }; - DType::Decimal( - DecimalDType::new(precision, SCALE), - Nullability::NonNullable, - ) - } } enum Values { @@ -261,49 +234,6 @@ fn execute_canonical(array: VortexResult, ctx: &mut ExecutionCtx) -> C .vortex_expect("the operation executes to canonical for the benchmark inputs") } -#[divan::bench(args = cases())] -fn compare_constant(bencher: Bencher, case: Case) { - let values = case.values(); - bench_case( - bencher, - || (values.array(case), case.constant(case.in_range())), - |(array, constant), ctx| execute_canonical(array.binary(constant, Operator::Lt), ctx), - ); -} - -#[divan::bench(args = cases())] -fn compare_constant_out_of_range(bencher: Bencher, case: Case) { - let values = case.values(); - bench_case( - bencher, - || (values.array(case), case.constant(case.out_of_range())), - |(array, constant), ctx| execute_canonical(array.binary(constant, Operator::Lt), ctx), - ); -} - -#[divan::bench(args = cases())] -fn sum(bencher: Bencher, case: Case) { - let values = case.values(); - bench_case( - bencher, - || values.array(case), - |array, ctx| sum::sum(&array, ctx).vortex_expect("the benchmark decimal sum is valid"), - ); -} - -#[divan::bench(args = cases())] -fn min_max(bencher: Bencher, case: Case) { - let values = case.values(); - bench_case( - bencher, - || values.array(case), - |array, ctx| { - min_max::min_max(&array, ctx, NumericalAggregateOpts::default()) - .vortex_expect("the benchmark decimal extrema are valid") - }, - ); -} - #[divan::bench(args = cases())] fn filter(bencher: Bencher, case: Case) { let values = case.values(); @@ -316,7 +246,7 @@ fn filter(bencher: Bencher, case: Case) { ); } -#[divan::bench(args = cases())] +#[divan::bench(args = widen_cases())] fn execute(bencher: Bencher, case: Case) { let values = case.values(); bench_case( @@ -325,13 +255,3 @@ fn execute(bencher: Bencher, case: Case) { |array, ctx| execute_canonical(Ok(array), ctx), ); } - -#[divan::bench(args = cases())] -fn cast_widen(bencher: Bencher, case: Case) { - let values = case.values(); - bench_case( - bencher, - || values.array(case), - |array, ctx| execute_canonical(array.cast(case.cast_target()), ctx), - ); -} From 732851062010ad642bf64fc2c3c99f22b7a75b7a Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 3 Nov 2026 08:35:31 -0500 Subject: [PATCH 3/3] Update CodSpeed counts for the decimal benchmark cases Signed-off-by: Connor Tsui --- scripts/codspeed_bench_counts.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/codspeed_bench_counts.json b/scripts/codspeed_bench_counts.json index 00ced92f623..c356c1d946f 100644 --- a/scripts/codspeed_bench_counts.json +++ b/scripts/codspeed_bench_counts.json @@ -75,7 +75,7 @@ "varbinview_try_new": 6, "varbinview_zip": 2, "widen": 46, - "widen_decimal": 42 + "widen_decimal": 6 }, "vortex-buffer": { "allocation": 35,