From 3ce4851eb2e5962a9d359711def813fdd8b5316a Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:22:54 -0400 Subject: [PATCH 01/12] feat: execute interleaved decimal arrays Signed-off-by: Connor Tsui --- .../arrays/interleave/execute/decimal/mod.rs | 113 +++++++ .../interleave/execute/decimal/tests.rs | 284 ++++++++++++++++++ .../arrays/interleave/execute/fixed_width.rs | 91 ++++++ .../src/arrays/interleave/execute/mod.rs | 4 + .../arrays/interleave/execute/primitive.rs | 84 +----- 5 files changed, 496 insertions(+), 80 deletions(-) create mode 100644 vortex-array/src/arrays/interleave/execute/decimal/mod.rs create mode 100644 vortex-array/src/arrays/interleave/execute/decimal/tests.rs create mode 100644 vortex-array/src/arrays/interleave/execute/fixed_width.rs diff --git a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs new file mode 100644 index 00000000000..ac69429a187 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution for decimal [`Interleave`](super::super::Interleave) values. +//! +//! Sources can use different stored widths for the same logical dtype. Decimal execution +//! normalizes them to the logical canonical width before gathering rows. + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use super::super::InterleaveSlots; +use super::fixed_width::Source; +use super::fixed_width::gather; +use crate::AnyColumnar; +use crate::array::Array; +use crate::arrays::Constant; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Primitive; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::copy_to_canonical_width; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; +use crate::match_each_decimal_value_type; +use crate::match_each_unsigned_integer_ptype; +use crate::require_child; + +pub(super) fn execute( + mut array: Array, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let num_values = array.num_values(); + array = + require_child!(array, array.array_indices(), InterleaveSlots::ARRAY_INDICES => Primitive); + array = require_child!(array, array.row_indices(), InterleaveSlots::ROW_INDICES => Primitive); + for i in 0..num_values { + array = require_child!(array, array.value(i), InterleaveSlots::VALUES_OFFSET + i => AnyColumnar); + } + + let validity = array.as_ref().validity()?; + let dtype = DecimalDType::try_from(array.dtype())?; + let output = match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + DecimalArray::new(gather_values::(&array, dtype, ctx)?, dtype, validity) + }); + + Ok(ExecutionResult::done(output)) +} + +fn gather_values( + array: &Array, + dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let allocator = ctx.allocator().clone(); + let values = (0..array.num_values()) + .map(|i| { + let value = array.value(i); + let len = value.len(); + let (data, row_mask) = if let Some(constant) = value.as_opt::() { + let payload = match constant.scalar().as_decimal().decimal_value() { + Some(value) => value.cast::().ok_or_else(|| { + vortex_err!("interleave decimal constant must fit {dtype}, got {value}") + })?, + None => T::default(), + }; + (Buffer::full_in(payload, 1, allocator.clone()), 0) + } else { + let decimal = value.as_::(); + let data = if decimal.values_type() == T::DECIMAL_TYPE { + decimal.buffer::() + } else { + // A decimal source can have noncanonical storage. Normalize before indexing + // rows so that every source has the same physical width. A narrowing copy + // checks the precision and zeros arbitrary null payloads. + let mask = decimal.validity()?.execute_mask(len, ctx)?; + let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { + copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? + }); + Buffer::::from_byte_buffer(copy.unwrap_host()) + }; + (data, usize::MAX) + }; + + Ok(Source { + data, + len, + row_mask, + }) + }) + .collect::>>()?; + let branches = array.array_indices().as_::(); + let rows = array.row_indices().as_::(); + + match_each_unsigned_integer_ptype!(branches.ptype(), |A| { + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather( + &values, + branches.as_slice::(), + rows.as_slice::(), + &allocator, + ) + }) + }) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/arrays/interleave/execute/decimal/tests.rs b/vortex-array/src/arrays/interleave/execute/decimal/tests.rs new file mode 100644 index 00000000000..8951d5e4791 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/decimal/tests.rs @@ -0,0 +1,284 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::InterleaveArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::SliceArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::assert_arrays_eq; +use crate::dtype::BigCast; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::validity::Validity; + +/// Executes the gather, checks its physical width and dtype, and compares it with a canonical oracle. +fn check(values: Vec, indices: &[(u8, u64)], expected: DecimalArray) -> VortexResult<()> { + let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch)); + let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row)); + let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?; + let mut ctx = VortexSession::empty().create_execution_ctx(); + let output = array.into_array().execute::(&mut ctx)?; + assert_eq!( + output.values_type(), + expected.decimal_dtype().canonical_values_type() + ); + assert_eq!(output.dtype(), expected.dtype()); + assert_arrays_eq!(output.into_array(), expected.into_array(), &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::i8(2)] +#[case::i16(4)] +#[case::i32(9)] +#[case::i64(18)] +#[case::i128(38)] +#[case::i256(76)] +fn test_decimal_interleave( + #[case] precision: u8, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, + #[values(-3, 0, 2)] scale: i8, +) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, scale); + let (first_validity, output_validity) = if nullability.is_nullable() { + ( + Validity::from_iter([true, false, true]), + Validity::from_iter([true, true, false, true, true, true]), + ) + } else { + (Validity::NonNullable, Validity::NonNullable) + }; + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + let min = T::MIN_BY_PRECISION[usize::from(precision)]; + let max = T::MAX_BY_PRECISION[usize::from(precision)]; + let zero = T::default(); + let middle = if nullability.is_nullable() { + T::MIN + } else { + zero + }; + let first = DecimalArray::new(Buffer::copy_from([min, middle, max]), dtype, first_validity); + let second = DecimalArray::new( + Buffer::copy_from([max, min, zero]), + dtype, + Validity::NonNullable, + ); + let expected = DecimalArray::new( + Buffer::copy_from([zero, min, zero, max, max, min]), + dtype, + output_validity, + ); + check( + vec![first.into_array(), second.into_array()], + &[(1, 2), (0, 0), (0, 1), (1, 0), (0, 2), (0, 0)], + expected, + ) + }) +} + +#[rstest] +#[case::i8(2)] +#[case::i16(4)] +#[case::i32(9)] +#[case::i64(18)] +#[case::i128(38)] +#[case::i256(76)] +fn test_decimal_interleave_constants( + #[case] precision: u8, + #[values(false, true)] null_constant: bool, + #[values(DecimalType::I8, DecimalType::I256)] scalar_storage: DecimalType, +) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 2); + let scalar = if null_constant { + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) + } else { + let payload = match_each_decimal_value_type!(scalar_storage, |S| { + DecimalValue::from( + ::from(-7i8).vortex_expect("-7 fits every decimal type"), + ) + }); + Scalar::decimal(payload, dtype, Nullability::NonNullable) + }; + let constant = ConstantArray::new(scalar, 3).into_array(); + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + let seven = ::from(7i8).vortex_expect("7 fits every decimal type"); + let minus_seven = ::from(-7i8).vortex_expect("-7 fits every decimal type"); + let column = DecimalArray::new(Buffer::copy_from([seven]), dtype, Validity::NonNullable) + .into_array(); + let expected = DecimalArray::new( + Buffer::copy_from([minus_seven, seven, minus_seven, minus_seven]), + dtype, + if null_constant { + Validity::from_iter([false, true, false, false]) + } else { + Validity::NonNullable + }, + ); + check( + vec![constant, column], + &[(0, 2), (1, 0), (0, 1), (0, 2)], + expected, + ) + }) +} + +#[rstest] +#[case::i8(2)] +#[case::i16(4)] +#[case::i32(9)] +#[case::i64(18)] +#[case::i128(38)] +#[case::i256(76)] +fn test_decimal_interleave_mixed_stored_widths(#[case] precision: u8) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 0); + let first = + DecimalArray::new(Buffer::copy_from([1i8, -2]), dtype, Validity::NonNullable).into_array(); + let second = DecimalArray::new( + Buffer::copy_from([i256::from_i128(3), i256::MIN, i256::MAX]), + dtype, + Validity::from_iter([true, false, false]), + ) + .into_array(); + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + let values = [3i8, -2, 0, 1, 0].map(|v| { + ::from(v).vortex_expect("small test values fit every decimal type") + }); + let expected = DecimalArray::new( + Buffer::copy_from(values), + dtype, + Validity::from_iter([true, true, false, true, false]), + ); + check( + vec![first, second], + &[(1, 0), (0, 1), (1, 1), (0, 0), (1, 2)], + expected, + ) + }) +} + +#[rstest] +#[case::empty(&[])] +#[case::all_null(&[(0, 0), (1, 1), (0, 1)])] +fn test_decimal_interleave_without_valid_values(#[case] indices: &[(u8, u64)]) -> VortexResult<()> { + let dtype = DecimalDType::new(9, -2); + let first = DecimalArray::new( + Buffer::copy_from([i256::MIN, i256::MAX]), + dtype, + Validity::AllInvalid, + ) + .into_array(); + let second = ConstantArray::new( + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), + 2, + ) + .into_array(); + let expected = DecimalArray::new( + Buffer::::zeroed(indices.len()), + dtype, + Validity::AllInvalid, + ); + check(vec![first, second], indices, expected) +} + +#[test] +fn test_decimal_interleave_empty_sources() -> VortexResult<()> { + let dtype = DecimalDType::new(9, 0); + let empty = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable).into_array(); + let expected = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable); + check(vec![empty.clone(), empty], &[], expected) +} + +#[test] +fn test_decimal_interleave_executes_lazy_children() -> VortexResult<()> { + let dtype = DecimalDType::new(2, 0); + let values = + DecimalArray::new(Buffer::copy_from([1i8, -2]), dtype, Validity::NonNullable).into_array(); + let lazy = SliceArray::try_new(values.clone(), 0..2)?.into_array(); + let branches = SliceArray::try_new(PrimitiveArray::from_iter([1u8, 0, 1]).into_array(), 0..3)? + .into_array(); + let rows = SliceArray::try_new(PrimitiveArray::from_iter([1u16, 0, 0]).into_array(), 0..3)? + .into_array(); + let interleaved = InterleaveArray::try_new(vec![values, lazy], branches, rows)?.into_array(); + let expected = DecimalArray::new( + Buffer::copy_from([-2i8, 1, 1]), + dtype, + Validity::NonNullable, + ) + .into_array(); + let mut ctx = VortexSession::empty().create_execution_ctx(); + assert_arrays_eq!(interleaved, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::array_index(2, 0, "array index out of bounds")] +#[case::row_index(0, 2, "row index out of bounds")] +#[case::constant_row_index(1, 2, "row index out of bounds")] +#[case::large_array_index(u64::MAX, 0, "array index out of bounds")] +fn test_decimal_interleave_rejects_out_of_bounds_selectors( + #[case] branch: u64, + #[case] row: u8, + #[case] message: &str, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, +) -> VortexResult<()> { + let dtype = DecimalDType::new(2, 0); + let values = DecimalArray::new( + Buffer::copy_from([1i8, -2]), + dtype, + Validity::from(nullability), + ) + .into_array(); + let constant = + ConstantArray::new(Scalar::decimal(1i8.into(), dtype, nullability), 2).into_array(); + let array = InterleaveArray::try_new( + vec![values, constant], + PrimitiveArray::from_iter([branch]).into_array(), + PrimitiveArray::from_iter([row]).into_array(), + )? + .into_array(); + let error = array + .execute::(&mut VortexSession::empty().create_execution_ctx()) + .unwrap_err(); + assert!(error.to_string().contains(message), "{error}"); + + Ok(()) +} + +#[test] +fn test_decimal_interleave_rejects_wide_values_outside_precision() -> VortexResult<()> { + let dtype = DecimalDType::new(2, 0); + let source = + DecimalArray::new(Buffer::copy_from([100i16]), dtype, Validity::NonNullable).into_array(); + let array = InterleaveArray::try_new( + vec![source.clone(), source], + PrimitiveArray::from_iter([0u8]).into_array(), + PrimitiveArray::from_iter([0u8]).into_array(), + )? + .into_array(); + let error = array + .execute::(&mut VortexSession::empty().create_execution_ctx()) + .unwrap_err(); + assert!(error.to_string().contains("precision"), "{error}"); + + Ok(()) +} diff --git a/vortex-array/src/arrays/interleave/execute/fixed_width.rs b/vortex-array/src/arrays/interleave/execute/fixed_width.rs new file mode 100644 index 00000000000..2f21a408c26 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/fixed_width.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared gather kernel for fixed-width [`Interleave`](super::super::Interleave) values. +//! +//! Each source holds a decoded buffer or one constant payload. Selector bounds are checked against +//! the logical source length before the kernel maps a row to its physical payload. + +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_buffer::BufferAllocatorRef; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +/// Physical storage for one fixed-width source array. +pub(super) struct Source { + /// One constant payload or a decoded payload for each logical row. + pub(super) data: Buffer, + /// The logical row count, which can exceed `data.len()` for a constant. + pub(super) len: usize, + /// Maps rows to `data`: zero for constants and `usize::MAX` for decoded buffers. + pub(super) row_mask: usize, +} + +pub(super) fn gather( + values: &[Source], + branches: &[A], + rows: &[R], + allocator: &BufferAllocatorRef, +) -> VortexResult> +where + T: Copy, + A: AsPrimitive, + R: AsPrimitive, +{ + // `zip` truncates to the shorter input. + vortex_ensure_eq!( + rows.len(), + branches.len(), + "interleave row_indices length does not match array_indices length", + ); + + // Write rows through a `&mut` slice. Yielding `Result` per row splits each `[u8; W]` copy + // into narrow moves (the `Ok` payload is misaligned), and pushing through `BufferMut` reloads + // its length on every row. + let len = branches.len(); + let mut output = BufferMut::with_capacity_in(len, allocator.clone()); + let slots = &mut output.spare_capacity_mut()[..len]; + for ((slot, branch), row) in slots.iter_mut().zip(branches).zip(rows) { + let Some(source) = values.get((*branch).as_()) else { + vortex_bail!("interleave array index out of bounds"); + }; + let row = (*row).as_(); + vortex_ensure!(row < source.len, "interleave row index out of bounds"); + slot.write(source.data[row & source.row_mask]); + } + // SAFETY: the loop returns on any error, so reaching here means it wrote all `len` slots: + // `slots`, `branches`, and `rows` all have length `len`. + unsafe { output.set_len(len) }; + Ok(output.freeze()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory::BufferAllocatorRef; + + #[test] + fn rejects_out_of_bounds_selectors() { + let values = [Source { + data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()), + len: 1, + row_mask: 0, + }]; + + let allocator = BufferAllocatorRef::statically_allocated(); + let error = gather(&values, &[1u8], &[0u8], &allocator).unwrap_err(); + assert!( + error.to_string().contains("array index out of bounds"), + "{error}" + ); + let error = gather(&values, &[0u8], &[1u8], &allocator).unwrap_err(); + assert!( + error.to_string().contains("row index out of bounds"), + "{error}" + ); + } +} diff --git a/vortex-array/src/arrays/interleave/execute/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index 558db9e3309..20133f64cfc 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -10,6 +10,8 @@ //! [`Interleave::check`]: super::Interleave::check mod bool; +mod decimal; +mod fixed_width; mod primitive; use vortex_error::VortexResult; @@ -27,6 +29,8 @@ pub(super) fn execute( ) -> VortexResult { if array.dtype().is_boolean() { bool::execute(array, ctx) + } else if array.dtype().is_decimal() { + decimal::execute(array, ctx) } else if array.dtype().is_primitive() { primitive::execute(array, ctx) } else { diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index f1595045383..85f7d87ccd8 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -3,19 +3,16 @@ //! Execution for primitive [`Interleave`] values. -use num_traits::AsPrimitive; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_ensure_eq; use super::super::Interleave; use super::super::InterleaveArrayExt; use super::super::InterleaveSlots; +use super::fixed_width::Source; +use super::fixed_width::gather; use crate::AnyColumnar; use crate::array::Array; use crate::arrays::Constant; @@ -57,14 +54,6 @@ pub(super) fn execute( Ok(ExecutionResult::done(output)) } -/// Physical storage for one primitive source array. -struct PrimitiveSource { - data: Buffer, - len: usize, - // Maps logical rows to `data`: zero for constants and identity for decoded buffers. - row_mask: usize, -} - fn gather_values( array: &Array, allocator: &BufferAllocatorRef, @@ -85,13 +74,13 @@ fn gather_values( bytes }) .unwrap_or([0; W]); - PrimitiveSource { + Source { data: Buffer::full_in(payload, 1, allocator.clone()), len, row_mask: 0, } } else { - PrimitiveSource { + Source { data: Buffer::<[u8; W]>::from_byte_buffer( value.as_::().buffer_handle().to_host_sync(), ), @@ -115,68 +104,3 @@ fn gather_values( }) }) } - -fn gather( - values: &[PrimitiveSource], - branches: &[A], - rows: &[R], - allocator: &BufferAllocatorRef, -) -> VortexResult> -where - T: Copy, - A: AsPrimitive, - R: AsPrimitive, -{ - // `zip` truncates to the shorter input. - vortex_ensure_eq!( - rows.len(), - branches.len(), - "interleave row_indices length does not match array_indices length", - ); - - // Write rows through a `&mut` slice. Yielding `Result` per row splits each `[u8; W]` copy - // into narrow moves (the `Ok` payload is misaligned), and pushing through `BufferMut` reloads - // its length on every row. - let len = branches.len(); - let mut output = BufferMut::with_capacity_in(len, allocator.clone()); - let slots = &mut output.spare_capacity_mut()[..len]; - for ((slot, branch), row) in slots.iter_mut().zip(branches).zip(rows) { - let Some(source) = values.get((*branch).as_()) else { - vortex_bail!("interleave array index out of bounds"); - }; - let row = (*row).as_(); - vortex_ensure!(row < source.len, "interleave row index out of bounds"); - slot.write(source.data[row & source.row_mask]); - } - // SAFETY: the loop returns on any error, so reaching here means it wrote all `len` slots: - // `slots`, `branches`, and `rows` all have length `len`. - unsafe { output.set_len(len) }; - Ok(output.freeze()) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::memory::BufferAllocatorRef; - - #[test] - fn rejects_out_of_bounds_selectors() { - let values = [PrimitiveSource { - data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()), - len: 1, - row_mask: 0, - }]; - - let allocator = BufferAllocatorRef::statically_allocated(); - let error = gather(&values, &[1u8], &[0u8], &allocator).unwrap_err(); - assert!( - error.to_string().contains("array index out of bounds"), - "{error}" - ); - let error = gather(&values, &[0u8], &[1u8], &allocator).unwrap_err(); - assert!( - error.to_string().contains("row index out of bounds"), - "{error}" - ); - } -} From 03f3087b0e09e1139944c01457962af6eb1c672d Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:36:23 -0400 Subject: [PATCH 02/12] Clarify decimal Interleave source preparation Prepare each decimal source as one typed gather input, with an early return for constants. Keep buffer normalization separate from selector dispatch. Parameterize bounds-check cases and report assertion failures at their callers. Signed-off-by: Connor Tsui --- .../arrays/interleave/execute/decimal/mod.rs | 78 +++++++++++-------- .../interleave/execute/decimal/tests.rs | 17 ++-- .../arrays/interleave/execute/fixed_width.rs | 31 ++++---- 3 files changed, 73 insertions(+), 53 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs index ac69429a187..6df5bc4ba0e 100644 --- a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs @@ -16,6 +16,7 @@ use super::super::InterleaveSlots; use super::fixed_width::Source; use super::fixed_width::gather; use crate::AnyColumnar; +use crate::ArrayRef; use crate::array::Array; use crate::arrays::Constant; use crate::arrays::Decimal; @@ -59,40 +60,7 @@ fn gather_values( ) -> VortexResult> { let allocator = ctx.allocator().clone(); let values = (0..array.num_values()) - .map(|i| { - let value = array.value(i); - let len = value.len(); - let (data, row_mask) = if let Some(constant) = value.as_opt::() { - let payload = match constant.scalar().as_decimal().decimal_value() { - Some(value) => value.cast::().ok_or_else(|| { - vortex_err!("interleave decimal constant must fit {dtype}, got {value}") - })?, - None => T::default(), - }; - (Buffer::full_in(payload, 1, allocator.clone()), 0) - } else { - let decimal = value.as_::(); - let data = if decimal.values_type() == T::DECIMAL_TYPE { - decimal.buffer::() - } else { - // A decimal source can have noncanonical storage. Normalize before indexing - // rows so that every source has the same physical width. A narrowing copy - // checks the precision and zeros arbitrary null payloads. - let mask = decimal.validity()?.execute_mask(len, ctx)?; - let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { - copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? - }); - Buffer::::from_byte_buffer(copy.unwrap_host()) - }; - (data, usize::MAX) - }; - - Ok(Source { - data, - len, - row_mask, - }) - }) + .map(|i| prepare_source::(array.value(i), dtype, ctx)) .collect::>>()?; let branches = array.array_indices().as_::(); let rows = array.row_indices().as_::(); @@ -109,5 +77,47 @@ fn gather_values( }) } +/// Normalizes one decimal source while retaining a single payload for a constant. +fn prepare_source( + value: &ArrayRef, + dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let len = value.len(); + if let Some(constant) = value.as_opt::() { + let payload = match constant.scalar().as_decimal().decimal_value() { + Some(value) => value.cast::().ok_or_else(|| { + vortex_err!("interleave decimal constant must fit {dtype}, got {value}") + })?, + None => T::default(), + }; + + return Ok(Source { + data: Buffer::full_in(payload, 1, ctx.allocator().clone()), + len, + row_mask: 0, + }); + } + + let decimal = value.as_::(); + let data = if decimal.values_type() == T::DECIMAL_TYPE { + decimal.buffer::() + } else { + // Decimal sources can still have noncanonical storage. A narrowing copy must check the + // precision and zero null slots before the gather reads the buffer. + let mask = decimal.validity()?.execute_mask(len, ctx)?; + let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { + copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? + }); + Buffer::::from_byte_buffer(copy.unwrap_host()) + }; + + Ok(Source { + data, + len, + row_mask: usize::MAX, + }) +} + #[cfg(test)] mod tests; diff --git a/vortex-array/src/arrays/interleave/execute/decimal/tests.rs b/vortex-array/src/arrays/interleave/execute/decimal/tests.rs index 8951d5e4791..76f6f560842 100644 --- a/vortex-array/src/arrays/interleave/execute/decimal/tests.rs +++ b/vortex-array/src/arrays/interleave/execute/decimal/tests.rs @@ -30,7 +30,12 @@ use crate::scalar::Scalar; use crate::validity::Validity; /// Executes the gather, checks its physical width and dtype, and compares it with a canonical oracle. -fn check(values: Vec, indices: &[(u8, u64)], expected: DecimalArray) -> VortexResult<()> { +#[track_caller] +fn assert_interleaved( + values: Vec, + indices: &[(u8, u64)], + expected: DecimalArray, +) -> VortexResult<()> { let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch)); let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row)); let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?; @@ -87,7 +92,7 @@ fn test_decimal_interleave( dtype, output_validity, ); - check( + assert_interleaved( vec![first.into_array(), second.into_array()], &[(1, 2), (0, 0), (0, 1), (1, 0), (0, 2), (0, 0)], expected, @@ -133,7 +138,7 @@ fn test_decimal_interleave_constants( Validity::NonNullable }, ); - check( + assert_interleaved( vec![constant, column], &[(0, 2), (1, 0), (0, 1), (0, 2)], expected, @@ -167,7 +172,7 @@ fn test_decimal_interleave_mixed_stored_widths(#[case] precision: u8) -> VortexR dtype, Validity::from_iter([true, true, false, true, false]), ); - check( + assert_interleaved( vec![first, second], &[(1, 0), (0, 1), (1, 1), (0, 0), (1, 2)], expected, @@ -196,7 +201,7 @@ fn test_decimal_interleave_without_valid_values(#[case] indices: &[(u8, u64)]) - dtype, Validity::AllInvalid, ); - check(vec![first, second], indices, expected) + assert_interleaved(vec![first, second], indices, expected) } #[test] @@ -204,7 +209,7 @@ fn test_decimal_interleave_empty_sources() -> VortexResult<()> { let dtype = DecimalDType::new(9, 0); let empty = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable).into_array(); let expected = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable); - check(vec![empty.clone(), empty], &[], expected) + assert_interleaved(vec![empty.clone(), empty], &[], expected) } #[test] diff --git a/vortex-array/src/arrays/interleave/execute/fixed_width.rs b/vortex-array/src/arrays/interleave/execute/fixed_width.rs index 2f21a408c26..f1c154625ce 100644 --- a/vortex-array/src/arrays/interleave/execute/fixed_width.rs +++ b/vortex-array/src/arrays/interleave/execute/fixed_width.rs @@ -65,11 +65,22 @@ where #[cfg(test)] mod tests { - use super::*; + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_error::VortexResult; + + use super::Source; + use super::gather; use crate::memory::BufferAllocatorRef; - #[test] - fn rejects_out_of_bounds_selectors() { + #[rstest] + #[case::array_index(1, 0, "array index out of bounds")] + #[case::row_index(0, 1, "row index out of bounds")] + fn rejects_out_of_bounds_selectors( + #[case] branch: u8, + #[case] row: u8, + #[case] message: &str, + ) -> VortexResult<()> { let values = [Source { data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()), len: 1, @@ -77,15 +88,9 @@ mod tests { }]; let allocator = BufferAllocatorRef::statically_allocated(); - let error = gather(&values, &[1u8], &[0u8], &allocator).unwrap_err(); - assert!( - error.to_string().contains("array index out of bounds"), - "{error}" - ); - let error = gather(&values, &[0u8], &[1u8], &allocator).unwrap_err(); - assert!( - error.to_string().contains("row index out of bounds"), - "{error}" - ); + let error = gather(&values, &[branch], &[row], &allocator).unwrap_err(); + assert!(error.to_string().contains(message), "{error}"); + + Ok(()) } } From c0f945072b72aae01dd4e7b5563abf3fc1845734 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 15:51:13 -0400 Subject: [PATCH 03/12] Fix the decimal Interleave doc link and test session The module imports `Interleave`, so rustdoc rejects the explicit link target as redundant. The source comment now covers both the widening and the narrowing copy, and the tests share one module session. Signed-off-by: Connor Tsui --- .../src/arrays/interleave/execute/decimal/mod.rs | 7 ++++--- .../src/arrays/interleave/execute/decimal/tests.rs | 12 ++++++++---- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs index 6df5bc4ba0e..3acd2fb4557 100644 --- a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Execution for decimal [`Interleave`](super::super::Interleave) values. +//! Execution for decimal [`Interleave`] values. //! //! Sources can use different stored widths for the same logical dtype. Decimal execution //! normalizes them to the logical canonical width before gathering rows. @@ -103,8 +103,9 @@ fn prepare_source( let data = if decimal.values_type() == T::DECIMAL_TYPE { decimal.buffer::() } else { - // Decimal sources can still have noncanonical storage. A narrowing copy must check the - // precision and zero null slots before the gather reads the buffer. + // A source can still store its values at a width other than the canonical width of its + // dtype (#10701). The copy widens such values, or checks their precision and narrows them + // with zeros in the null slots. let mask = decimal.validity()?.execute_mask(len, ctx)?; let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? diff --git a/vortex-array/src/arrays/interleave/execute/decimal/tests.rs b/vortex-array/src/arrays/interleave/execute/decimal/tests.rs index 76f6f560842..1095e6966c9 100644 --- a/vortex-array/src/arrays/interleave/execute/decimal/tests.rs +++ b/vortex-array/src/arrays/interleave/execute/decimal/tests.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::LazyLock; + use rstest::rstest; use vortex_buffer::Buffer; use vortex_error::VortexExpect; @@ -29,6 +31,8 @@ use crate::scalar::DecimalValue; use crate::scalar::Scalar; use crate::validity::Validity; +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + /// Executes the gather, checks its physical width and dtype, and compares it with a canonical oracle. #[track_caller] fn assert_interleaved( @@ -39,7 +43,7 @@ fn assert_interleaved( let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch)); let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row)); let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?; - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); let output = array.into_array().execute::(&mut ctx)?; assert_eq!( output.values_type(), @@ -229,7 +233,7 @@ fn test_decimal_interleave_executes_lazy_children() -> VortexResult<()> { Validity::NonNullable, ) .into_array(); - let mut ctx = VortexSession::empty().create_execution_ctx(); + let mut ctx = SESSION.create_execution_ctx(); assert_arrays_eq!(interleaved, expected, &mut ctx); Ok(()) @@ -262,7 +266,7 @@ fn test_decimal_interleave_rejects_out_of_bounds_selectors( )? .into_array(); let error = array - .execute::(&mut VortexSession::empty().create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) .unwrap_err(); assert!(error.to_string().contains(message), "{error}"); @@ -281,7 +285,7 @@ fn test_decimal_interleave_rejects_wide_values_outside_precision() -> VortexResu )? .into_array(); let error = array - .execute::(&mut VortexSession::empty().create_execution_ctx()) + .execute::(&mut SESSION.create_execution_ctx()) .unwrap_err(); assert!(error.to_string().contains("precision"), "{error}"); From 5a1347fef712cf1bf407f0443273061405cfbda8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 16:33:17 -0400 Subject: [PATCH 04/12] Build Interleave sources through constructors and share the selector dispatch Primitive and decimal execution both repeated the dispatch over the two selector types. `gather` now owns it, and the slice kernel becomes `gather_rows`. `Source::constant` and `Source::decoded` set `row_mask` and the row count, so callers no longer write these fields by hand. The unit test now also covers the row check of a decoded source. Signed-off-by: Connor Tsui --- .../arrays/interleave/execute/decimal/mod.rs | 32 +---- .../arrays/interleave/execute/fixed_width.rs | 109 ++++++++++++++---- .../arrays/interleave/execute/primitive.rs | 33 +----- 3 files changed, 96 insertions(+), 78 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs index 3acd2fb4557..79d160bcc09 100644 --- a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs @@ -29,7 +29,6 @@ use crate::dtype::NativeDecimalType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::match_each_decimal_value_type; -use crate::match_each_unsigned_integer_ptype; use crate::require_child; pub(super) fn execute( @@ -58,23 +57,11 @@ fn gather_values( dtype: DecimalDType, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let allocator = ctx.allocator().clone(); - let values = (0..array.num_values()) + let sources = (0..array.num_values()) .map(|i| prepare_source::(array.value(i), dtype, ctx)) .collect::>>()?; - let branches = array.array_indices().as_::(); - let rows = array.row_indices().as_::(); - match_each_unsigned_integer_ptype!(branches.ptype(), |A| { - match_each_unsigned_integer_ptype!(rows.ptype(), |R| { - gather( - &values, - branches.as_slice::(), - rows.as_slice::(), - &allocator, - ) - }) - }) + gather(array, &sources, ctx.allocator()) } /// Normalizes one decimal source while retaining a single payload for a constant. @@ -83,7 +70,6 @@ fn prepare_source( dtype: DecimalDType, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let len = value.len(); if let Some(constant) = value.as_opt::() { let payload = match constant.scalar().as_decimal().decimal_value() { Some(value) => value.cast::().ok_or_else(|| { @@ -92,11 +78,7 @@ fn prepare_source( None => T::default(), }; - return Ok(Source { - data: Buffer::full_in(payload, 1, ctx.allocator().clone()), - len, - row_mask: 0, - }); + return Ok(Source::constant(payload, value.len(), ctx.allocator())); } let decimal = value.as_::(); @@ -106,18 +88,14 @@ fn prepare_source( // A source can still store its values at a width other than the canonical width of its // dtype (#10701). The copy widens such values, or checks their precision and narrows them // with zeros in the null slots. - let mask = decimal.validity()?.execute_mask(len, ctx)?; + let mask = decimal.validity()?.execute_mask(value.len(), ctx)?; let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? }); Buffer::::from_byte_buffer(copy.unwrap_host()) }; - Ok(Source { - data, - len, - row_mask: usize::MAX, - }) + Ok(Source::decoded(data)) } #[cfg(test)] diff --git a/vortex-array/src/arrays/interleave/execute/fixed_width.rs b/vortex-array/src/arrays/interleave/execute/fixed_width.rs index f1c154625ce..dfd59fd87c7 100644 --- a/vortex-array/src/arrays/interleave/execute/fixed_width.rs +++ b/vortex-array/src/arrays/interleave/execute/fixed_width.rs @@ -1,10 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Shared gather kernel for fixed-width [`Interleave`](super::super::Interleave) values. +//! Shared gather kernel for fixed-width [`Interleave`] values. //! -//! Each source holds a decoded buffer or one constant payload. Selector bounds are checked against -//! the logical source length before the kernel maps a row to its physical payload. +//! Primitive and decimal execution turn each value into a [`Source`] and call [`gather`]. They +//! differ only in how they read a constant payload or a decoded buffer, and in the array that they +//! build from the gathered payloads. use num_traits::AsPrimitive; use vortex_buffer::Buffer; @@ -15,18 +16,77 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -/// Physical storage for one fixed-width source array. +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use crate::array::Array; +use crate::arrays::Primitive; +use crate::match_each_unsigned_integer_ptype; + +/// Physical storage for one fixed-width value of an [`Interleave`] array. +/// +/// A constant keeps one payload instead of one payload for each row. `row_mask` maps each row to +/// its payload, so one gather loop reads constants and decoded buffers without a branch. pub(super) struct Source { - /// One constant payload or a decoded payload for each logical row. - pub(super) data: Buffer, - /// The logical row count, which can exceed `data.len()` for a constant. - pub(super) len: usize, - /// Maps rows to `data`: zero for constants and `usize::MAX` for decoded buffers. - pub(super) row_mask: usize, + /// One payload for each row, or a single payload for every row of a constant. + data: Buffer, + /// The number of logical rows. A constant can have more rows than `data` has payloads. + len: usize, + /// `usize::MAX` for a decoded buffer and zero for a constant, so that `row & row_mask` is the + /// position of the payload of `row` in `data`. + row_mask: usize, +} + +impl Source { + /// Returns a source of `len` rows that all hold `payload`. + pub(super) fn constant(payload: T, len: usize, allocator: &BufferAllocatorRef) -> Self { + Self { + data: Buffer::full_in(payload, 1, allocator.clone()), + len, + row_mask: 0, + } + } + + /// Returns a source with one row for each payload in `data`. + pub(super) fn decoded(data: Buffer) -> Self { + Self { + len: data.len(), + data, + row_mask: usize::MAX, + } + } } -pub(super) fn gather( - values: &[Source], +/// Gathers the payload of each output row of `array` from `sources`. +/// +/// `sources` holds one [`Source`] for each value of `array`, in order. Both selectors of `array` +/// **must** be [`Primitive`] arrays, which the kernels ensure with `require_child!`. +/// +/// # Errors +/// +/// Returns an error if the selectors have different lengths, or if a selector names a value or a +/// row that does not exist. +pub(super) fn gather( + array: &Array, + sources: &[Source], + allocator: &BufferAllocatorRef, +) -> VortexResult> { + let branches = array.array_indices().as_::(); + let rows = array.row_indices().as_::(); + + match_each_unsigned_integer_ptype!(branches.ptype(), |A| { + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather_rows( + sources, + branches.as_slice::(), + rows.as_slice::(), + allocator, + ) + }) + }) +} + +fn gather_rows( + sources: &[Source], branches: &[A], rows: &[R], allocator: &BufferAllocatorRef, @@ -50,7 +110,7 @@ where let mut output = BufferMut::with_capacity_in(len, allocator.clone()); let slots = &mut output.spare_capacity_mut()[..len]; for ((slot, branch), row) in slots.iter_mut().zip(branches).zip(rows) { - let Some(source) = values.get((*branch).as_()) else { + let Some(source) = sources.get((*branch).as_()) else { vortex_bail!("interleave array index out of bounds"); }; let row = (*row).as_(); @@ -66,29 +126,30 @@ where #[cfg(test)] mod tests { use rstest::rstest; - use vortex_buffer::Buffer; + use vortex_buffer::buffer; use vortex_error::VortexResult; use super::Source; - use super::gather; + use super::gather_rows; use crate::memory::BufferAllocatorRef; #[rstest] - #[case::array_index(1, 0, "array index out of bounds")] - #[case::row_index(0, 1, "row index out of bounds")] + #[case::array_index(2, 0, "array index out of bounds")] + #[case::constant_row_index(0, 1, "row index out of bounds")] + #[case::decoded_row_index(1, 1, "row index out of bounds")] fn rejects_out_of_bounds_selectors( #[case] branch: u8, #[case] row: u8, #[case] message: &str, ) -> VortexResult<()> { - let values = [Source { - data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()), - len: 1, - row_mask: 0, - }]; - let allocator = BufferAllocatorRef::statically_allocated(); - let error = gather(&values, &[branch], &[row], &allocator).unwrap_err(); + let sources = [ + Source::constant(1u32, 1, &allocator), + Source::decoded(buffer![2u32]), + ]; + + let error = gather_rows(&sources, &[branch], &[row], &allocator) + .expect_err("every case names a value or a row that does not exist"); assert!(error.to_string().contains(message), "{error}"); Ok(()) diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index 85f7d87ccd8..c758d49b47b 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -23,7 +23,6 @@ use crate::dtype::ToBytes; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::match_each_ptype_width; -use crate::match_each_unsigned_integer_ptype; use crate::require_child; pub(super) fn execute( @@ -58,10 +57,9 @@ fn gather_values( array: &Array, allocator: &BufferAllocatorRef, ) -> VortexResult> { - let values = (0..array.num_values()) + let sources = (0..array.num_values()) .map(|i| { let value = array.value(i); - let len = value.len(); if let Some(constant) = value.as_opt::() { // Validity carries nullness; a null constant's payload is never observed. let payload = constant @@ -74,33 +72,14 @@ fn gather_values( bytes }) .unwrap_or([0; W]); - Source { - data: Buffer::full_in(payload, 1, allocator.clone()), - len, - row_mask: 0, - } + Source::constant(payload, value.len(), allocator) } else { - Source { - data: Buffer::<[u8; W]>::from_byte_buffer( - value.as_::().buffer_handle().to_host_sync(), - ), - len, - row_mask: usize::MAX, - } + Source::decoded(Buffer::<[u8; W]>::from_byte_buffer( + value.as_::().buffer_handle().to_host_sync(), + )) } }) .collect::>(); - let branches = array.array_indices().as_::(); - let rows = array.row_indices().as_::(); - match_each_unsigned_integer_ptype!(branches.ptype(), |A| { - match_each_unsigned_integer_ptype!(rows.ptype(), |R| { - gather( - &values, - branches.as_slice::(), - rows.as_slice::(), - allocator, - ) - }) - }) + gather(array, &sources, allocator) } From c0eb004aae214aeb1373f906c8ef4e969c1efd2c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 16:35:59 -0400 Subject: [PATCH 05/12] Flatten the decimal Interleave module and keep only its own tests The decimal tests repeated cases that other tests already cover. The scale and nullability dimensions do not reach decimal code, the selector bounds checks belong to the shared `gather`, and `copy_to_canonical_width` has tests for each width. The remaining tests cover each canonical width, constants stored at other widths, sources stored at other widths, empty input and a value outside the precision. The lazy-children test in `interleave/mod.rs` now covers the decimal kernel too. With fewer tests, the kernel is a flat `decimal.rs` like `bool.rs` and `primitive.rs`. Signed-off-by: Connor Tsui --- .../src/arrays/interleave/execute/decimal.rs | 297 ++++++++++++++++++ .../arrays/interleave/execute/decimal/mod.rs | 102 ------ .../interleave/execute/decimal/tests.rs | 293 ----------------- vortex-array/src/arrays/interleave/mod.rs | 43 ++- 4 files changed, 329 insertions(+), 406 deletions(-) create mode 100644 vortex-array/src/arrays/interleave/execute/decimal.rs delete mode 100644 vortex-array/src/arrays/interleave/execute/decimal/mod.rs delete mode 100644 vortex-array/src/arrays/interleave/execute/decimal/tests.rs diff --git a/vortex-array/src/arrays/interleave/execute/decimal.rs b/vortex-array/src/arrays/interleave/execute/decimal.rs new file mode 100644 index 00000000000..793b17e3cb4 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/decimal.rs @@ -0,0 +1,297 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution for decimal [`Interleave`] values. +//! +//! Sources can use different stored widths for the same logical dtype. Decimal execution +//! normalizes them to the logical canonical width before gathering rows. + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use super::super::InterleaveSlots; +use super::fixed_width::Source; +use super::fixed_width::gather; +use crate::AnyColumnar; +use crate::ArrayRef; +use crate::array::Array; +use crate::arrays::Constant; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::Primitive; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::decimal::copy_to_canonical_width; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; +use crate::match_each_decimal_value_type; +use crate::require_child; + +pub(super) fn execute( + mut array: Array, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let num_values = array.num_values(); + array = + require_child!(array, array.array_indices(), InterleaveSlots::ARRAY_INDICES => Primitive); + array = require_child!(array, array.row_indices(), InterleaveSlots::ROW_INDICES => Primitive); + for i in 0..num_values { + array = require_child!(array, array.value(i), InterleaveSlots::VALUES_OFFSET + i => AnyColumnar); + } + + let validity = array.as_ref().validity()?; + let dtype = DecimalDType::try_from(array.dtype())?; + let output = match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + DecimalArray::new(gather_values::(&array, dtype, ctx)?, dtype, validity) + }); + + Ok(ExecutionResult::done(output)) +} + +fn gather_values( + array: &Array, + dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let sources = (0..array.num_values()) + .map(|i| prepare_source::(array.value(i), dtype, ctx)) + .collect::>>()?; + + gather(array, &sources, ctx.allocator()) +} + +/// Normalizes one decimal source while retaining a single payload for a constant. +fn prepare_source( + value: &ArrayRef, + dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if let Some(constant) = value.as_opt::() { + let payload = match constant.scalar().as_decimal().decimal_value() { + Some(value) => value.cast::().ok_or_else(|| { + vortex_err!("interleave decimal constant must fit {dtype}, got {value}") + })?, + None => T::default(), + }; + + return Ok(Source::constant(payload, value.len(), ctx.allocator())); + } + + let decimal = value.as_::(); + let data = if decimal.values_type() == T::DECIMAL_TYPE { + decimal.buffer::() + } else { + // A source can still store its values at a width other than the canonical width of its + // dtype (#10701). The copy widens such values, or checks their precision and narrows them + // with zeros in the null slots. + let mask = decimal.validity()?.execute_mask(value.len(), ctx)?; + let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { + copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? + }); + Buffer::::from_byte_buffer(copy.unwrap_host()) + }; + + Ok(Source::decoded(data)) +} + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_session::VortexSession; + + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::arrays::ConstantArray; + use crate::arrays::DecimalArray; + use crate::arrays::InterleaveArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::decimal::DecimalArrayExt; + use crate::assert_arrays_eq; + use crate::dtype::BigCast; + use crate::dtype::DType; + use crate::dtype::DecimalDType; + use crate::dtype::NativeDecimalType; + use crate::dtype::Nullability; + use crate::dtype::i256; + use crate::match_each_decimal_value_type; + use crate::scalar::DecimalValue; + use crate::scalar::Scalar; + use crate::validity::Validity; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + /// Interleaves `values` with one `(array_index, row_index)` pair for each output row, and + /// checks that the result has the canonical width of its dtype and equals `expected`. + #[track_caller] + fn assert_interleaved( + values: Vec, + indices: &[(u8, u8)], + expected: DecimalArray, + ) -> VortexResult<()> { + let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch)); + let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row)); + let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?; + let mut ctx = SESSION.create_execution_ctx(); + + let output = array.into_array().execute::(&mut ctx)?; + assert_eq!( + output.values_type(), + expected.decimal_dtype().canonical_values_type() + ); + assert_eq!(output.dtype(), expected.dtype()); + assert_arrays_eq!(output.into_array(), expected.into_array(), &mut ctx); + + Ok(()) + } + + /// Gathers rows out of order, with repeats and nulls, at each canonical width. The null slot of + /// the nullable source holds `T::MIN`, which can be outside the precision. + #[rstest] + #[case::i8(2)] + #[case::i16(4)] + #[case::i32(9)] + #[case::i64(18)] + #[case::i128(38)] + #[case::i256(76)] + fn test_decimal_interleave(#[case] precision: u8) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 2); + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + let min = T::MIN_BY_PRECISION[usize::from(precision)]; + let max = T::MAX_BY_PRECISION[usize::from(precision)]; + let zero = T::default(); + let nullable = DecimalArray::new( + Buffer::copy_from([min, T::MIN, max]), + dtype, + Validity::from_iter([true, false, true]), + ); + let non_nullable = DecimalArray::new( + Buffer::copy_from([max, min, zero]), + dtype, + Validity::NonNullable, + ); + let expected = DecimalArray::new( + Buffer::copy_from([zero, min, zero, max, max, min]), + dtype, + Validity::from_iter([true, true, false, true, true, true]), + ); + + assert_interleaved( + vec![nullable.into_array(), non_nullable.into_array()], + &[(1, 2), (0, 0), (0, 1), (1, 0), (0, 2), (0, 0)], + expected, + ) + }) + } + + /// A constant keeps one payload at the canonical width, also when its scalar stores the value + /// at another width. + #[rstest] + #[case::narrower_scalar(76, Some(DecimalValue::I8(-7)))] + #[case::wider_scalar(2, Some(DecimalValue::I256(i256::from_i128(-7))))] + #[case::null(9, None)] + fn test_decimal_interleave_constant( + #[case] precision: u8, + #[case] payload: Option, + ) -> VortexResult<()> { + let dtype = DecimalDType::new(precision, 2); + let (scalar, validity) = match payload { + Some(payload) => ( + Scalar::decimal(payload, dtype, Nullability::NonNullable), + Validity::NonNullable, + ), + None => ( + Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), + Validity::from_iter([false, true, false, false]), + ), + }; + let constant = ConstantArray::new(scalar, 3).into_array(); + match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { + let seven = ::from(7i8).vortex_expect("7 fits every decimal type"); + let minus_seven = + ::from(-7i8).vortex_expect("-7 fits every decimal type"); + let column = DecimalArray::new(buffer![seven], dtype, Validity::NonNullable); + let expected = DecimalArray::new( + buffer![minus_seven, seven, minus_seven, minus_seven], + dtype, + validity, + ); + + assert_interleaved( + vec![constant, column.into_array()], + &[(0, 2), (1, 0), (0, 1), (0, 2)], + expected, + ) + }) + } + + /// Sources that store their values at another width are copied to the canonical width first. + /// The `i8` source widens, and the `i256` source narrows with garbage in its null slots. + #[test] + fn test_decimal_interleave_mixed_stored_widths() -> VortexResult<()> { + let dtype = DecimalDType::new(9, 0); + let narrow = DecimalArray::new(buffer![1i8, -2], dtype, Validity::NonNullable); + let wide = DecimalArray::new( + buffer![i256::from_i128(3), i256::MIN, i256::MAX], + dtype, + Validity::from_iter([true, false, false]), + ); + let expected = DecimalArray::new( + buffer![3i32, -2, 0, 1, 0], + dtype, + Validity::from_iter([true, true, false, true, false]), + ); + + assert_interleaved( + vec![narrow.into_array(), wide.into_array()], + &[(1, 0), (0, 1), (1, 1), (0, 0), (1, 2)], + expected, + ) + } + + #[test] + fn test_decimal_interleave_empty() -> VortexResult<()> { + let dtype = DecimalDType::new(9, 0); + let empty = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable); + let expected = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable); + + assert_interleaved( + vec![empty.clone().into_array(), empty.into_array()], + &[], + expected, + ) + } + + #[test] + fn test_decimal_interleave_rejects_values_outside_precision() -> VortexResult<()> { + let source = DecimalArray::new( + buffer![100i16], + DecimalDType::new(2, 0), + Validity::NonNullable, + ) + .into_array(); + let array = InterleaveArray::try_new( + vec![source.clone(), source], + PrimitiveArray::from_iter([0u8]).into_array(), + PrimitiveArray::from_iter([0u8]).into_array(), + )? + .into_array(); + + let error = array + .execute::(&mut SESSION.create_execution_ctx()) + .expect_err("the narrowing copy to `i8` rejects 100 at precision 2"); + assert!(error.to_string().contains("precision"), "{error}"); + + Ok(()) + } +} diff --git a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs b/vortex-array/src/arrays/interleave/execute/decimal/mod.rs deleted file mode 100644 index 79d160bcc09..00000000000 --- a/vortex-array/src/arrays/interleave/execute/decimal/mod.rs +++ /dev/null @@ -1,102 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Execution for decimal [`Interleave`] values. -//! -//! Sources can use different stored widths for the same logical dtype. Decimal execution -//! normalizes them to the logical canonical width before gathering rows. - -use vortex_buffer::Buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_err; - -use super::super::Interleave; -use super::super::InterleaveArrayExt; -use super::super::InterleaveSlots; -use super::fixed_width::Source; -use super::fixed_width::gather; -use crate::AnyColumnar; -use crate::ArrayRef; -use crate::array::Array; -use crate::arrays::Constant; -use crate::arrays::Decimal; -use crate::arrays::DecimalArray; -use crate::arrays::Primitive; -use crate::arrays::decimal::DecimalArrayExt; -use crate::arrays::decimal::copy_to_canonical_width; -use crate::dtype::DecimalDType; -use crate::dtype::NativeDecimalType; -use crate::executor::ExecutionCtx; -use crate::executor::ExecutionResult; -use crate::match_each_decimal_value_type; -use crate::require_child; - -pub(super) fn execute( - mut array: Array, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let num_values = array.num_values(); - array = - require_child!(array, array.array_indices(), InterleaveSlots::ARRAY_INDICES => Primitive); - array = require_child!(array, array.row_indices(), InterleaveSlots::ROW_INDICES => Primitive); - for i in 0..num_values { - array = require_child!(array, array.value(i), InterleaveSlots::VALUES_OFFSET + i => AnyColumnar); - } - - let validity = array.as_ref().validity()?; - let dtype = DecimalDType::try_from(array.dtype())?; - let output = match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { - DecimalArray::new(gather_values::(&array, dtype, ctx)?, dtype, validity) - }); - - Ok(ExecutionResult::done(output)) -} - -fn gather_values( - array: &Array, - dtype: DecimalDType, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let sources = (0..array.num_values()) - .map(|i| prepare_source::(array.value(i), dtype, ctx)) - .collect::>>()?; - - gather(array, &sources, ctx.allocator()) -} - -/// Normalizes one decimal source while retaining a single payload for a constant. -fn prepare_source( - value: &ArrayRef, - dtype: DecimalDType, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - if let Some(constant) = value.as_opt::() { - let payload = match constant.scalar().as_decimal().decimal_value() { - Some(value) => value.cast::().ok_or_else(|| { - vortex_err!("interleave decimal constant must fit {dtype}, got {value}") - })?, - None => T::default(), - }; - - return Ok(Source::constant(payload, value.len(), ctx.allocator())); - } - - let decimal = value.as_::(); - let data = if decimal.values_type() == T::DECIMAL_TYPE { - decimal.buffer::() - } else { - // A source can still store its values at a width other than the canonical width of its - // dtype (#10701). The copy widens such values, or checks their precision and narrows them - // with zeros in the null slots. - let mask = decimal.validity()?.execute_mask(value.len(), ctx)?; - let copy = match_each_decimal_value_type!(decimal.values_type(), |S| { - copy_to_canonical_width(decimal.buffer::().as_slice(), &mask, dtype)? - }); - Buffer::::from_byte_buffer(copy.unwrap_host()) - }; - - Ok(Source::decoded(data)) -} - -#[cfg(test)] -mod tests; diff --git a/vortex-array/src/arrays/interleave/execute/decimal/tests.rs b/vortex-array/src/arrays/interleave/execute/decimal/tests.rs deleted file mode 100644 index 1095e6966c9..00000000000 --- a/vortex-array/src/arrays/interleave/execute/decimal/tests.rs +++ /dev/null @@ -1,293 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::LazyLock; - -use rstest::rstest; -use vortex_buffer::Buffer; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_session::VortexSession; - -use crate::ArrayRef; -use crate::IntoArray; -use crate::VortexSessionExecute; -use crate::arrays::ConstantArray; -use crate::arrays::DecimalArray; -use crate::arrays::InterleaveArray; -use crate::arrays::PrimitiveArray; -use crate::arrays::SliceArray; -use crate::arrays::decimal::DecimalArrayExt; -use crate::assert_arrays_eq; -use crate::dtype::BigCast; -use crate::dtype::DType; -use crate::dtype::DecimalDType; -use crate::dtype::DecimalType; -use crate::dtype::NativeDecimalType; -use crate::dtype::Nullability; -use crate::dtype::i256; -use crate::match_each_decimal_value_type; -use crate::scalar::DecimalValue; -use crate::scalar::Scalar; -use crate::validity::Validity; - -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - -/// Executes the gather, checks its physical width and dtype, and compares it with a canonical oracle. -#[track_caller] -fn assert_interleaved( - values: Vec, - indices: &[(u8, u64)], - expected: DecimalArray, -) -> VortexResult<()> { - let branches = PrimitiveArray::from_iter(indices.iter().map(|&(branch, _)| branch)); - let rows = PrimitiveArray::from_iter(indices.iter().map(|&(_, row)| row)); - let array = InterleaveArray::try_new(values, branches.into_array(), rows.into_array())?; - let mut ctx = SESSION.create_execution_ctx(); - let output = array.into_array().execute::(&mut ctx)?; - assert_eq!( - output.values_type(), - expected.decimal_dtype().canonical_values_type() - ); - assert_eq!(output.dtype(), expected.dtype()); - assert_arrays_eq!(output.into_array(), expected.into_array(), &mut ctx); - - Ok(()) -} - -#[rstest] -#[case::i8(2)] -#[case::i16(4)] -#[case::i32(9)] -#[case::i64(18)] -#[case::i128(38)] -#[case::i256(76)] -fn test_decimal_interleave( - #[case] precision: u8, - #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, - #[values(-3, 0, 2)] scale: i8, -) -> VortexResult<()> { - let dtype = DecimalDType::new(precision, scale); - let (first_validity, output_validity) = if nullability.is_nullable() { - ( - Validity::from_iter([true, false, true]), - Validity::from_iter([true, true, false, true, true, true]), - ) - } else { - (Validity::NonNullable, Validity::NonNullable) - }; - match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { - let min = T::MIN_BY_PRECISION[usize::from(precision)]; - let max = T::MAX_BY_PRECISION[usize::from(precision)]; - let zero = T::default(); - let middle = if nullability.is_nullable() { - T::MIN - } else { - zero - }; - let first = DecimalArray::new(Buffer::copy_from([min, middle, max]), dtype, first_validity); - let second = DecimalArray::new( - Buffer::copy_from([max, min, zero]), - dtype, - Validity::NonNullable, - ); - let expected = DecimalArray::new( - Buffer::copy_from([zero, min, zero, max, max, min]), - dtype, - output_validity, - ); - assert_interleaved( - vec![first.into_array(), second.into_array()], - &[(1, 2), (0, 0), (0, 1), (1, 0), (0, 2), (0, 0)], - expected, - ) - }) -} - -#[rstest] -#[case::i8(2)] -#[case::i16(4)] -#[case::i32(9)] -#[case::i64(18)] -#[case::i128(38)] -#[case::i256(76)] -fn test_decimal_interleave_constants( - #[case] precision: u8, - #[values(false, true)] null_constant: bool, - #[values(DecimalType::I8, DecimalType::I256)] scalar_storage: DecimalType, -) -> VortexResult<()> { - let dtype = DecimalDType::new(precision, 2); - let scalar = if null_constant { - Scalar::null(DType::Decimal(dtype, Nullability::Nullable)) - } else { - let payload = match_each_decimal_value_type!(scalar_storage, |S| { - DecimalValue::from( - ::from(-7i8).vortex_expect("-7 fits every decimal type"), - ) - }); - Scalar::decimal(payload, dtype, Nullability::NonNullable) - }; - let constant = ConstantArray::new(scalar, 3).into_array(); - match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { - let seven = ::from(7i8).vortex_expect("7 fits every decimal type"); - let minus_seven = ::from(-7i8).vortex_expect("-7 fits every decimal type"); - let column = DecimalArray::new(Buffer::copy_from([seven]), dtype, Validity::NonNullable) - .into_array(); - let expected = DecimalArray::new( - Buffer::copy_from([minus_seven, seven, minus_seven, minus_seven]), - dtype, - if null_constant { - Validity::from_iter([false, true, false, false]) - } else { - Validity::NonNullable - }, - ); - assert_interleaved( - vec![constant, column], - &[(0, 2), (1, 0), (0, 1), (0, 2)], - expected, - ) - }) -} - -#[rstest] -#[case::i8(2)] -#[case::i16(4)] -#[case::i32(9)] -#[case::i64(18)] -#[case::i128(38)] -#[case::i256(76)] -fn test_decimal_interleave_mixed_stored_widths(#[case] precision: u8) -> VortexResult<()> { - let dtype = DecimalDType::new(precision, 0); - let first = - DecimalArray::new(Buffer::copy_from([1i8, -2]), dtype, Validity::NonNullable).into_array(); - let second = DecimalArray::new( - Buffer::copy_from([i256::from_i128(3), i256::MIN, i256::MAX]), - dtype, - Validity::from_iter([true, false, false]), - ) - .into_array(); - match_each_decimal_value_type!(dtype.canonical_values_type(), |T| { - let values = [3i8, -2, 0, 1, 0].map(|v| { - ::from(v).vortex_expect("small test values fit every decimal type") - }); - let expected = DecimalArray::new( - Buffer::copy_from(values), - dtype, - Validity::from_iter([true, true, false, true, false]), - ); - assert_interleaved( - vec![first, second], - &[(1, 0), (0, 1), (1, 1), (0, 0), (1, 2)], - expected, - ) - }) -} - -#[rstest] -#[case::empty(&[])] -#[case::all_null(&[(0, 0), (1, 1), (0, 1)])] -fn test_decimal_interleave_without_valid_values(#[case] indices: &[(u8, u64)]) -> VortexResult<()> { - let dtype = DecimalDType::new(9, -2); - let first = DecimalArray::new( - Buffer::copy_from([i256::MIN, i256::MAX]), - dtype, - Validity::AllInvalid, - ) - .into_array(); - let second = ConstantArray::new( - Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), - 2, - ) - .into_array(); - let expected = DecimalArray::new( - Buffer::::zeroed(indices.len()), - dtype, - Validity::AllInvalid, - ); - assert_interleaved(vec![first, second], indices, expected) -} - -#[test] -fn test_decimal_interleave_empty_sources() -> VortexResult<()> { - let dtype = DecimalDType::new(9, 0); - let empty = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable).into_array(); - let expected = DecimalArray::new(Buffer::::empty(), dtype, Validity::NonNullable); - assert_interleaved(vec![empty.clone(), empty], &[], expected) -} - -#[test] -fn test_decimal_interleave_executes_lazy_children() -> VortexResult<()> { - let dtype = DecimalDType::new(2, 0); - let values = - DecimalArray::new(Buffer::copy_from([1i8, -2]), dtype, Validity::NonNullable).into_array(); - let lazy = SliceArray::try_new(values.clone(), 0..2)?.into_array(); - let branches = SliceArray::try_new(PrimitiveArray::from_iter([1u8, 0, 1]).into_array(), 0..3)? - .into_array(); - let rows = SliceArray::try_new(PrimitiveArray::from_iter([1u16, 0, 0]).into_array(), 0..3)? - .into_array(); - let interleaved = InterleaveArray::try_new(vec![values, lazy], branches, rows)?.into_array(); - let expected = DecimalArray::new( - Buffer::copy_from([-2i8, 1, 1]), - dtype, - Validity::NonNullable, - ) - .into_array(); - let mut ctx = SESSION.create_execution_ctx(); - assert_arrays_eq!(interleaved, expected, &mut ctx); - - Ok(()) -} - -#[rstest] -#[case::array_index(2, 0, "array index out of bounds")] -#[case::row_index(0, 2, "row index out of bounds")] -#[case::constant_row_index(1, 2, "row index out of bounds")] -#[case::large_array_index(u64::MAX, 0, "array index out of bounds")] -fn test_decimal_interleave_rejects_out_of_bounds_selectors( - #[case] branch: u64, - #[case] row: u8, - #[case] message: &str, - #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, -) -> VortexResult<()> { - let dtype = DecimalDType::new(2, 0); - let values = DecimalArray::new( - Buffer::copy_from([1i8, -2]), - dtype, - Validity::from(nullability), - ) - .into_array(); - let constant = - ConstantArray::new(Scalar::decimal(1i8.into(), dtype, nullability), 2).into_array(); - let array = InterleaveArray::try_new( - vec![values, constant], - PrimitiveArray::from_iter([branch]).into_array(), - PrimitiveArray::from_iter([row]).into_array(), - )? - .into_array(); - let error = array - .execute::(&mut SESSION.create_execution_ctx()) - .unwrap_err(); - assert!(error.to_string().contains(message), "{error}"); - - Ok(()) -} - -#[test] -fn test_decimal_interleave_rejects_wide_values_outside_precision() -> VortexResult<()> { - let dtype = DecimalDType::new(2, 0); - let source = - DecimalArray::new(Buffer::copy_from([100i16]), dtype, Validity::NonNullable).into_array(); - let array = InterleaveArray::try_new( - vec![source.clone(), source], - PrimitiveArray::from_iter([0u8]).into_array(), - PrimitiveArray::from_iter([0u8]).into_array(), - )? - .into_array(); - let error = array - .execute::(&mut SESSION.create_execution_ctx()) - .unwrap_err(); - assert!(error.to_string().contains("precision"), "{error}"); - - Ok(()) -} diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index 28c40c8d804..b9c5c7374ee 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -322,10 +322,12 @@ mod tests { use crate::VortexSessionExecute; use crate::arrays::Bool; use crate::arrays::BoolArray; + use crate::arrays::DecimalArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::SliceArray; use crate::assert_arrays_eq; + use crate::dtype::DecimalDType; use crate::dtype::PType; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -337,16 +339,25 @@ mod tests { ThirdValue, } + /// The value type, which picks the execution kernel. + #[derive(Clone, Copy)] + enum ValueKind { + Bool, + Primitive, + Decimal, + } + #[rstest] #[case::array_selector(LazyChild::ArraySelector)] #[case::row_selector(LazyChild::RowSelector)] #[case::third_value(LazyChild::ThirdValue)] fn executes_lazy_children( #[case] lazy_child: LazyChild, - #[values(false, true)] boolean: bool, + #[values(ValueKind::Bool, ValueKind::Primitive, ValueKind::Decimal)] value_kind: ValueKind, ) -> VortexResult<()> { - let (mut values, expected) = if boolean { - ( + let decimal_dtype = DecimalDType::new(9, 2); + let (mut values, expected) = match value_kind { + ValueKind::Bool => ( vec![ BoolArray::from_iter([true, false]).into_array(), BoolArray::from_iter([false, true]).into_array(), @@ -354,9 +365,8 @@ mod tests { ], BoolArray::from_iter([None, Some(true), Some(true), Some(false), Some(false)]) .into_array(), - ) - } else { - ( + ), + ValueKind::Primitive => ( vec![ PrimitiveArray::from_iter([10i32, 11]).into_array(), PrimitiveArray::from_iter([20i32, 21]).into_array(), @@ -364,7 +374,19 @@ mod tests { ], PrimitiveArray::from_option_iter([None, Some(10i32), Some(21), Some(30), Some(11)]) .into_array(), - ) + ), + ValueKind::Decimal => ( + vec![ + DecimalArray::from_iter([10i32, 11], decimal_dtype).into_array(), + DecimalArray::from_iter([20i32, 21], decimal_dtype).into_array(), + DecimalArray::from_option_iter([Some(30i32), None], decimal_dtype).into_array(), + ], + DecimalArray::from_option_iter( + [None, Some(10i32), Some(21), Some(30), Some(11)], + decimal_dtype, + ) + .into_array(), + ), }; let mut array_indices = PrimitiveArray::from_iter([2u8, 0, 1, 2, 0]).into_array(); let mut row_indices = PrimitiveArray::from_iter([1u32, 0, 1, 0, 1]).into_array(); @@ -384,10 +406,9 @@ mod tests { !matches!(lazy_child, LazyChild::RowSelector) ); for (i, value) in values.iter().enumerate() { - let matches = if boolean { - value.is::() - } else { - value.is::() + let matches = match value_kind { + ValueKind::Bool => value.is::(), + ValueKind::Primitive | ValueKind::Decimal => value.is::(), }; assert_eq!( matches, From 9beccd72d766f19926fcfa1dd28f9b79dd0f1ffc Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:11:12 -0400 Subject: [PATCH 06/12] Add canonical baselines for decimal Widen compute Signed-off-by: Connor Tsui --- vortex-array/Cargo.toml | 4 + vortex-array/benches/widen_decimal.rs | 261 ++++++++++++++++++++++++++ 2 files changed, 265 insertions(+) create mode 100644 vortex-array/benches/widen_decimal.rs diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 10733c48d8b..9432ddb964a 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -337,6 +337,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..225efb4cc22 --- /dev/null +++ b/vortex-array/benches/widen_decimal.rs @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compute on canonical decimals that can later use a narrower decimal representation. +//! +//! 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. + +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::DecimalArray; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +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 = CASES[0]; + case.values() + .array(case) + .execute::(&mut SESSION.create_execution_ctx()) + .vortex_expect("the benchmark input is a canonical decimal array"); + + 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, +} + +const CASES: &[Case] = &[ + Case { + precision: 18, + child_precision: 9, + }, + Case { + precision: 10, + child_precision: 5, + }, + Case { + precision: 38, + child_precision: 18, + }, +]; + +impl fmt::Display for Case { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "p{}_childp{}_canonical", + self.precision, self.child_precision, + ) + } +} + +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))); + + if self.precision <= 18 { + Values::I64(values.collect()) + } else { + Values::I128(values.map(i128::from).collect()) + } + } + + fn constant(self, value: i64) -> ArrayRef { + let scalar = Scalar::decimal(value.into(), self.decimal_dtype(), Nullability::NonNullable); + ConstantArray::new(scalar, LEN).into_array() + } + + 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 { + I64(Buffer), + I128(Buffer), +} + +impl Values { + fn array(&self, case: Case) -> ArrayRef { + match self { + Self::I64(values) => { + DecimalArray::new(values.clone(), case.decimal_dtype(), Validity::NonNullable) + .into_array() + } + Self::I128(values) => { + DecimalArray::new(values.clone(), case.decimal_dtype(), Validity::NonNullable) + .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), + ); +} From 23fa131c8c6326d4130ab6a03f8c740e1c93b843 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:15:56 -0400 Subject: [PATCH 07/12] Support decimal values in Widen arrays Signed-off-by: Connor Tsui --- vortex-array/benches/widen_decimal.rs | 114 ++++++++-- vortex-array/src/arrays/widen/array.rs | 44 ++-- vortex-array/src/arrays/widen/encode.rs | 88 ++++++++ vortex-array/src/arrays/widen/mod.rs | 12 +- .../src/arrays/widen/tests/decimal.rs | 205 ++++++++++++++++++ vortex-array/src/arrays/widen/tests/mod.rs | 12 + vortex-array/src/arrays/widen/vtable/mod.rs | 31 ++- 7 files changed, 464 insertions(+), 42 deletions(-) create mode 100644 vortex-array/src/arrays/widen/tests/decimal.rs diff --git a/vortex-array/benches/widen_decimal.rs b/vortex-array/benches/widen_decimal.rs index 225efb4cc22..c1e41ffb23a 100644 --- a/vortex-array/benches/widen_decimal.rs +++ b/vortex-array/benches/widen_decimal.rs @@ -1,12 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Compute on canonical decimals that can later use a narrower decimal representation. +//! 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. +//! 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; @@ -26,10 +27,15 @@ 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; @@ -50,11 +56,14 @@ fn main() { LazyLock::force(&SESSION); // Keep process-wide execution initialization outside the first measured operation. - let case = CASES[0]; + let case = Case { + input: Input::Widen, + ..CASES[0] + }; case.values() .array(case) .execute::(&mut SESSION.create_execution_ctx()) - .vortex_expect("the benchmark input is a canonical decimal array"); + .vortex_expect("the benchmark input executes to canonical decimals"); divan::main(); } @@ -73,29 +82,59 @@ const SCALE: i8 = 2; 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{}_canonical", - self.precision, self.child_precision, + "p{}_childp{}_{}", + self.precision, + self.child_precision, + match self.input { + Input::Canonical => "canonical", + Input::Widen => "widen", + }, ) } } @@ -119,10 +158,33 @@ impl Case { .into_iter() .chain((2..LEN).map(|_| rng.random_range(minimum..=maximum))); - if self.precision <= 18 { + 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"), } } @@ -149,19 +211,29 @@ impl Case { } enum Values { + I32(Buffer), I64(Buffer), I128(Buffer), } impl Values { - fn array(&self, case: Case) -> ArrayRef { + fn decimal_array(&self, dtype: DecimalDType) -> DecimalArray { match self { - Self::I64(values) => { - DecimalArray::new(values.clone(), case.decimal_dtype(), Validity::NonNullable) - .into_array() - } - Self::I128(values) => { - DecimalArray::new(values.clone(), case.decimal_dtype(), Validity::NonNullable) + 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() } } @@ -185,7 +257,7 @@ fn execute_canonical(array: VortexResult, ctx: &mut ExecutionCtx) -> C .vortex_expect("the operation executes to canonical for the benchmark inputs") } -#[divan::bench(args = CASES)] +#[divan::bench(args = cases())] fn compare_constant(bencher: Bencher, case: Case) { let values = case.values(); bench_case( @@ -195,7 +267,7 @@ fn compare_constant(bencher: Bencher, case: Case) { ); } -#[divan::bench(args = CASES)] +#[divan::bench(args = cases())] fn compare_constant_out_of_range(bencher: Bencher, case: Case) { let values = case.values(); bench_case( @@ -205,7 +277,7 @@ fn compare_constant_out_of_range(bencher: Bencher, case: Case) { ); } -#[divan::bench(args = CASES)] +#[divan::bench(args = cases())] fn sum(bencher: Bencher, case: Case) { let values = case.values(); bench_case( @@ -215,7 +287,7 @@ fn sum(bencher: Bencher, case: Case) { ); } -#[divan::bench(args = CASES)] +#[divan::bench(args = cases())] fn min_max(bencher: Bencher, case: Case) { let values = case.values(); bench_case( @@ -228,7 +300,7 @@ fn min_max(bencher: Bencher, case: Case) { ); } -#[divan::bench(args = CASES)] +#[divan::bench(args = cases())] fn filter(bencher: Bencher, case: Case) { let values = case.values(); let mut rng = StdRng::seed_from_u64(1); @@ -240,7 +312,7 @@ fn filter(bencher: Bencher, case: Case) { ); } -#[divan::bench(args = CASES)] +#[divan::bench(args = cases())] fn execute(bencher: Bencher, case: Case) { let values = case.values(); bench_case( @@ -250,7 +322,7 @@ fn execute(bencher: Bencher, case: Case) { ); } -#[divan::bench(args = CASES)] +#[divan::bench(args = cases())] fn cast_widen(bencher: Bencher, case: Case) { let values = case.values(); bench_case( diff --git a/vortex-array/src/arrays/widen/array.rs b/vortex-array/src/arrays/widen/array.rs index ed08a50dc47..260a39cf853 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,10 @@ 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. This trusts the child dtype and reads no values. When `values` is + /// itself a `Widen`, the new array stores its values instead, so that `Widen` arrays do not nest. 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 +53,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}" + ); } - _ => vortex_bail!("Widen requires integer dtypes, got {values} values for {logical}"), - }; + (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 or decimal dtypes, got {values} values for {logical}" + ), + } - 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..3c437abdf24 100644 --- a/vortex-array/src/arrays/widen/encode.rs +++ b/vortex-array/src/arrays/widen/encode.rs @@ -4,6 +4,7 @@ use std::ops::RangeInclusive; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use crate::ArrayRef; @@ -11,12 +12,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 +40,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 parent = array.decimal_dtype(); + let validity = array.validity()?; + let mask = validity.execute_mask(array.len(), ctx)?; + let range = match_each_decimal_value_type!(array.values_type(), |T| { + valid_range(array.buffer::().as_slice(), &mask).map(|(min, max)| { + ( + DecimalValue::from(min).as_i256(), + DecimalValue::from(max).as_i256(), + ) + }) + }); + let child_dtype = smallest_decimal_dtype(range, parent)?; + if child_dtype.canonical_values_type() == parent.canonical_values_type() { + return Ok(array.into_array()); + } + + let buffer = match_each_decimal_value_type!(array.values_type(), |T| { + copy_to_canonical_width(array.buffer::().as_slice(), &mask, child_dtype)? + }); + let dtype = array.dtype().clone(); + + // SAFETY: the range proves that every valid value fits the child precision. The copy + // preserves the scale and length, uses the child canonical width, and keeps its validity. + let values = unsafe { + DecimalArray::new_unchecked_handle( + buffer, + child_dtype.canonical_values_type(), + child_dtype, + validity, + ) + }; + + Ok(WidenArray::try_new(values.into_array(), dtype)?.into_array()) + } + /// Returns `array` cast to the type that [`Self::encode`] chooses for its values, without a /// `Widen` around them. /// @@ -107,3 +162,36 @@ 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 fn smallest_decimal_dtype( + range: Option<(i256, i256)>, + parent: DecimalDType, +) -> VortexResult { + let minimum_precision = parent.scale().max(1) as u8; + let Some((min, max)) = range else { + return Ok(DecimalDType::new(minimum_precision, parent.scale())); + }; + 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..b873655c7ab 100644 --- a/vortex-array/src/arrays/widen/mod.rs +++ b/vortex-array/src/arrays/widen/mod.rs @@ -1,7 +1,7 @@ // 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 @@ -9,9 +9,14 @@ //! values, but not `u8` values. [`WidenArray::encode`] picks the narrowest such type for a //! [`PrimitiveArray`]. //! +//! Decimal children keep the logical scale and nullability, with a strictly narrower canonical +//! width. [`WidenArray::encode_decimal`] selects the smallest precision that holds the valid values +//! and the scale. Decimal execution copies the child values into the logical canonical width, +//! without casting, because a widening decimal cast can itself return a `Widen`. +//! //! 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. @@ -31,6 +36,7 @@ pub use vtable::WidenArray; pub(crate) mod compute; mod encode; +pub use encode::smallest_decimal_dtype; mod vtable; use vortex_session::VortexSession; 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..658ce2384bc --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal.rs @@ -0,0 +1,205 @@ +// 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 super::stored_decimal_dtype; +use crate::ArrayParts; +use crate::ArrayRef; +use crate::Canonical; +use crate::EmptyArrayData; +use crate::IntoArray; +use crate::VortexSessionExecute; +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::WidenSlots; +use crate::arrays::widen::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::i32_to_i64(dtype(9, 2, NonNullable), dtype(18, 2, NonNullable), true)] +#[case::i128_to_i256(dtype(38, 0, Nullable), dtype(76, 0, Nullable), true)] +#[case::same_i64(dtype(10, 2, NonNullable), dtype(18, 2, NonNullable), false)] +#[case::same_i256(dtype(75, 0, NonNullable), dtype(76, 0, NonNullable), false)] +#[case::scale(dtype(9, 1, NonNullable), dtype(18, 2, NonNullable), false)] +#[case::nullability(dtype(9, 2, Nullable), dtype(18, 2, NonNullable), false)] +#[case::higher_precision(dtype(18, 2, NonNullable), dtype(9, 2, NonNullable), false)] +#[case::integer_child(PType::I32.into(), dtype(18, 2, NonNullable), false)] +#[case::integer_parent(dtype(9, 2, NonNullable), PType::I64.into(), false)] +fn test_try_new_decimal_dtypes( + #[case] child: DType, + #[case] parent: DType, + #[case] accepted: bool, +) -> VortexResult<()> { + let values = Canonical::empty(&child).into_array(); + let parts = ArrayParts::new( + Widen, + parent.clone(), + 0, + EmptyArrayData, + WidenSlots { + values: values.clone(), + } + .into_slots(), + ); + let constructed = WidenArray::try_new(values, parent.clone()); + let validated = WidenArray::try_from_parts(parts); + assert_eq!(constructed.is_ok(), accepted); + assert_eq!(validated.is_ok(), accepted); + if !accepted { + let error = constructed + .expect_err("the dtype pair is invalid") + .to_string(); + assert!(error.contains(&child.to_string()), "{error}"); + assert!(error.contains(&parent.to_string()), "{error}"); + } + + Ok(()) +} + +#[test] +fn test_decimal_widen_flattens() -> VortexResult<()> { + let values = DecimalArray::from_iter([1i16, 99], DecimalDType::new(4, 0)).into_array(); + let inner = WidenArray::try_new(values.clone(), dtype(9, 0, NonNullable))?.into_array(); + let outer = WidenArray::try_new(inner, dtype(18, 0, NonNullable))?; + assert!(ArrayRef::ptr_eq(outer.values(), &values)); + assert!(!outer.values().is::()); + + Ok(()) +} + +#[rstest] +#[case::canonical(DecimalArray::from_iter([-999i32, 0, 999], DecimalDType::new(9, 2)), [-999, 0, 999])] +#[case::nullable(DecimalArray::from_option_iter([Some(-999i32), None, Some(999)], DecimalDType::new(9, 2)), [-999, 0, 999])] +#[case::narrow_stored(DecimalArray::from_iter([-99i8, 0, 99], DecimalDType::new(9, 2)), [-99, 0, 99])] +#[case::parent_width_stored(DecimalArray::from_iter([-999i64, 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.clone().into_array(), parent.clone())?.into_array(); + let executed = widened.clone().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.clone(), &mut ctx); + assert_arrays_eq!(widened, reference, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::scale_raises_precision(vec![Some(1)], 7, 7, true)] +#[case::negative_boundary(vec![Some(-128), Some(99)], 0, 3, true)] +#[case::negative_scale(vec![Some(-99), Some(99)], -3, 2, true)] +#[case::all_null(vec![None, None], 0, 1, true)] +#[case::empty(vec![], 0, 1, true)] +#[case::no_narrower_width(vec![Some(1_000_000_000)], 2, 10, false)] +fn test_encode_decimal_best_precision( + #[case] values: Vec>, + #[case] scale: i8, + #[case] child_precision: u8, + #[case] is_widen: bool, +) -> 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)?; + assert_eq!(encoded.is::(), is_widen); + if let Some(widen) = encoded.as_opt::() { + assert_eq!(stored_decimal_dtype(&encoded).precision(), child_precision); + let child = widen.values().as_::(); + 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(()) +} + +#[rstest] +#[case::empty_positive_scale(None, 18, 12, Some(12))] +#[case::empty_negative_scale(None, 18, -4, Some(1))] +#[case::negative_boundary(Some((-128, 0)), 18, 0, Some(3))] +#[case::decimal_boundary(Some((-999, 999)), 18, 0, Some(3))] +#[case::outside_parent(Some((-1_000, 1)), 3, 0, None)] +#[case::reversed(Some((2, 1)), 18, 0, None)] +fn test_smallest_decimal_dtype( + #[case] range: Option<(i64, i64)>, + #[case] parent_precision: u8, + #[case] scale: i8, + #[case] precision: Option, +) -> VortexResult<()> { + let range = range.map(|(min, max)| (i256::from_i128(min.into()), i256::from_i128(max.into()))); + let result = smallest_decimal_dtype(range, DecimalDType::new(parent_precision, scale)); + match precision { + Some(precision) => assert_eq!(result?, DecimalDType::new(precision, scale)), + None => assert!(result.is_err()), + } + + Ok(()) +} + +#[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/mod.rs b/vortex-array/src/arrays/widen/tests/mod.rs index 6a959e199d3..587cccf7333 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -8,6 +8,7 @@ mod aggregates; mod compute; +mod decimal; mod encoding; mod numeric; @@ -27,6 +28,7 @@ use crate::arrays::Widen; use crate::arrays::WidenArray; use crate::arrays::widen::WidenArraySlotsExt; use crate::dtype::DType; +use crate::dtype::DecimalDType; use crate::dtype::PType; use crate::validity::Validity; @@ -48,6 +50,16 @@ fn stored_ptype(array: &ArrayRef) -> PType { } } +/// Returns the decimal dtype of the values child, or of `array` when it is not a `Widen`. +fn stored_decimal_dtype(array: &ArrayRef) -> DecimalDType { + let dtype = match array.as_opt::() { + Some(widen) => widen.values().dtype().as_decimal_opt().copied(), + None => array.dtype().as_decimal_opt().copied(), + }; + + dtype.vortex_expect("the fixture has a decimal dtype") +} + /// Returns `array` executed to its canonical form. fn canonical(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { Ok(array.clone().execute::(ctx)?.into_array()) diff --git a/vortex-array/src/arrays/widen/vtable/mod.rs b/vortex-array/src/arrays/widen/vtable/mod.rs index 3af2178b523..37fdb4b47d2 100644 --- a/vortex-array/src/arrays/widen/vtable/mod.rs +++ b/vortex-array/src/arrays/widen/vtable/mod.rs @@ -24,6 +24,9 @@ use crate::array::ArrayView; use crate::array::VTable; use crate::array::ValidityVTableFromChild; use crate::array::with_empty_buffers; +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 +34,14 @@ use crate::arrays::widen::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::match_each_decimal_value_type; 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 { @@ -112,7 +116,30 @@ impl VTable for Widen { WidenSlots::NAMES[idx].to_string() } - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { + fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + if let DType::Decimal(decimal_dtype, _) = array.dtype() { + let values = array.values().clone().execute::(ctx)?; + let validity = values.validity()?; + let mask = validity.execute_mask(values.len(), ctx)?; + let buffer = match_each_decimal_value_type!(values.values_type(), |T| { + copy_to_canonical_width(values.buffer::().as_slice(), &mask, *decimal_dtype)? + }); + + // SAFETY: the copy has the parent canonical width and the child length. Widen + // construction guarantees the same scale and a parent precision that holds the child, + // and the unchanged child validity has the same length. + let widened = unsafe { + DecimalArray::new_unchecked_handle( + buffer, + decimal_dtype.canonical_values_type(), + *decimal_dtype, + validity, + ) + }; + + return Ok(ExecutionResult::done(widened.into_array())); + } + // 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)) From 9d55543c3e2fab757f6f9573178839fcbc544dd5 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:19:37 -0400 Subject: [PATCH 08/12] Push decimal comparisons and fitting fills through Widen Signed-off-by: Connor Tsui --- .../src/arrays/widen/compute/between.rs | 6 +- .../src/arrays/widen/compute/compare.rs | 31 +- .../src/arrays/widen/compute/fill_null.rs | 18 ++ vortex-array/src/arrays/widen/compute/mod.rs | 28 +- .../src/arrays/widen/tests/decimal_compare.rs | 302 ++++++++++++++++++ vortex-array/src/arrays/widen/tests/mod.rs | 1 + 6 files changed, 375 insertions(+), 11 deletions(-) create mode 100644 vortex-array/src/arrays/widen/tests/decimal_compare.rs diff --git a/vortex-array/src/arrays/widen/compute/between.rs b/vortex-array/src/arrays/widen/compute/between.rs index 9f6068651ce..33d73ce5f9a 100644 --- a/vortex-array/src/arrays/widen/compute/between.rs +++ b/vortex-array/src/arrays/widen/compute/between.rs @@ -9,7 +9,7 @@ use crate::arrays::Constant; use crate::arrays::Widen; use crate::arrays::widen::WidenArraySlotsExt; use crate::arrays::widen::compute::constant_fits_values; -use crate::arrays::widen::compute::constant_in_values_ptype; +use crate::arrays::widen::compute::constant_in_values_dtype; use crate::builtins::ArrayBuiltins; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::BetweenReduce; @@ -28,8 +28,8 @@ impl BetweenReduce for Widen { && constant_fits_values(lower, values.dtype()) == Some(true) && constant_fits_values(upper, values.dtype()) == Some(true) { - let lower = constant_in_values_ptype(lower, values, array.len())?; - let upper = constant_in_values_ptype(upper, values, array.len())?; + let lower = constant_in_values_dtype(lower, values, array.len())?; + let upper = constant_in_values_dtype(upper, values, array.len())?; return values .clone() .between(lower, upper, options.clone()) diff --git a/vortex-array/src/arrays/widen/compute/compare.rs b/vortex-array/src/arrays/widen/compute/compare.rs index 9f3b5c4985d..3ad8967b6d3 100644 --- a/vortex-array/src/arrays/widen/compute/compare.rs +++ b/vortex-array/src/arrays/widen/compute/compare.rs @@ -15,9 +15,11 @@ use crate::arrays::Widen; use crate::arrays::widen::WidenArraySlotsExt; use crate::arrays::widen::compute::cast_keeping_nullability; use crate::arrays::widen::compute::constant_fits_values; -use crate::arrays::widen::compute::constant_in_values_ptype; +use crate::arrays::widen::compute::constant_in_values_dtype; use crate::arrays::widen::compute::ptype_holding_both_values; use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; use crate::scalar::Scalar; use crate::scalar_fn::fns::binary::CompareKernel; use crate::scalar_fn::fns::operators::CompareOperator; @@ -50,7 +52,7 @@ impl CompareKernel for Widen { return compare_out_of_range(lhs, constant, operator, ctx).map(Some); } - let constant = constant_in_values_ptype(constant, values, lhs.len())?; + let constant = constant_in_values_dtype(constant, values, lhs.len())?; values.binary(constant, Operator::from(operator)).map(Some) } @@ -62,6 +64,31 @@ fn compare_narrow_values( rhs: ArrayView<'_, Widen>, operator: CompareOperator, ) -> VortexResult> { + if let (DType::Decimal(lhs_dtype, _), DType::Decimal(rhs_dtype, _)) = + (lhs.values().dtype(), rhs.values().dtype()) + { + if lhs_dtype.scale() != rhs_dtype.scale() { + return Ok(None); + } + + let decimal_dtype = DecimalDType::new( + lhs_dtype.precision().max(rhs_dtype.precision()), + lhs_dtype.scale(), + ); + let lhs_values = lhs.values().cast(DType::Decimal( + decimal_dtype, + lhs.values().dtype().nullability(), + ))?; + let rhs_values = rhs.values().cast(DType::Decimal( + decimal_dtype, + rhs.values().dtype().nullability(), + ))?; + + return lhs_values + .binary(rhs_values, Operator::from(operator)) + .map(Some); + } + let Some(ptype) = ptype_holding_both_values(lhs, rhs) else { return Ok(None); }; diff --git a/vortex-array/src/arrays/widen/compute/fill_null.rs b/vortex-array/src/arrays/widen/compute/fill_null.rs index e869b4ebc40..644e3ec34c9 100644 --- a/vortex-array/src/arrays/widen/compute/fill_null.rs +++ b/vortex-array/src/arrays/widen/compute/fill_null.rs @@ -8,6 +8,7 @@ use crate::array::ArrayView; use crate::arrays::Widen; use crate::arrays::widen::WidenArraySlotsExt; use crate::arrays::widen::compute::cast_keeping_nullability; +use crate::arrays::widen::compute::constant_fits_values; use crate::arrays::widen::compute::ptype_holding_values_and_scalar; use crate::arrays::widen::compute::wrap_values; use crate::builtins::ArrayBuiltins; @@ -20,6 +21,23 @@ impl FillNullReduce for Widen { array: ArrayView<'_, Self>, fill_value: &Scalar, ) -> VortexResult> { + if matches!(array.values().dtype(), DType::Decimal(..)) { + if !fill_value.dtype().eq_ignore_nullability(array.dtype()) + || constant_fits_values(fill_value, array.values().dtype()) != Some(true) + { + return Ok(None); + } + + let fill_dtype = array + .values() + .dtype() + .with_nullability(fill_value.dtype().nullability()); + let filled = array.values().fill_null(fill_value.cast(&fill_dtype)?)?; + let dtype = array.dtype().with_nullability(filled.dtype().nullability()); + + return wrap_values(filled, dtype).map(Some); + } + let Some(ptype) = ptype_holding_values_and_scalar(array, fill_value)? else { return Ok(None); }; diff --git a/vortex-array/src/arrays/widen/compute/mod.rs b/vortex-array/src/arrays/widen/compute/mod.rs index e409beed88b..5d43ef4fcff 100644 --- a/vortex-array/src/arrays/widen/compute/mod.rs +++ b/vortex-array/src/arrays/widen/compute/mod.rs @@ -100,11 +100,27 @@ pub(super) fn cast_keeping_nullability(array: &ArrayRef, ptype: PType) -> Vortex array.cast(DType::Primitive(ptype, array.dtype().nullability())) } -/// Returns whether the range of the integer dtype `values` holds `constant`. +/// Returns whether the dtype range of `values` holds `constant`. /// -/// Returns `None` when `values` is not an integer dtype, when `constant` is null, and when -/// `constant` is not an integer with the signedness of `values`. +/// Integers must have the same signedness, and decimals must have the same scale. Returns `None` +/// for null constants, unsupported dtypes, or incompatible constants. pub(super) fn constant_fits_values(constant: &Scalar, values: &DType) -> Option { + if let DType::Decimal(decimal_dtype, _) = values { + let DType::Decimal(constant_dtype, _) = constant.dtype() else { + return None; + }; + if constant_dtype.scale() != decimal_dtype.scale() { + return None; + } + + return Some( + constant + .as_decimal() + .decimal_value()? + .fits_in_precision(*decimal_dtype), + ); + } + let DType::Primitive(ptype, _) = values else { return None; }; @@ -126,8 +142,8 @@ pub(super) fn constant_fits_values(constant: &Scalar, values: &DType) -> Option< /// Returns the wider of the value types of two [`Widen`] arrays, or `None` when a value type is /// not primitive. /// -/// The logical dtypes **must** have the same signedness, which a debug build asserts. The value -/// types then have that signedness too, so the wider one holds every value of the narrower one. +/// The logical dtypes **must** have the same signedness. The value types then have that signedness +/// too, so the wider one holds every value of the narrower one. pub(super) fn ptype_holding_both_values( lhs: ArrayView<'_, Widen>, rhs: ArrayView<'_, Widen>, @@ -147,7 +163,7 @@ pub(super) fn ptype_holding_both_values( /// Returns `constant`, which **must** fit the type of `values`, as a constant array of that type /// with the nullability of `constant`. -pub(super) fn constant_in_values_ptype( +pub(super) fn constant_in_values_dtype( constant: &Scalar, values: &ArrayRef, len: usize, diff --git a/vortex-array/src/arrays/widen/tests/decimal_compare.rs b/vortex-array/src/arrays/widen/tests/decimal_compare.rs new file mode 100644 index 00000000000..dc5d8b25979 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_compare.rs @@ -0,0 +1,302 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal comparison, bounds, and fill-null pushdown tests. +//! +//! Canonical decimal operations provide the reference results, including nullable operands. + +use rstest::rstest; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::arrays::widen::compute::constant_fits_values; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::i256; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::between::BetweenOptions; +use crate::scalar_fn::fns::between::StrictComparison; +use crate::scalar_fn::fns::binary::CompareKernel; +use crate::scalar_fn::fns::fill_null::FillNullReduce; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::fns::operators::Operator; + +/// Checks both operand orders against fully decoded decimal operands. +#[track_caller] +fn assert_compare_matches_canonical( + lhs: &ArrayRef, + rhs: &ArrayRef, + operator: Operator, + ctx: &mut ExecutionCtx, +) -> VortexResult<()> { + let canonical_lhs = lhs.clone().execute::(ctx)?.into_array(); + let canonical_rhs = rhs.clone().execute::(ctx)?.into_array(); + let expected = canonical_lhs.binary(canonical_rhs.clone(), operator)?; + assert_arrays_eq!(lhs.binary(rhs.clone(), operator)?, expected, ctx); + + let expected = canonical_rhs.binary(canonical_lhs, operator)?; + assert_arrays_eq!(rhs.binary(lhs.clone(), operator)?, expected, ctx); + + Ok(()) +} + +fn widen(values: impl IntoArray) -> VortexResult { + let values = values.into_array(); + let dtype = DType::Decimal(DecimalDType::new(18, 2), values.dtype().nullability()); + + WidenArray::try_new(values, dtype) +} + +fn decimal_scalar(value: Option, nullability: Nullability) -> Scalar { + let dtype = DecimalDType::new(18, 2); + match value { + Some(value) => Scalar::decimal(DecimalValue::I64(value), dtype, nullability), + None => Scalar::null(DType::Decimal(dtype, Nullability::Nullable)), + } +} + +#[rstest] +#[case::below_precision(Some(-100_000))] +#[case::minimum(Some(-99_999))] +#[case::negative_inside(Some(-32_768))] +#[case::zero(Some(0))] +#[case::positive_inside(Some(32_768))] +#[case::maximum(Some(99_999))] +#[case::above_precision(Some(100_000))] +#[case::null(None)] +fn test_compare_constant( + #[case] value: Option, + #[values(false, true)] nullable: bool, + #[values( + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::Lte, + Operator::Gt, + Operator::Gte + )] + operator: Operator, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(5, 2); + let values = if nullable { + DecimalArray::from_option_iter( + [Some(-99_999i32), None, Some(0), Some(32_768), Some(99_999)], + decimal_dtype, + ) + } else { + DecimalArray::from_iter([-99_999i32, -32_768, 0, 32_768, 99_999], decimal_dtype) + }; + let array = widen(values)?.into_array(); + let constant = + ConstantArray::new(decimal_scalar(value, Nullability::Nullable), array.len()).into_array(); + + assert_compare_matches_canonical(&array, &constant, operator, &mut ctx) +} + +#[rstest] +fn test_compare_two_widen( + #[values( + Operator::Eq, + Operator::NotEq, + Operator::Lt, + Operator::Lte, + Operator::Gt, + Operator::Gte + )] + operator: Operator, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = widen(DecimalArray::from_iter( + [-99_999i32, -1, 0, 99_999], + DecimalDType::new(5, 2), + ))? + .into_array(); + let rhs = widen(DecimalArray::from_option_iter( + [Some(-100_000i32), None, Some(0), Some(999_999_999)], + DecimalDType::new(9, 2), + ))? + .into_array(); + + assert_compare_matches_canonical(&lhs, &rhs, operator, &mut ctx) +} + +#[test] +fn test_compare_uses_child_precision_and_nullability() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = widen(DecimalArray::from_iter( + [-99_999i32, 0, 99_999], + DecimalDType::new(5, 2), + ))?; + let rhs = widen(DecimalArray::from_option_iter( + [Some(-100_000i32), None, Some(999_999_999)], + DecimalDType::new(9, 2), + ))? + .into_array(); + + let compared = Widen::compare(lhs.as_view(), &rhs, CompareOperator::Lt, &mut ctx)? + .vortex_expect("Widen compares decimal children at their common precision"); + let lhs_values = compared.nth_child(0).vortex_expect("compare has a lhs"); + let rhs_values = compared.nth_child(1).vortex_expect("compare has a rhs"); + assert_eq!( + lhs_values.dtype(), + &DType::Decimal(DecimalDType::new(9, 2), Nullability::NonNullable) + ); + assert_eq!( + rhs_values.dtype(), + &DType::Decimal(DecimalDType::new(9, 2), Nullability::Nullable) + ); + + Ok(()) +} + +#[test] +fn test_constant_fits_precision_and_scale() { + let dtype = DType::Decimal(DecimalDType::new(75, 2), Nullability::NonNullable); + let max = i256::MAX_BY_PRECISION[75]; + for (value, fits) in [ + (max, true), + (-max, true), + (max + i256::ONE, false), + (-max - i256::ONE, false), + ] { + let scalar = Scalar::decimal( + DecimalValue::I256(value), + DecimalDType::new(76, 2), + Nullability::NonNullable, + ); + assert_eq!(constant_fits_values(&scalar, &dtype), Some(fits)); + } + + let different_scale = Scalar::decimal( + DecimalValue::I8(1), + DecimalDType::new(18, 3), + Nullability::NonNullable, + ); + assert_eq!(constant_fits_values(&different_scale, &dtype), None); + assert_eq!( + constant_fits_values(&Scalar::null(dtype.as_nullable()), &dtype), + None + ); + assert_eq!(constant_fits_values(&Scalar::from(1i64), &dtype), None); +} + +#[rstest] +#[case::inside(Some(-99_999), Some(99_999))] +#[case::outside(Some(-100_000), Some(100_000))] +#[case::below(Some(-100_001), Some(-100_000))] +#[case::above(Some(100_000), Some(100_001))] +#[case::reversed(Some(99_999), Some(-99_999))] +#[case::null_lower(None, Some(99_999))] +#[case::null_upper(Some(-99_999), None)] +fn test_between( + #[case] lower: Option, + #[case] upper: Option, + #[values(StrictComparison::Strict, StrictComparison::NonStrict)] lower_strict: StrictComparison, + #[values(StrictComparison::Strict, StrictComparison::NonStrict)] upper_strict: StrictComparison, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = widen(DecimalArray::from_option_iter( + [Some(-99_999i32), None, Some(0), Some(99_999)], + DecimalDType::new(5, 2), + ))? + .into_array(); + let lower = + ConstantArray::new(decimal_scalar(lower, Nullability::Nullable), array.len()).into_array(); + let upper = + ConstantArray::new(decimal_scalar(upper, Nullability::Nullable), array.len()).into_array(); + let options = BetweenOptions { + lower_strict, + upper_strict, + }; + + let result = array + .clone() + .between(lower.clone(), upper.clone(), options.clone())?; + let expected = array + .execute::(&mut ctx)? + .into_array() + .between(lower, upper, options)?; + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::negative_fits(-99_999, true)] +#[case::positive_fits(99_999, true)] +#[case::negative_outside(-100_000, false)] +#[case::positive_outside(100_000, false)] +fn test_fill_null( + #[case] value: i64, + #[case] fits: bool, + #[values(Nullability::NonNullable, Nullability::Nullable)] nullability: Nullability, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = widen(DecimalArray::from_option_iter( + [Some(-99_999i32), None, Some(0), Some(99_999)], + DecimalDType::new(5, 2), + ))?; + let fill = decimal_scalar(Some(value), nullability); + + let reduced = Widen::fill_null(array.as_view(), &fill)?; + assert_eq!(reduced.is_some(), fits); + if let Some(reduced) = reduced { + let widen = reduced.as_::(); + let child = widen.values(); + assert_eq!( + child.dtype(), + &DType::Decimal(DecimalDType::new(5, 2), nullability) + ); + } + + let array = array.into_array(); + let result = array.fill_null(fill.clone())?; + let expected = array + .execute::(&mut ctx)? + .into_array() + .fill_null(fill)?; + assert_arrays_eq!(result, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_fill_null_rejects_null_scalar() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = widen(DecimalArray::from_option_iter( + [Some(1i32), None], + DecimalDType::new(5, 2), + ))?; + let fill = decimal_scalar(None, Nullability::Nullable); + + assert!(Widen::fill_null(array.as_view(), &fill)?.is_none()); + let error = array + .into_array() + .fill_null(fill) + .and_then(|result| result.execute::(&mut ctx)) + .expect_err("fill_null requires a non-null scalar"); + assert!( + error + .to_string() + .contains("fill_null requires a non-null fill value") + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/mod.rs b/vortex-array/src/arrays/widen/tests/mod.rs index 587cccf7333..ff457a5833c 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -9,6 +9,7 @@ mod aggregates; mod compute; mod decimal; +mod decimal_compare; mod encoding; mod numeric; From 781f85ca3dd4ea62bac072d634d2d1ead82edaad Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:22:27 -0400 Subject: [PATCH 09/12] Push ordered decimal aggregates through Widen Signed-off-by: Connor Tsui --- .../src/arrays/widen/compute/aggregate.rs | 30 ++- .../src/arrays/widen/tests/aggregates.rs | 29 +- .../arrays/widen/tests/decimal_aggregates.rs | 248 ++++++++++++++++++ vortex-array/src/arrays/widen/tests/mod.rs | 1 + 4 files changed, 292 insertions(+), 16 deletions(-) create mode 100644 vortex-array/src/arrays/widen/tests/decimal_aggregates.rs diff --git a/vortex-array/src/arrays/widen/compute/aggregate.rs b/vortex-array/src/arrays/widen/compute/aggregate.rs index aeca44dc042..e00838144f0 100644 --- a/vortex-array/src/arrays/widen/compute/aggregate.rs +++ b/vortex-array/src/arrays/widen/compute/aggregate.rs @@ -7,6 +7,11 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::fns::is_constant::IsConstant; +use crate::aggregate_fn::fns::is_sorted::IsSorted; +use crate::aggregate_fn::fns::max::Max; +use crate::aggregate_fn::fns::min::Min; +use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::arrays::Widen; @@ -19,7 +24,10 @@ use crate::scalar::Scalar; /// /// The values run the same accumulator code as the canonical array, so edge cases such as an empty /// batch or a single null match by construction, and the values keep their own cached statistics -/// and kernels. [`Sum`] is served only when [`is_exact_sum`] holds. +/// and kernels. Decimal widening preserves order, equality and validity, so extrema, constancy and +/// sortedness have the same partial after widening. Decimal sums must use the logical precision +/// to select their accumulator, so they fall back to execution. Integer [`Sum`] is served only +/// when [`is_exact_sum`] holds. #[derive(Debug)] pub(crate) struct WidenArrayAggregate; @@ -34,10 +42,19 @@ impl DynAggregateKernel for WidenArrayAggregate { return Ok(None); }; - // NB: Exactness rests on integer decoding. A decimal `Widen` (CD4 in #10733) needs its own - // argument, because `Sum` checks a decimal for overflow against the precision of its input. - if !batch.dtype().is_int() { - return Ok(None); + match batch.dtype() { + DType::Decimal(..) => { + if !(aggregate_fn.is::() + || aggregate_fn.is::() + || aggregate_fn.is::() + || aggregate_fn.is::() + || aggregate_fn.is::()) + { + return Ok(None); + } + } + dtype if dtype.is_int() => {} + _ => return Ok(None), } let values = widen.values(); @@ -49,8 +66,7 @@ impl DynAggregateKernel for WidenArrayAggregate { accumulator.accumulate(values, ctx)?; let partial_dtype = aggregate_fn.state_dtype(batch.dtype()).vortex_expect( - "an aggregate that accepts the values accepts the logical dtype, which is a wider \ - integer of the same signedness", + "the served aggregate accepts the logical dtype when it accepts the narrower values", ); Ok(Some(accumulator.flush()?.cast(&partial_dtype)?)) diff --git a/vortex-array/src/arrays/widen/tests/aggregates.rs b/vortex-array/src/arrays/widen/tests/aggregates.rs index b91640151b3..46e07b200a2 100644 --- a/vortex-array/src/arrays/widen/tests/aggregates.rs +++ b/vortex-array/src/arrays/widen/tests/aggregates.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fmt::Debug; use std::num::NonZeroUsize; use rstest::rstest; @@ -99,7 +100,7 @@ fn canonical_batch( /// Runs the kernel that the session registers for `aggregate` on `Widen` over `array`, and returns /// its partial, or `None` when no kernel is registered or the kernel declines. -fn kernel_partial( +pub(super) fn kernel_partial( aggregate: &AggregateFnRef, array: &ArrayRef, ctx: &mut ExecutionCtx, @@ -113,26 +114,28 @@ fn kernel_partial( } } -/// Accumulates `batches` of `dtype` with `aggregate` once as [`WidenArray::encode`] encodes them -/// and once as canonical arrays, and asserts that the running partials and the results are equal. +/// Accumulates `batches` of `dtype` with `aggregate` once as `encode` encodes them and once as +/// canonical arrays, and asserts that the running partials and the results are equal. /// /// On each batch that encodes to a `Widen`, the kernel must return a partial exactly when /// `served`, and the partial must equal the partial of a fresh accumulator on the canonical batch. /// Encoding and each accumulator get their own canonical batch, because both cache statistics on /// the array they read, and an accumulator reads a cached statistic instead of the data. #[track_caller] -fn assert_aggregate_batches( +pub(super) fn assert_aggregate_batches( aggregate: &AggregateFnRef, served: bool, - batches: &[Vec>], + batches: &[T], dtype: &DType, + mut build_batch: impl FnMut(&T, &DType, &mut ExecutionCtx) -> VortexResult, + mut encode: impl FnMut(C, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult<()> { let mut actual = aggregate.accumulator(dtype)?; let mut expected = aggregate.accumulator(dtype)?; for values in batches { - let encoded = WidenArray::encode(canonical_batch(values, dtype, ctx)?, ctx)?; + let encoded = encode(build_batch(values, dtype, ctx)?, ctx)?; if encoded.is::() { let partial = kernel_partial(aggregate, &encoded, ctx)?; @@ -140,13 +143,13 @@ fn assert_aggregate_batches( if let Some(partial) = partial { let mut reference = aggregate.accumulator(dtype)?; - reference.accumulate(&canonical_batch(values, dtype, ctx)?.into_array(), ctx)?; + reference.accumulate(&build_batch(values, dtype, ctx)?.into_array(), ctx)?; assert_eq!(partial, reference.flush()?, "{aggregate} on {values:?}"); } } actual.accumulate(&encoded, ctx)?; - expected.accumulate(&canonical_batch(values, dtype, ctx)?.into_array(), ctx)?; + expected.accumulate(&build_batch(values, dtype, ctx)?.into_array(), ctx)?; assert_eq!( actual.partial_scalar()?, expected.partial_scalar()?, @@ -238,7 +241,15 @@ fn test_aggregate_partials( let dtype = DType::Primitive(logical, nullability); for (aggregate, served) in aggregates() { - assert_aggregate_batches(&aggregate, served, &batches, &dtype, &mut ctx)?; + assert_aggregate_batches( + &aggregate, + served, + &batches, + &dtype, + |values, dtype, ctx| canonical_batch(values, dtype, ctx), + WidenArray::encode, + &mut ctx, + )?; } Ok(()) diff --git a/vortex-array/src/arrays/widen/tests/decimal_aggregates.rs b/vortex-array/src/arrays/widen/tests/decimal_aggregates.rs new file mode 100644 index 00000000000..d6e60518275 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_aggregates.rs @@ -0,0 +1,248 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal aggregate equivalence, including the precision-dependent sum fallback. +//! +//! Each batch uses fresh canonical arrays so cached statistics cannot hide the kernel path. + +use rstest::rstest; +use vortex_error::VortexResult; + +use super::SESSION; +use super::aggregates::assert_aggregate_batches; +use super::aggregates::kernel_partial; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VTable; +use crate::VortexSessionExecute; +use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnVTableExt; +use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::NumericalAggregateOpts; +use crate::aggregate_fn::fns::first::First; +use crate::aggregate_fn::fns::is_constant::IsConstant; +use crate::aggregate_fn::fns::is_sorted::IsSorted; +use crate::aggregate_fn::fns::is_sorted::IsSortedOptions; +use crate::aggregate_fn::fns::last::Last; +use crate::aggregate_fn::fns::max::Max; +use crate::aggregate_fn::fns::min::Min; +use crate::aggregate_fn::fns::min_max::MinMax; +use crate::aggregate_fn::fns::sum::Sum; +use crate::aggregate_fn::fns::sum_v2::SumV2; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::dtype::BigCast; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::dtype::Nullability::NonNullable; +use crate::dtype::Nullability::Nullable; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; + +fn aggregates() -> Vec<(AggregateFnRef, bool)> { + let mut aggregates = vec![ + (IsConstant.bind(EmptyOptions), true), + (IsSorted.bind(IsSortedOptions { strict: false }), true), + (IsSorted.bind(IsSortedOptions { strict: true }), true), + (First.bind(EmptyOptions), false), + (Last.bind(EmptyOptions), false), + ]; + + for skip_nans in [true, false] { + let options = NumericalAggregateOpts { skip_nans }; + aggregates.extend([ + (MinMax.bind(options), true), + (Min.bind(options), true), + (Max.bind(options), true), + (Sum.bind(options), false), + (SumV2.bind(options), false), + ]); + } + + aggregates +} + +fn canonical_batch(values: &[Option], dtype: &DType) -> DecimalArray { + let decimal = *dtype + .as_decimal_opt() + .expect("decimal fixture has a decimal dtype"); + + match_each_decimal_value_type!(decimal.canonical_values_type(), |T| { + let values = values.iter().map(|value| { + value.map(|value| { + ::from(value).expect("fixture values fit the parent canonical width") + }) + }); + match dtype.nullability() { + Nullable => DecimalArray::from_option_iter(values, decimal), + NonNullable => DecimalArray::from_iter( + values.map(|value| value.expect("non-nullable fixture contains only valid values")), + decimal, + ), + } + }) +} + +/// Every fixture must exercise the `Widen` path, including empty and all-null batches. +#[track_caller] +fn encode_batch(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { + let encoded = WidenArray::encode_decimal(array, ctx)?; + assert!(encoded.is::()); + + Ok(encoded) +} + +#[rstest] +#[case::extrema( + vec![ + vec![Some(-99), None, Some(99)], // + vec![Some(-9999), Some(9999)], // + vec![Some(-999_999_999), Some(999_999_999)], // + vec![Some(-999_999_999_999_999_999), Some(999_999_999_999_999_999)], // + vec![Some(-(10i128.pow(38) - 1)), Some(10i128.pow(38) - 1)], // + ], + Nullable, +)] +#[case::ascending( + vec![vec![Some(-99), Some(-1)], vec![Some(0), Some(99)]], + NonNullable, +)] +#[case::descending_boundary( + vec![vec![Some(0), Some(99)], vec![Some(-99), Some(-1)]], + Nullable, +)] +#[case::unsorted_batch(vec![vec![Some(99), Some(-99)]], NonNullable)] +#[case::equal_boundary(vec![vec![Some(7)], vec![Some(7)]], Nullable)] +#[case::constant(vec![vec![Some(7), Some(7)], vec![Some(7)]], NonNullable)] +#[case::changed_constant(vec![vec![Some(7), Some(7)], vec![Some(8)]], Nullable)] +#[case::constant_with_null(vec![vec![Some(7), None, Some(7)]], Nullable)] +#[case::empty(vec![vec![], vec![]], Nullable)] +#[case::all_null(vec![vec![None], vec![None, None]], Nullable)] +#[case::single_null(vec![vec![None]], Nullable)] +#[case::single_null_then_sorted(vec![vec![None], vec![Some(1), Some(2)]], Nullable)] +#[case::empty_then_constant(vec![vec![], vec![Some(7), Some(7)]], Nullable)] +#[case::leading_and_trailing_nulls( + vec![vec![None, Some(-7)], vec![Some(7), None]], + Nullable, +)] +fn test_decimal_aggregate_partials( + #[case] batches: Vec>>, + #[case] nullability: Nullability, + #[values(-3, 0, 2)] scale: i8, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dtype = DType::Decimal(DecimalDType::new(76, scale), nullability); + + for (aggregate, served) in aggregates() { + assert_aggregate_batches( + &aggregate, + served, + &batches, + &dtype, + |values, dtype, _ctx| Ok(canonical_batch(values, dtype)), + encode_batch, + &mut ctx, + )?; + } + + Ok(()) +} + +#[rstest] +#[case::sum_skip_nans(Sum.bind(NumericalAggregateOpts { skip_nans: true }))] +#[case::sum_include_nans(Sum.bind(NumericalAggregateOpts { skip_nans: false }))] +#[case::sum_v2_skip_nans(SumV2.bind(NumericalAggregateOpts { skip_nans: true }))] +#[case::sum_v2_include_nans(SumV2.bind(NumericalAggregateOpts { skip_nans: false }))] +fn test_decimal_sum_uses_logical_precision(#[case] aggregate: AggregateFnRef) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let child_decimal = DecimalDType::new(2, 0); + let parent_decimal = DecimalDType::new(18, 0); + let parent_dtype = DType::Decimal(parent_decimal, NonNullable); + let child_scalar = Scalar::decimal(DecimalValue::I8(99), child_decimal, NonNullable); + let parent_scalar = Scalar::decimal(DecimalValue::I64(99), parent_decimal, NonNullable); + let len = 1usize << 34; + let huge_child = ConstantArray::new(child_scalar, len).into_array(); + let huge_widen = WidenArray::try_new(huge_child.clone(), parent_dtype.clone())?.into_array(); + + // Calling the registered kernel directly must decline without executing this huge array. + let kernel = ctx + .aggregate_kernels() + .find_aggregate_kernel(Widen.id(), aggregate.id()); + if aggregate.is::() { + let kernel = kernel.expect("the session registers the Widen Sum kernel"); + assert!( + kernel + .aggregate(&aggregate, &huge_widen, &mut ctx)? + .is_none() + ); + } else { + assert!(kernel.is_none(), "SumV2 has no Widen kernel"); + } + + let mut parent = aggregate.accumulator(&parent_dtype)?; + parent.accumulate( + &ConstantArray::new(parent_scalar, len).into_array(), + &mut ctx, + )?; + let expected = Scalar::decimal( + DecimalValue::I128(1_700_807_049_216), + DecimalDType::new(28, 0), + Nullable, + ); + assert_eq!(parent.finish()?, expected); + + let mut child = aggregate.accumulator(huge_child.dtype())?; + child.accumulate(&huge_child, &mut ctx)?; + assert_eq!( + child.finish()?, + Scalar::null(DType::Decimal(DecimalDType::new(12, 0), Nullable)) + ); + + let batches = vec![ + vec![], // + vec![Some(99); 256], // + vec![None, Some(-99)], // + ]; + assert_aggregate_batches( + &aggregate, + false, + &batches, + &parent_dtype.as_nullable(), + |values, dtype, _ctx| Ok(canonical_batch(values, dtype)), + encode_batch, + &mut ctx, + )?; + + let values = DecimalArray::from_iter([99i8], child_decimal); + let widen = WidenArray::try_new(values.into_array(), parent_dtype.clone())?.into_array(); + assert!(kernel_partial(&aggregate, &widen, &mut ctx)?.is_none()); + let canonical = DecimalArray::from_iter([99i64], parent_decimal).into_array(); + let mut actual = aggregate.accumulator(&parent_dtype)?; + let mut reference = aggregate.accumulator(&parent_dtype)?; + actual.accumulate(&widen, &mut ctx)?; + reference.accumulate(&canonical, &mut ctx)?; + assert_eq!(actual.partial_scalar()?, reference.partial_scalar()?); + + // Doubling real batch partials reaches the same logical length without allocating its rows. + for doubling in 1..=34 { + let actual_partial = actual.partial_scalar()?; + let reference_partial = reference.partial_scalar()?; + actual.combine_partials(&actual_partial)?; + reference.combine_partials(&reference_partial)?; + assert_eq!( + actual.partial_scalar()?, + reference.partial_scalar()?, + "{aggregate} after {doubling} doublings" + ); + } + + assert_eq!(actual.finish()?, expected); + assert_eq!(reference.finish()?, expected); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/mod.rs b/vortex-array/src/arrays/widen/tests/mod.rs index ff457a5833c..7280243c8cb 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -9,6 +9,7 @@ mod aggregates; mod compute; mod decimal; +mod decimal_aggregates; mod decimal_compare; mod encoding; mod numeric; From ef6125743d1e2d0c34f7022d75366f59a9eaf69f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:24:22 -0400 Subject: [PATCH 10/12] Keep same-scale decimal widening casts lazy Signed-off-by: Connor Tsui --- .../src/arrays/decimal/compute/cast.rs | 41 ++- vortex-array/src/arrays/widen/compute/cast.rs | 7 + .../src/arrays/widen/tests/decimal_cast.rs | 257 ++++++++++++++++++ vortex-array/src/arrays/widen/tests/mod.rs | 1 + 4 files changed, 303 insertions(+), 3 deletions(-) create mode 100644 vortex-array/src/arrays/widen/tests/decimal_cast.rs diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index c8f561af0b7..c3fe5894fde 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -21,6 +21,7 @@ use crate::arrays::ConstantArray; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; +use crate::arrays::WidenArray; use crate::arrays::decimal::DecimalArrayExt; use crate::dtype::BigCast; use crate::dtype::DType; @@ -44,8 +45,6 @@ use crate::validity::Validity; impl CastReduce for Decimal { fn cast(array: ArrayView<'_, Decimal>, dtype: &DType) -> VortexResult> { - // Only nullability changes within the same decimal dtype are reducible without execution. - // Precision/scale changes need the kernel. let DType::Decimal(to_decimal_dtype, to_nullability) = dtype else { return Ok(None); }; @@ -56,7 +55,8 @@ impl CastReduce for Decimal { ); }; - if from_decimal_dtype != to_decimal_dtype { + let widens_values = widens_canonical_decimal_values(*from_decimal_dtype, *to_decimal_dtype); + if from_decimal_dtype != to_decimal_dtype && !widens_values { return Ok(None); } @@ -67,6 +67,10 @@ impl CastReduce for Decimal { return Ok(None); }; + if widens_values { + return widen_decimal_values(array, dtype, new_validity).map(Some); + } + // SAFETY: validity has the same length, only its nullability tag changes. unsafe { Ok(Some( @@ -121,6 +125,10 @@ impl CastKernel for Decimal { .clone() .cast_nullability(*to_nullability, array.len(), ctx)?; + if widens_canonical_decimal_values(*from_decimal_dtype, *to_decimal_dtype) { + return widen_decimal_values(array, dtype, new_validity).map(Some); + } + // Reuse the values buffer untouched when no rescale is required, the target precision // only widens (so every value still fits), and the current physical type is already wide // enough to hold the target precision. This keeps the common precision-widening cast @@ -169,6 +177,33 @@ impl CastKernel for Decimal { } } +/// Returns whether a same-scale precision increase needs wider canonical storage. +fn widens_canonical_decimal_values(from: DecimalDType, to: DecimalDType) -> bool { + from.scale() == to.scale() + && from.precision() < to.precision() + && from.canonical_values_type().byte_width() < to.canonical_values_type().byte_width() +} + +/// Gives the child the target nullability before constructing a decimal `Widen`. +fn widen_decimal_values( + array: ArrayView<'_, Decimal>, + dtype: &DType, + validity: Validity, +) -> VortexResult { + // SAFETY: values and decimal dtype are unchanged, and the caller cast validity without + // changing its length. + let values = unsafe { + DecimalArray::new_unchecked_handle( + array.buffer_handle().clone(), + array.values_type(), + array.decimal_dtype(), + validity, + ) + }; + + Ok(WidenArray::try_new(values.into_array(), dtype.clone())?.into_array()) +} + fn cast_to_integer( array: ArrayView<'_, Decimal>, nullability: Nullability, diff --git a/vortex-array/src/arrays/widen/compute/cast.rs b/vortex-array/src/arrays/widen/compute/cast.rs index 6f7131b2f90..771f89316af 100644 --- a/vortex-array/src/arrays/widen/compute/cast.rs +++ b/vortex-array/src/arrays/widen/compute/cast.rs @@ -16,6 +16,13 @@ use crate::scalar_fn::fns::cast::CastReduce; impl CastReduce for Widen { fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult> { let values = array.values(); + if matches!(values.dtype(), DType::Decimal(..)) { + return match dtype { + DType::Decimal(..) | DType::Primitive(..) => values.cast(dtype.clone()).map(Some), + _ => Ok(None), + }; + } + let (DType::Primitive(target_ptype, nullability), DType::Primitive(values_ptype, _)) = (dtype, values.dtype()) else { diff --git a/vortex-array/src/arrays/widen/tests/decimal_cast.rs b/vortex-array/src/arrays/widen/tests/decimal_cast.rs new file mode 100644 index 00000000000..dc228aea727 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_cast.rs @@ -0,0 +1,257 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal casts preserve lazy widening, checked precision and nullability. +//! +//! These cases also cover decimal buffers whose stored width differs from their canonical width. + +use rstest::rstest; +use vortex_buffer::buffer; +use vortex_error::VortexResult; + +use super::SESSION; +use super::canonical; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::arrays::BoolArray; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::DictArray; +use crate::arrays::Widen; +use crate::arrays::WidenArray; +use crate::arrays::decimal::DecimalArrayExt; +use crate::arrays::widen::WidenArraySlotsExt; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::dtype::i256; +use crate::scalar_fn::fns::cast::CastKernel; +use crate::scalar_fn::fns::cast::CastReduce; +use crate::validity::Validity; + +#[rstest] +#[case::i8(DecimalArray::from_iter([12i8, -12], DecimalDType::new(2, 0)), 4)] +#[case::i16(DecimalArray::from_iter([12i16, -12], DecimalDType::new(4, 0)), 9)] +#[case::i32(DecimalArray::from_iter([12i32, -12], DecimalDType::new(9, 0)), 18)] +#[case::i64(DecimalArray::from_iter([12i64, -12], DecimalDType::new(18, 0)), 38)] +#[case::i128(DecimalArray::from_iter([12i128, -12], DecimalDType::new(38, 0)), 76)] +#[case::stored_target_width(DecimalArray::from_iter([12i16, -12], DecimalDType::new(2, 0)), 4)] +#[case::stored_wider_than_target(DecimalArray::from_iter([12i128, -12], DecimalDType::new(2, 0)), 4)] +#[case::stored_i256(DecimalArray::from_iter([i256::from_i128(12), i256::from_i128(-12)], DecimalDType::new(2, 0)), 4)] +fn test_precision_growth_is_lazy( + #[case] array: DecimalArray, + #[case] precision: u8, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source_dtype = array.dtype().clone(); + let source_values_type = array.values_type(); + let decimal_dtype = DecimalDType::new(precision, 0); + let target = DType::Decimal(decimal_dtype, Nullability::NonNullable); + + let casted = array.into_array().cast(target.clone())?; + + assert!(casted.is::()); + let widen = casted.as_::(); + let values = widen.values(); + assert_eq!(values.dtype(), &source_dtype); + assert_eq!(values.as_::().values_type(), source_values_type); + let executed = casted.execute::(&mut ctx)?; + assert_eq!(executed.dtype(), &target); + assert_eq!( + executed.values_type(), + decimal_dtype.canonical_values_type() + ); + let expected = DecimalArray::from_iter([12i128, -12], decimal_dtype); + assert_arrays_eq!(executed, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::to_nullable(Validity::NonNullable, Nullability::Nullable)] +#[case::to_non_nullable(Validity::AllValid, Nullability::NonNullable)] +#[case::keeps_nullable(Validity::AllValid, Nullability::Nullable)] +fn test_precision_growth_sets_child_nullability( + #[case] validity: Validity, + #[case] nullability: Nullability, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = DecimalArray::new(buffer![12i8, -12], DecimalDType::new(2, 0), validity); + let target_decimal = DecimalDType::new(4, 0); + let target = DType::Decimal(target_decimal, nullability); + + let casted = array.into_array().cast(target.clone())?; + + assert!(casted.is::()); + assert_eq!( + casted.as_::().values().dtype().nullability(), + nullability + ); + assert_eq!(casted.dtype(), &target); + let expected_validity = match nullability { + Nullability::Nullable => Validity::AllValid, + Nullability::NonNullable => Validity::NonNullable, + }; + let expected = DecimalArray::new(buffer![12i16, -12], target_decimal, expected_validity); + assert_arrays_eq!(casted, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::all_valid(true)] +#[case::null_row(false)] +fn test_encoded_validity_falls_back_to_kernel(#[case] second_valid: bool) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let validity = DictArray::try_new( + buffer![0u8, 1].into_array(), + BoolArray::from_iter([true, second_valid]).into_array(), + )?; + let array = DecimalArray::new( + buffer![12i8, -12], + DecimalDType::new(2, 0), + Validity::Array(validity.into_array()), + ); + let target_decimal = DecimalDType::new(4, 0); + let target = DType::Decimal(target_decimal, Nullability::NonNullable); + + assert!(::cast(array.as_view(), &target)?.is_none()); + let result = ::cast(array.as_view(), &target, &mut ctx); + if !second_valid { + let error = result.expect_err("the encoded validity contains a null row"); + assert!(error.to_string().contains("invalid values to non-nullable")); + return Ok(()); + } + + let casted = result?.expect("same-scale precision growth has a cast kernel"); + assert!(casted.is::()); + assert_eq!( + casted.as_::().values().dtype().nullability(), + Nullability::NonNullable + ); + let expected = DecimalArray::from_iter([12i16, -12], target_decimal); + assert_arrays_eq!(casted, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_same_canonical_width_reuses_stored_buffer() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = DecimalArray::from_iter([12i128, -12], DecimalDType::new(10, 0)); + let source_ptr = array.buffer::().as_ptr(); + let target_decimal = DecimalDType::new(18, 0); + let target = DType::Decimal(target_decimal, Nullability::NonNullable); + + let casted = array + .into_array() + .cast(target)? + .execute::(&mut ctx)?; + + assert_eq!(casted.buffer::().as_ptr(), source_ptr); + assert_arrays_eq!( + casted, + DecimalArray::from_iter([12i128, -12], target_decimal), + &mut ctx + ); + + Ok(()) +} + +#[rstest] +#[case::smaller_decimal(DType::Decimal(DecimalDType::new(3, 2), Nullability::Nullable))] +#[case::wider_decimal(DType::Decimal(DecimalDType::new(38, 2), Nullability::Nullable))] +#[case::rescale(DType::Decimal(DecimalDType::new(18, 3), Nullability::Nullable))] +#[case::integer(DType::Primitive(PType::I64, Nullability::Nullable))] +#[case::float(DType::Primitive(PType::F64, Nullability::Nullable))] +fn test_decimal_widen_cast_matches_canonical(#[case] target: DType) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = + DecimalArray::from_option_iter([Some(125i16), None, Some(-250)], DecimalDType::new(4, 2)); + let array = WidenArray::try_new( + values.into_array(), + DType::Decimal(DecimalDType::new(18, 2), Nullability::Nullable), + )? + .into_array(); + let expected = canonical(&array, &mut ctx)?.cast(target.clone())?; + + let casted = array.cast(target)?; + + assert_arrays_eq!(casted, expected, &mut ctx); + + Ok(()) +} + +#[test] +fn test_precision_narrowing_checks_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::from_iter([1000i16], DecimalDType::new(4, 0)); + let array = WidenArray::try_new( + values.into_array(), + DType::Decimal(DecimalDType::new(18, 0), Nullability::NonNullable), + )? + .into_array(); + let target = DType::Decimal(DecimalDType::new(3, 0), Nullability::NonNullable); + + let error = array + .cast(target)? + .execute::(&mut ctx) + .expect_err("1000 does not fit a precision of three digits"); + + assert!(error.to_string().contains("does not fit in precision")); + + Ok(()) +} + +#[test] +fn test_repeated_casts_terminate_without_nested_widen() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let source = DecimalArray::from_iter([12i8, -12], DecimalDType::new(2, 0)).into_array(); + let mut array = source.clone(); + for precision in [4, 9, 18, 38, 76, 4, 2] { + let decimal_dtype = DecimalDType::new(precision, 0); + let target = DType::Decimal(decimal_dtype, Nullability::NonNullable); + array = array.cast(target.clone())?; + assert!(ArrayRef::ptr_eq(&array.cast(target)?, &array)); + if let Some(widen) = array.as_opt::() { + assert!(!widen.values().is::()); + assert_eq!(widen.values().dtype(), source.dtype()); + } + assert_arrays_eq!( + array.clone(), + DecimalArray::from_iter([12i128, -12], decimal_dtype), + &mut ctx + ); + } + + assert_arrays_eq!(array, source, &mut ctx); + + Ok(()) +} + +#[test] +fn test_same_canonical_width_copies_narrow_stored_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = DecimalArray::from_iter([125i32, -250], DecimalDType::new(10, 2)); + let target_decimal = DecimalDType::new(18, 2); + let target = DType::Decimal(target_decimal, Nullability::NonNullable); + let casted = array.into_array().cast(target)?; + assert!(!casted.is::()); + + let executed = casted.execute::(&mut ctx)?; + assert_eq!( + executed.values_type(), + target_decimal.canonical_values_type() + ); + assert_arrays_eq!( + executed, + DecimalArray::from_iter([125i64, -250], target_decimal), + &mut ctx + ); + + Ok(()) +} diff --git a/vortex-array/src/arrays/widen/tests/mod.rs b/vortex-array/src/arrays/widen/tests/mod.rs index 7280243c8cb..8d4b560bdc9 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -10,6 +10,7 @@ mod aggregates; mod compute; mod decimal; mod decimal_aggregates; +mod decimal_cast; mod decimal_compare; mod encoding; mod numeric; From 011dbc8f4e5b62721dd777d54778218a3c8538eb Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 15:26:52 -0400 Subject: [PATCH 11/12] Cover decimal Widen fallbacks and default file writes Signed-off-by: Connor Tsui --- .../src/arrays/widen/tests/decimal.rs | 4 +- .../arrays/widen/tests/decimal_fallback.rs | 188 ++++++++++++++++++ vortex-array/src/arrays/widen/tests/mod.rs | 1 + vortex-file/src/tests.rs | 33 +++ 4 files changed, 224 insertions(+), 2 deletions(-) create mode 100644 vortex-array/src/arrays/widen/tests/decimal_fallback.rs diff --git a/vortex-array/src/arrays/widen/tests/decimal.rs b/vortex-array/src/arrays/widen/tests/decimal.rs index 658ce2384bc..1d7c2b58dba 100644 --- a/vortex-array/src/arrays/widen/tests/decimal.rs +++ b/vortex-array/src/arrays/widen/tests/decimal.rs @@ -110,12 +110,12 @@ fn test_decimal_execute_widens_and_terminates( DecimalDType::new(18, 2), validity, ); - let widened = WidenArray::try_new(child.clone().into_array(), parent.clone())?.into_array(); + let widened = WidenArray::try_new(child.into_array(), parent.clone())?.into_array(); let executed = widened.clone().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.clone(), &mut ctx); + assert_arrays_eq!(executed, reference, &mut ctx); assert_arrays_eq!(widened, reference, &mut ctx); 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..ae9ff486248 --- /dev/null +++ b/vortex-array/src/arrays/widen/tests/decimal_fallback.rs @@ -0,0 +1,188 @@ +// 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] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_arithmetic_falls_back_to_canonical( + #[case] validity: Validity, + #[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)?; + 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_eq!(actual.dtype(), expected.dtype()); + 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_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_chunked_execution_matches_canonical( + #[case] validity: Validity, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(validity)?; + 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_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +fn test_decimal_interleave_falls_back_to_canonical(#[case] validity: Validity) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = decimal_widen(validity)?; + 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_eq!(actual.dtype(), expected.dtype()); + 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 8d4b560bdc9..7ad3a1ef5cf 100644 --- a/vortex-array/src/arrays/widen/tests/mod.rs +++ b/vortex-array/src/arrays/widen/tests/mod.rs @@ -12,6 +12,7 @@ mod decimal; mod decimal_aggregates; mod decimal_cast; mod decimal_compare; +mod decimal_fallback; mod encoding; mod numeric; diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index ce4d7e54ba5..7e1becc2a30 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -35,6 +35,8 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::WidenArray; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; @@ -258,6 +260,37 @@ async fn test_read_simple() { assert_arrays_eq!(result, st, &mut SESSION.create_execution_ctx()); } +#[rstest] +#[case::non_nullable(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false, true, true, false]))] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_decimal_widen_default_writer_roundtrip( + #[case] validity: Validity, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_dtype = DecimalDType::new(18, 2); + let values = DecimalArray::new( + buffer![-125i16, 50, 0, 250, 900], + DecimalDType::new(4, 2), + validity.clone(), + ); + let dtype = DType::Decimal(decimal_dtype, validity.nullability()); + let widened = WidenArray::try_new(values.into_array(), dtype)?.into_array(); + assert!(widened.is::()); + let array = StructArray::from_fields(&[("decimal", widened)])?.into_array(); + let expected_values = + DecimalArray::new(buffer![-125i64, 50, 0, 250, 900], decimal_dtype, validity); + let expected = StructArray::from_fields(&[("decimal", expected_values.into_array())])?; + + let actual = write_read_roundtrip(array).await?; + + assert_eq!(actual.dtype(), expected.dtype()); + assert_arrays_eq!(actual, expected, &mut ctx); + + Ok(()) +} + /// 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. From 9577faf6749402386f0c2415b4b4e6e01fb2f3c0 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 15:54:47 -0400 Subject: [PATCH 12/12] Tidy decimal Widen execution, casts and tests Move decimal execution into its own function, with a note on why it copies instead of yielding a cast. The leaf cast compares canonical widths directly and says which casts reduce without execution. `smallest_decimal_dtype` takes a `RangeInclusive`, like `smallest_integer_ptype`. The two-`Widen` compare test now covers children with different canonical widths. Signed-off-by: Connor Tsui --- .../src/arrays/decimal/compute/cast.rs | 9 ++-- vortex-array/src/arrays/widen/encode.rs | 10 ++-- .../src/arrays/widen/tests/decimal.rs | 11 ++-- .../src/arrays/widen/tests/decimal_compare.rs | 30 +++++++---- vortex-array/src/arrays/widen/vtable/mod.rs | 50 ++++++++++++------- 5 files changed, 66 insertions(+), 44 deletions(-) diff --git a/vortex-array/src/arrays/decimal/compute/cast.rs b/vortex-array/src/arrays/decimal/compute/cast.rs index c3fe5894fde..af7e2dfd757 100644 --- a/vortex-array/src/arrays/decimal/compute/cast.rs +++ b/vortex-array/src/arrays/decimal/compute/cast.rs @@ -55,6 +55,8 @@ impl CastReduce for Decimal { ); }; + // Without execution, only a nullability change and a same-scale cast that grows the + // canonical width reduce. Every other change of the decimal dtype needs the kernel. let widens_values = widens_canonical_decimal_values(*from_decimal_dtype, *to_decimal_dtype); if from_decimal_dtype != to_decimal_dtype && !widens_values { return Ok(None); @@ -177,11 +179,10 @@ impl CastKernel for Decimal { } } -/// Returns whether a same-scale precision increase needs wider canonical storage. +/// Returns whether a cast keeps the scale and grows the canonical width, which implies a higher +/// precision. fn widens_canonical_decimal_values(from: DecimalDType, to: DecimalDType) -> bool { - from.scale() == to.scale() - && from.precision() < to.precision() - && from.canonical_values_type().byte_width() < to.canonical_values_type().byte_width() + from.scale() == to.scale() && from.canonical_values_type() < to.canonical_values_type() } /// Gives the child the target nullability before constructing a decimal `Widen`. diff --git a/vortex-array/src/arrays/widen/encode.rs b/vortex-array/src/arrays/widen/encode.rs index 3c437abdf24..90d3d966ce3 100644 --- a/vortex-array/src/arrays/widen/encode.rs +++ b/vortex-array/src/arrays/widen/encode.rs @@ -55,10 +55,7 @@ impl WidenArray { let mask = validity.execute_mask(array.len(), ctx)?; let range = match_each_decimal_value_type!(array.values_type(), |T| { valid_range(array.buffer::().as_slice(), &mask).map(|(min, max)| { - ( - DecimalValue::from(min).as_i256(), - DecimalValue::from(max).as_i256(), - ) + DecimalValue::from(min).as_i256()..=DecimalValue::from(max).as_i256() }) }); let child_dtype = smallest_decimal_dtype(range, parent)?; @@ -174,13 +171,14 @@ pub(super) fn integer_value(scalar: &Scalar) -> VortexResult { /// /// Returns an error if the minimum exceeds the maximum or an endpoint exceeds the parent precision. pub fn smallest_decimal_dtype( - range: Option<(i256, i256)>, + range: Option>, parent: DecimalDType, ) -> VortexResult { let minimum_precision = parent.scale().max(1) as u8; - let Some((min, max)) = range else { + let Some(range) = range else { 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}" diff --git a/vortex-array/src/arrays/widen/tests/decimal.rs b/vortex-array/src/arrays/widen/tests/decimal.rs index 1d7c2b58dba..dc17440ec22 100644 --- a/vortex-array/src/arrays/widen/tests/decimal.rs +++ b/vortex-array/src/arrays/widen/tests/decimal.rs @@ -16,6 +16,7 @@ use crate::Canonical; use crate::EmptyArrayData; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::Widen; use crate::arrays::WidenArray; @@ -113,7 +114,7 @@ fn test_decimal_execute_widens_and_terminates( let widened = WidenArray::try_new(child.into_array(), parent.clone())?.into_array(); let executed = widened.clone().execute::(&mut ctx)?.into_array(); assert_eq!(executed.dtype(), &parent); - let decimal = executed.as_::(); + let decimal = executed.as_::(); assert_eq!(decimal.values_type(), DecimalType::I64); assert_arrays_eq!(executed, reference, &mut ctx); assert_arrays_eq!(widened, reference, &mut ctx); @@ -140,7 +141,7 @@ fn test_encode_decimal_best_precision( assert_eq!(encoded.is::(), is_widen); if let Some(widen) = encoded.as_opt::() { assert_eq!(stored_decimal_dtype(&encoded).precision(), child_precision); - let child = widen.values().as_::(); + let child = widen.values().as_::(); assert_eq!( child.values_type(), child.decimal_dtype().canonical_values_type() @@ -184,7 +185,7 @@ fn test_smallest_decimal_dtype( #[case] scale: i8, #[case] precision: Option, ) -> VortexResult<()> { - let range = range.map(|(min, max)| (i256::from_i128(min.into()), i256::from_i128(max.into()))); + let range = range.map(|(min, max)| i256::from_i128(min.into())..=i256::from_i128(max.into())); let result = smallest_decimal_dtype(range, DecimalDType::new(parent_precision, scale)); match precision { Some(precision) => assert_eq!(result?, DecimalDType::new(precision, scale)), @@ -197,9 +198,9 @@ fn test_smallest_decimal_dtype( #[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])); + 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()); + assert!(smallest_decimal_dtype(Some(i256::MIN..=i256::MAX), parent).is_err()); Ok(()) } diff --git a/vortex-array/src/arrays/widen/tests/decimal_compare.rs b/vortex-array/src/arrays/widen/tests/decimal_compare.rs index dc5d8b25979..56cc208d4c9 100644 --- a/vortex-array/src/arrays/widen/tests/decimal_compare.rs +++ b/vortex-array/src/arrays/widen/tests/decimal_compare.rs @@ -110,8 +110,26 @@ fn test_compare_constant( assert_compare_matches_canonical(&array, &constant, operator, &mut ctx) } +/// In `different_widths`, the cast of the `i8` child to the common precision returns a `Widen` of +/// its own, so the compare runs through the fallback. #[rstest] +#[case::same_width( + DecimalArray::from_iter([-99_999i32, -1, 0, 99_999], DecimalDType::new(5, 2)), + DecimalArray::from_option_iter( + [Some(-100_000i32), None, Some(0), Some(999_999_999)], + DecimalDType::new(9, 2), + ), +)] +#[case::different_widths( + DecimalArray::from_iter([-99i8, -1, 0, 99], DecimalDType::new(2, 2)), + DecimalArray::from_option_iter( + [Some(-100i32), None, Some(0), Some(999_999_999)], + DecimalDType::new(9, 2), + ), +)] fn test_compare_two_widen( + #[case] lhs: DecimalArray, + #[case] rhs: DecimalArray, #[values( Operator::Eq, Operator::NotEq, @@ -123,16 +141,8 @@ fn test_compare_two_widen( operator: Operator, ) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); - let lhs = widen(DecimalArray::from_iter( - [-99_999i32, -1, 0, 99_999], - DecimalDType::new(5, 2), - ))? - .into_array(); - let rhs = widen(DecimalArray::from_option_iter( - [Some(-100_000i32), None, Some(0), Some(999_999_999)], - DecimalDType::new(9, 2), - ))? - .into_array(); + let lhs = widen(lhs)?.into_array(); + let rhs = widen(rhs)?.into_array(); assert_compare_matches_canonical(&lhs, &rhs, operator, &mut ctx) } diff --git a/vortex-array/src/arrays/widen/vtable/mod.rs b/vortex-array/src/arrays/widen/vtable/mod.rs index 37fdb4b47d2..1ee74c6b931 100644 --- a/vortex-array/src/arrays/widen/vtable/mod.rs +++ b/vortex-array/src/arrays/widen/vtable/mod.rs @@ -34,6 +34,7 @@ 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::serde::ArrayChildren; @@ -117,26 +118,10 @@ impl VTable for Widen { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { + // NB: Decimals copy their values instead of yielding a cast. A same-scale decimal cast that + // grows the canonical width returns a `Widen`, so a yielded cast would execute to itself. if let DType::Decimal(decimal_dtype, _) = array.dtype() { - let values = array.values().clone().execute::(ctx)?; - let validity = values.validity()?; - let mask = validity.execute_mask(values.len(), ctx)?; - let buffer = match_each_decimal_value_type!(values.values_type(), |T| { - copy_to_canonical_width(values.buffer::().as_slice(), &mask, *decimal_dtype)? - }); - - // SAFETY: the copy has the parent canonical width and the child length. Widen - // construction guarantees the same scale and a parent precision that holds the child, - // and the unchanged child validity has the same length. - let widened = unsafe { - DecimalArray::new_unchecked_handle( - buffer, - decimal_dtype.canonical_values_type(), - *decimal_dtype, - validity, - ) - }; - + let widened = execute_decimal(&array, *decimal_dtype, ctx)?; return Ok(ExecutionResult::done(widened.into_array())); } @@ -163,3 +148,30 @@ impl VTable for Widen { PARENT_RULES.evaluate(array, parent, child_idx) } } + +/// Executes the values of `array` and copies them into the canonical width of `decimal_dtype`, the +/// logical dtype of `array`. +fn execute_decimal( + array: &WidenArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let values = array.values().clone().execute::(ctx)?; + let validity = values.validity()?; + let mask = validity.execute_mask(values.len(), ctx)?; + let buffer = match_each_decimal_value_type!(values.values_type(), |T| { + copy_to_canonical_width(values.buffer::().as_slice(), &mask, decimal_dtype)? + }); + + // SAFETY: the copy has the parent canonical width and the child length. Widen construction + // guarantees the same scale and a parent precision that holds the child, and the unchanged + // child validity has the same length. + Ok(unsafe { + DecimalArray::new_unchecked_handle( + buffer, + decimal_dtype.canonical_values_type(), + decimal_dtype, + validity, + ) + }) +}