From 3ce4851eb2e5962a9d359711def813fdd8b5316a Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 14:22:54 -0400 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 4e054c4501d8c6b7fa238c8658eaba33f5a9a81c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 28 Oct 2026 17:08:01 -0400 Subject: [PATCH 6/6] Note that the non-canonical decimal Interleave copy goes away Once every `DecimalArray` stores its values at the canonical width of its dtype (#10701), all sources of a decimal `Interleave` share that width and the copy branch cannot run. Signed-off-by: Connor Tsui --- vortex-array/src/arrays/interleave/execute/decimal.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/decimal.rs b/vortex-array/src/arrays/interleave/execute/decimal.rs index 793b17e3cb4..6e6fbd3e492 100644 --- a/vortex-array/src/arrays/interleave/execute/decimal.rs +++ b/vortex-array/src/arrays/interleave/execute/decimal.rs @@ -85,8 +85,10 @@ fn prepare_source( 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 + // TODO(connor): Remove this branch once every `DecimalArray` stores its values at the + // canonical width of its dtype (#10701). Then all sources share the canonical width `T`, + // because they share the dtype up to nullability. Until then, a source can store its values + // at another width. 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| {