From e2254624fc86e3a3bf4b566ea44c6a3f4be38168 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 29 Oct 2026 14:23:30 -0400 Subject: [PATCH 1/3] Export decimal Widen arrays to Arrow with one copy Export the narrower decimal child directly with the logical precision and scale, avoiding the intermediate canonical-width copy. Skip null payloads when narrowing while still rejecting valid values outside the Arrow width. Keep decimal export coverage in a separate test module. Signed-off-by: Connor Tsui --- vortex-arrow/src/executor/decimal.rs | 422 --------------------- vortex-arrow/src/executor/decimal/mod.rs | 259 +++++++++++++ vortex-arrow/src/executor/decimal/tests.rs | 289 ++++++++++++++ 3 files changed, 548 insertions(+), 422 deletions(-) delete mode 100644 vortex-arrow/src/executor/decimal.rs create mode 100644 vortex-arrow/src/executor/decimal/mod.rs create mode 100644 vortex-arrow/src/executor/decimal/tests.rs diff --git a/vortex-arrow/src/executor/decimal.rs b/vortex-arrow/src/executor/decimal.rs deleted file mode 100644 index d86406277bd..00000000000 --- a/vortex-arrow/src/executor/decimal.rs +++ /dev/null @@ -1,422 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use arrow_array::ArrayRef as ArrowArrayRef; -use arrow_array::Decimal32Array as ArrowDecimal32Array; -use arrow_array::Decimal64Array as ArrowDecimal64Array; -use arrow_array::Decimal128Array as ArrowDecimal128Array; -use arrow_array::Decimal256Array as ArrowDecimal256Array; -use arrow_buffer::i256; -use arrow_schema::DataType; -use itertools::Itertools; -use num_traits::AsPrimitive; -use num_traits::ToPrimitive; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::arrays::DecimalArray; -use vortex_array::dtype::DType; -use vortex_array::dtype::DecimalType; -use vortex_buffer::Buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; - -use crate::null_buffer::to_null_buffer; - -pub(super) fn to_arrow_decimal( - array: ArrayRef, - data_type: &DataType, - ctx: &mut ExecutionCtx, -) -> VortexResult { - vortex_ensure!( - matches!(array.dtype(), DType::Decimal(..)), - "Cannot convert Vortex array with dtype {} to an Arrow {} array", - array.dtype(), - data_type - ); - - // Execute the array as a DecimalArray. - let decimal_array = array.execute::(ctx)?; - - match data_type { - DataType::Decimal32(..) => to_arrow_decimal32(decimal_array, ctx), - DataType::Decimal64(..) => to_arrow_decimal64(decimal_array, ctx), - DataType::Decimal128(..) => to_arrow_decimal128(decimal_array, ctx), - DataType::Decimal256(..) => to_arrow_decimal256(decimal_array, ctx), - _ => unreachable!("to_arrow_decimal called with non-decimal type"), - } -} - -fn to_arrow_decimal32(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { - let null_buffer = to_null_buffer( - array - .as_ref() - .validity()? - .execute_mask(array.as_ref().len(), ctx)?, - ); - let buffer: Buffer = match array.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => array.buffer::(), - DecimalType::I64 => array - .buffer::() - .into_iter() - .map(|x| { - x.to_i32() - .ok_or_else(|| vortex_err!("i64 to i32 narrowing cannot be done safely")) - }) - .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, - DecimalType::I128 => array - .buffer::() - .into_iter() - .map(|x| { - x.to_i32() - .ok_or_else(|| vortex_err!("i128 to i32 narrowing cannot be done safely")) - }) - .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, - DecimalType::I256 => array - .buffer::() - .into_iter() - .map(|x| { - x.to_i32() - .ok_or_else(|| vortex_err!("i256 to i32 narrowing cannot be done safely")) - }) - .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, - }; - Ok(Arc::new( - ArrowDecimal32Array::new(buffer.into_arrow_scalar_buffer(), null_buffer) - .with_precision_and_scale( - array.decimal_dtype().precision(), - array.decimal_dtype().scale(), - )?, - )) -} - -fn to_arrow_decimal64(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { - let null_buffer = to_null_buffer( - array - .as_ref() - .validity()? - .execute_mask(array.as_ref().len(), ctx)?, - ); - let buffer: Buffer = match array.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I64 => array.buffer::(), - DecimalType::I128 => array - .buffer::() - .into_iter() - .map(|x| { - x.to_i64() - .ok_or_else(|| vortex_err!("i128 to i64 narrowing cannot be done safely")) - }) - .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, - DecimalType::I256 => array - .buffer::() - .into_iter() - .map(|x| { - x.to_i64() - .ok_or_else(|| vortex_err!("i256 to i64 narrowing cannot be done safely")) - }) - .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, - }; - Ok(Arc::new( - ArrowDecimal64Array::new(buffer.into_arrow_scalar_buffer(), null_buffer) - .with_precision_and_scale( - array.decimal_dtype().precision(), - array.decimal_dtype().scale(), - )?, - )) -} - -fn to_arrow_decimal128(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { - let null_buffer = to_null_buffer( - array - .as_ref() - .validity()? - .execute_mask(array.as_ref().len(), ctx)?, - ); - let buffer: Buffer = match array.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I64 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I128 => array.buffer::(), - DecimalType::I256 => array - .buffer::() - .into_iter() - .map(|x| { - x.to_i128() - .ok_or_else(|| vortex_err!("i256 to i128 narrowing cannot be done safely")) - }) - .process_results(|iter| Buffer::from_trusted_len_iter(iter))?, - }; - Ok(Arc::new( - ArrowDecimal128Array::new(buffer.into_arrow_scalar_buffer(), null_buffer) - .with_precision_and_scale( - array.decimal_dtype().precision(), - array.decimal_dtype().scale(), - )?, - )) -} - -fn to_arrow_decimal256(array: DecimalArray, ctx: &mut ExecutionCtx) -> VortexResult { - let null_buffer = to_null_buffer( - array - .as_ref() - .validity()? - .execute_mask(array.as_ref().len(), ctx)?, - ); - let buffer: Buffer = match array.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I64 => { - Buffer::from_trusted_len_iter(array.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I128 => Buffer::from_trusted_len_iter( - array - .buffer::() - .into_iter() - .map(|x| vortex_array::dtype::i256::from_i128(x).into()), - ), - DecimalType::I256 => { - Buffer::::from_byte_buffer(array.buffer_handle().clone().into_host_sync()) - } - }; - Ok(Arc::new( - ArrowDecimal256Array::new(buffer.into_arrow_scalar_buffer(), null_buffer) - .with_precision_and_scale( - array.decimal_dtype().precision(), - array.decimal_dtype().scale(), - )?, - )) -} - -#[cfg(test)] -mod tests { - use std::sync::LazyLock; - - use arrow_array::Array; - use arrow_array::Decimal128Array; - use arrow_array::Decimal256Array; - use arrow_buffer::i256; - use arrow_schema::DataType; - use rstest::rstest; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::builders::ArrayBuilder; - use vortex_array::builders::DecimalBuilder; - use vortex_array::dtype::DecimalDType; - use vortex_array::dtype::NativeDecimalType; - use vortex_array::validity::Validity; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_session::VortexSession; - - use crate::ArrowArrayExecutor; - use crate::executor::decimal::DecimalArray; - - static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - - #[test] - fn decimal_to_arrow() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - // Make a very simple i128 and i256 array. - let decimal_vortex = DecimalArray::new( - buffer![1i128, 2i128, 3i128, 4i128, 5i128], - DecimalDType::new(19, 2), - Validity::NonNullable, - ); - let arrow = decimal_vortex - .into_array() - .execute_arrow(Some(&DataType::Decimal128(19, 2)), &mut ctx)?; - assert_eq!(arrow.data_type(), &DataType::Decimal128(19, 2)); - let decimal_array = arrow.as_any().downcast_ref::().unwrap(); - assert_eq!( - decimal_array.values().as_ref(), - &[1i128, 2i128, 3i128, 4i128, 5i128] - ); - assert_eq!( - decimal_array, - &Decimal128Array::from(vec![1, 2, 3, 4, 5]).with_precision_and_scale(19, 2)? - ); - Ok(()) - } - - #[rstest] - #[case(0i8)] - #[case(0i16)] - #[case(0i32)] - #[case(0i64)] - #[case(0i128)] - #[case(vortex_array::dtype::i256::ZERO)] - fn test_to_arrow_decimal128( - #[case] _decimal_type: T, - ) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let mut decimal = DecimalBuilder::new_in::( - DecimalDType::new(2, 1), - false.into(), - vortex_buffer::BufferAllocatorRef::static_ref(), - ); - decimal.append_value(10); - decimal.append_value(11); - decimal.append_value(12); - let decimal = decimal.finish(); - - let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal128(2, 1)), &mut ctx)?; - let arrow_decimal = arrow_array - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(arrow_decimal.value(0), 10); - assert_eq!(arrow_decimal.value(1), 11); - assert_eq!(arrow_decimal.value(2), 12); - assert_eq!( - arrow_decimal, - &Decimal128Array::from(vec![10, 11, 12]).with_precision_and_scale(2, 1)? - ); - Ok(()) - } - - #[rstest] - #[case(0i8)] - #[case(0i16)] - #[case(0i32)] - #[case(0i64)] - #[case(0i128)] - #[case(vortex_array::dtype::i256::ZERO)] - fn test_to_arrow_decimal32(#[case] _decimal_type: T) -> VortexResult<()> { - use arrow_array::Decimal32Array; - - let mut ctx = SESSION.create_execution_ctx(); - let mut decimal = DecimalBuilder::new_in::( - DecimalDType::new(2, 1), - false.into(), - vortex_buffer::BufferAllocatorRef::static_ref(), - ); - decimal.append_value(10); - decimal.append_value(11); - decimal.append_value(12); - let decimal = decimal.finish(); - - let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal32(2, 1)), &mut ctx)?; - let arrow_decimal = arrow_array - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(arrow_decimal.value(0), 10); - assert_eq!(arrow_decimal.value(1), 11); - assert_eq!(arrow_decimal.value(2), 12); - assert_eq!( - arrow_decimal, - &Decimal32Array::from(vec![10, 11, 12]).with_precision_and_scale(2, 1)? - ); - Ok(()) - } - - #[rstest] - #[case(0i8)] - #[case(0i16)] - #[case(0i32)] - #[case(0i64)] - #[case(0i128)] - #[case(vortex_array::dtype::i256::ZERO)] - fn test_to_arrow_decimal64(#[case] _decimal_type: T) -> VortexResult<()> { - use arrow_array::Decimal64Array; - - let mut ctx = SESSION.create_execution_ctx(); - let mut decimal = DecimalBuilder::new_in::( - DecimalDType::new(2, 1), - false.into(), - vortex_buffer::BufferAllocatorRef::static_ref(), - ); - decimal.append_value(10); - decimal.append_value(11); - decimal.append_value(12); - let decimal = decimal.finish(); - - let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal64(2, 1)), &mut ctx)?; - let arrow_decimal = arrow_array - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(arrow_decimal.value(0), 10); - assert_eq!(arrow_decimal.value(1), 11); - assert_eq!(arrow_decimal.value(2), 12); - assert_eq!( - arrow_decimal, - &Decimal64Array::from(vec![10, 11, 12]).with_precision_and_scale(2, 1)? - ); - Ok(()) - } - - #[rstest] - #[case(0i8)] - #[case(0i16)] - #[case(0i32)] - #[case(0i64)] - #[case(0i128)] - #[case(vortex_array::dtype::i256::ZERO)] - fn test_to_arrow_decimal256( - #[case] _decimal_type: T, - ) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let mut decimal = DecimalBuilder::new_in::( - DecimalDType::new(2, 1), - false.into(), - vortex_buffer::BufferAllocatorRef::static_ref(), - ); - decimal.append_value(10); - decimal.append_value(11); - decimal.append_value(12); - let decimal = decimal.finish(); - - let arrow_array = decimal.execute_arrow(Some(&DataType::Decimal256(2, 1)), &mut ctx)?; - let arrow_decimal = arrow_array - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(arrow_decimal.value(0), i256::from_i128(10)); - assert_eq!(arrow_decimal.value(1), i256::from_i128(11)); - assert_eq!(arrow_decimal.value(2), i256::from_i128(12)); - assert_eq!( - arrow_decimal, - &Decimal256Array::from(vec![ - i256::from_i128(10), - i256::from_i128(11), - i256::from_i128(12) - ]) - .with_precision_and_scale(2, 1)? - ); - Ok(()) - } -} diff --git a/vortex-arrow/src/executor/decimal/mod.rs b/vortex-arrow/src/executor/decimal/mod.rs new file mode 100644 index 00000000000..41d44fdafd7 --- /dev/null +++ b/vortex-arrow/src/executor/decimal/mod.rs @@ -0,0 +1,259 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal export from stored values with the logical precision and scale. +//! +//! Exporting a `Widen` child directly avoids an intermediate widening copy. Narrowing conversions +//! ignore null payloads because those bytes do not constrain the Arrow value width. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::Decimal32Array as ArrowDecimal32Array; +use arrow_array::Decimal64Array as ArrowDecimal64Array; +use arrow_array::Decimal128Array as ArrowDecimal128Array; +use arrow_array::Decimal256Array as ArrowDecimal256Array; +use arrow_buffer::i256; +use arrow_schema::DataType; +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::widen::WidenArraySlotsExt; +use vortex_array::dtype::BigCast; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativeDecimalType; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::null_buffer::to_null_buffer; + +pub(super) fn to_arrow_decimal( + array: ArrayRef, + data_type: &DataType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let decimal_dtype = *array.dtype().as_decimal_opt().ok_or_else(|| { + vortex_err!( + "Cannot convert Vortex array with dtype {} to an Arrow {} array", + array.dtype(), + data_type + ) + })?; + + // A `Widen` holds the same values in a narrower child. Exporting the child copies each value + // once, into the Arrow width. Executing the `Widen` first would also copy each value into the + // canonical width of `decimal_dtype`. + let array = array.execute_until::(ctx)?; + let values = match array.as_opt::() { + Some(widen) => widen.values().clone(), + None => array, + }; + let values = values.execute::(ctx)?; + + match data_type { + DataType::Decimal32(..) => to_arrow_decimal32(values, decimal_dtype, ctx), + DataType::Decimal64(..) => to_arrow_decimal64(values, decimal_dtype, ctx), + DataType::Decimal128(..) => to_arrow_decimal128(values, decimal_dtype, ctx), + DataType::Decimal256(..) => to_arrow_decimal256(values, decimal_dtype, ctx), + _ => unreachable!("to_arrow_decimal called with non-decimal type"), + } +} + +/// Exports `values` as an Arrow `Decimal32` with the precision and scale of `decimal_dtype`. +/// +/// `values` can be the narrower child of a [`Widen`], whose own decimal dtype has a smaller +/// precision. The Arrow precision and scale therefore come from `decimal_dtype`, never from +/// `values`. +fn to_arrow_decimal32( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => values.buffer::(), + DecimalType::I64 => narrow_values(values.buffer::(), &validity)?, + DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, + DecimalType::I256 => { + narrow_values(values.buffer::(), &validity)? + } + }; + + Ok(Arc::new( + ArrowDecimal32Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal64`. +fn to_arrow_decimal64( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I64 => values.buffer::(), + DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, + DecimalType::I256 => { + narrow_values(values.buffer::(), &validity)? + } + }; + + Ok(Arc::new( + ArrowDecimal64Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal128`. +fn to_arrow_decimal128( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I64 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I128 => values.buffer::(), + DecimalType::I256 => { + narrow_values(values.buffer::(), &validity)? + } + }; + + Ok(Arc::new( + ArrowDecimal128Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal256`. +fn to_arrow_decimal256( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I64 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I128 => Buffer::from_trusted_len_iter( + values + .buffer::() + .into_iter() + .map(|x| vortex_array::dtype::i256::from_i128(x).into()), + ), + DecimalType::I256 => { + Buffer::::from_byte_buffer(values.buffer_handle().clone().into_host_sync()) + } + }; + + Ok(Arc::new( + ArrowDecimal256Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Narrows the decimal `values` to `T`, and writes zero into every null row. +/// +/// Every valid value **must** fit in `T`, or this returns an error. A null row can hold any value, +/// including one that does not fit in `T`, so this never converts it. The zeroed null rows match +/// [`copy_to_canonical_width`], but unlike it, this checks only that each valid value fits in `T`, +/// not that it fits the precision. +/// +/// [`copy_to_canonical_width`]: vortex_array::arrays::decimal::copy_to_canonical_width +fn narrow_values( + values: Buffer, + validity: &Mask, +) -> VortexResult> { + let narrow = |value: S| { + ::from(value).ok_or_else(|| { + vortex_err!( + "every valid {} value must fit in {} to narrow, got {value}", + S::DECIMAL_TYPE, + T::DECIMAL_TYPE, + ) + }) + }; + + match validity.bit_buffer() { + AllOr::All => values + .into_iter() + .map(narrow) + .process_results(|iter| Buffer::from_trusted_len_iter(iter)), + AllOr::None => Ok(Buffer::zeroed(values.len())), + AllOr::Some(is_valid) => values + .into_iter() + .zip(is_valid.iter()) + .map(|(value, is_valid)| { + if is_valid { + narrow(value) + } else { + Ok(T::default()) + } + }) + .process_results(|iter| Buffer::from_trusted_len_iter(iter)), + } +} + +#[cfg(test)] +mod tests; diff --git a/vortex-arrow/src/executor/decimal/tests.rs b/vortex-arrow/src/executor/decimal/tests.rs new file mode 100644 index 00000000000..c67da312228 --- /dev/null +++ b/vortex-arrow/src/executor/decimal/tests.rs @@ -0,0 +1,289 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal export equivalence and width checks for canonical arrays and `Widen` children. + +use std::sync::Arc; +use std::sync::LazyLock; + +use arrow_array::Array; +use arrow_array::Decimal32Array; +use arrow_array::Decimal64Array; +use arrow_array::Decimal128Array; +use arrow_array::Decimal256Array; +use arrow_array::StructArray as ArrowStructArray; +use arrow_buffer::i256; +use arrow_schema::DataType; +use rstest::rstest; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::WidenArray; +use vortex_array::builders::ArrayBuilder; +use vortex_array::builders::DecimalBuilder; +use vortex_array::dtype::DType; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::NativeDecimalType; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_panic; +use vortex_session::VortexSession; + +use crate::ArrowArrayExecutor; + +static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + +/// Checks the concrete Arrow array type and compares its dtype, values, and nulls. +#[track_caller] +fn assert_arrow_eq(actual: &arrow_array::ArrayRef, expected: T) { + assert!(actual.as_any().is::()); + let expected: arrow_array::ArrayRef = Arc::new(expected); + assert_eq!(actual, &expected); +} + +/// Presents `values` as a `Widen` with the logical `precision`, and the scale and nullability +/// of `values`. +fn widen(values: DecimalArray, precision: u8) -> VortexResult { + let decimal_dtype = DecimalDType::new(precision, values.decimal_dtype().scale()); + let dtype = DType::Decimal(decimal_dtype, values.dtype().nullability()); + + Ok(WidenArray::try_new(values.into_array(), dtype)?.into_array()) +} + +/// A child at the largest precision that `T` holds, with the smallest and largest values of +/// that precision in its valid rows. +/// +/// A nullable child holds `garbage` in its null rows. Arrow equality ignores null rows, so the +/// export must match the canonical export whatever the null rows hold. +fn widen_child(garbage: T, nullable: bool) -> DecimalArray { + let decimal_dtype = DecimalDType::new(T::MAX_PRECISION, 1); + let min = T::MIN_BY_PRECISION[usize::from(T::MAX_PRECISION)]; + let max = T::MAX_BY_PRECISION[usize::from(T::MAX_PRECISION)]; + + if nullable { + DecimalArray::new( + Buffer::from_iter([min, garbage, max, min, garbage]), + decimal_dtype, + Validity::from_iter([true, false, true, true, false]), + ) + } else { + DecimalArray::new( + Buffer::from_iter([min, max, max, min, max]), + decimal_dtype, + Validity::NonNullable, + ) + } +} + +#[test] +fn decimal_to_arrow() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_vortex = DecimalArray::new( + buffer![1i128, 2i128, 3i128, 4i128, 5i128], + DecimalDType::new(19, 2), + Validity::NonNullable, + ); + let arrow = decimal_vortex + .into_array() + .execute_arrow(Some(&DataType::Decimal128(19, 2)), &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal128Array::from(vec![1, 2, 3, 4, 5]).with_precision_and_scale(19, 2)?, + ); + + Ok(()) +} + +#[rstest] +#[case::i8(0i8)] +#[case::i16(0i16)] +#[case::i32(0i32)] +#[case::i64(0i64)] +#[case::i128(0i128)] +#[case::i256(vortex_array::dtype::i256::ZERO)] +fn test_to_arrow_decimal( + #[case] _decimal_type: T, + #[values( + DataType::Decimal32(2, 1), + DataType::Decimal64(2, 1), + DataType::Decimal128(2, 1), + DataType::Decimal256(2, 1) + )] + data_type: DataType, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut decimal = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + decimal.append_value(10); + decimal.append_value(11); + decimal.append_value(12); + let decimal = decimal.finish(); + + let actual = decimal.execute_arrow(Some(&data_type), &mut ctx)?; + let expected: arrow_array::ArrayRef = match data_type { + DataType::Decimal32(precision, scale) => Arc::new( + Decimal32Array::from(vec![10, 11, 12]).with_precision_and_scale(precision, scale)?, + ), + DataType::Decimal64(precision, scale) => Arc::new( + Decimal64Array::from(vec![10, 11, 12]).with_precision_and_scale(precision, scale)?, + ), + DataType::Decimal128(precision, scale) => Arc::new( + Decimal128Array::from(vec![10, 11, 12]).with_precision_and_scale(precision, scale)?, + ), + DataType::Decimal256(precision, scale) => Arc::new( + Decimal256Array::from(vec![ + i256::from_i128(10), + i256::from_i128(11), + i256::from_i128(12), + ]) + .with_precision_and_scale(precision, scale)?, + ), + data_type => vortex_panic!("the test cases only select decimal types, got {data_type}"), + }; + assert_eq!(&actual, &expected); + + Ok(()) +} + +#[rstest] +#[case::i8_p4_decimal32(i8::MAX, 4, DataType::Decimal32(4, 1))] +#[case::i8_p4_decimal64(i8::MAX, 4, DataType::Decimal64(4, 1))] +#[case::i8_p4_decimal128(i8::MAX, 4, DataType::Decimal128(4, 1))] +#[case::i8_p4_decimal256(i8::MAX, 4, DataType::Decimal256(4, 1))] +#[case::i32_p18_decimal64(i32::MAX, 18, DataType::Decimal64(18, 1))] +#[case::i32_p18_decimal128(i32::MAX, 18, DataType::Decimal128(18, 1))] +#[case::i32_p18_decimal256(i32::MAX, 18, DataType::Decimal256(18, 1))] +#[case::i64_p38_decimal128(i64::MAX, 38, DataType::Decimal128(38, 1))] +#[case::i64_p38_decimal256(i64::MAX, 38, DataType::Decimal256(38, 1))] +#[case::i128_p76_decimal256(i128::MAX, 76, DataType::Decimal256(76, 1))] +fn test_widen_to_arrow_matches_canonical( + #[case] garbage: T, + #[case] precision: u8, + #[case] data_type: DataType, + #[values(false, true)] nullable: bool, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = widen(widen_child(garbage, nullable), precision)?; + let canonical = array + .clone() + .execute::(&mut ctx)? + .into_array(); + + let actual = array.execute_arrow(Some(&data_type), &mut ctx)?; + let expected = canonical.execute_arrow(Some(&data_type), &mut ctx)?; + assert_eq!(actual.data_type(), &data_type); + assert_eq!(&actual, &expected); + + Ok(()) +} + +#[test] +fn test_widen_to_arrow_uses_logical_precision() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new( + buffer![-99i8, 0, 99], + DecimalDType::new(2, 1), + Validity::NonNullable, + ); + + let arrow = widen(values, 4)?.execute_arrow(None, &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal128Array::from(vec![-99, 0, 99]).with_precision_and_scale(4, 1)?, + ); + + Ok(()) +} + +#[test] +fn test_sliced_widen_to_arrow() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new( + buffer![10i32, i32::MAX, 30, 40, 50], + DecimalDType::new(9, 2), + Validity::from_iter([true, false, true, true, true]), + ); + + let sliced = widen(values, 18)?.slice(1..4)?; + let arrow = sliced.execute_arrow(Some(&DataType::Decimal128(18, 2)), &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal128Array::from(vec![None, Some(30), Some(40)]).with_precision_and_scale(18, 2)?, + ); + + Ok(()) +} + +#[test] +fn test_struct_with_widen_field_to_arrow() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new( + buffer![1i32, 2, 3], + DecimalDType::new(9, 2), + Validity::NonNullable, + ); + let array = StructArray::from_fields(&[("price", widen(values, 18)?)])?.into_array(); + + let arrow = array.execute_arrow(None, &mut ctx)?; + let arrow = arrow + .as_any() + .downcast_ref::() + .vortex_expect("the requested Arrow type fixes the concrete array type"); + assert_arrow_eq( + arrow.column(0), + Decimal128Array::from(vec![1, 2, 3]).with_precision_and_scale(18, 2)?, + ); + + Ok(()) +} + +/// A child stored wider than its canonical width, with a value that does not fit in the Arrow +/// width in a null row. Some or all of the rows are null. The export narrows the child, so it +/// must skip the null rows. +#[rstest] +#[case::partially_null(Validity::from_iter([true, false]), vec![Some(1), None])] +#[case::all_null(Validity::AllInvalid, vec![None, None])] +fn test_widen_to_arrow_skips_null_rows_when_narrowing( + #[case] validity: Validity, + #[case] expected: Vec>, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new(buffer![1i128, i128::MAX], DecimalDType::new(9, 2), validity); + + let arrow = widen(values, 18)?.execute_arrow(Some(&DataType::Decimal64(18, 2)), &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal64Array::from(expected).with_precision_and_scale(18, 2)?, + ); + + Ok(()) +} + +#[rstest] +#[case::all_valid(Validity::NonNullable)] +#[case::partially_valid(Validity::from_iter([true, false]))] +fn test_to_arrow_rejects_valid_value_that_does_not_fit( + #[case] validity: Validity, +) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new(buffer![i128::MAX, 0], DecimalDType::new(9, 2), validity); + + let err = values + .into_array() + .execute_arrow(Some(&DataType::Decimal64(9, 2)), &mut ctx) + .expect_err("a valid value outside the target width must fail"); + assert!( + err.to_string() + .contains("every valid i128 value must fit in i64 to narrow") + ); + + Ok(()) +} From cd94f2446dd66e0960514afd426a5505f7fbdbcb Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 3 Nov 2026 13:38:55 -0500 Subject: [PATCH 2/3] Keep decimal Arrow export tests inline Signed-off-by: Connor Tsui --- vortex-arrow/src/executor/decimal.rs | 552 +++++++++++++++++++++ vortex-arrow/src/executor/decimal/mod.rs | 259 ---------- vortex-arrow/src/executor/decimal/tests.rs | 289 ----------- 3 files changed, 552 insertions(+), 548 deletions(-) create mode 100644 vortex-arrow/src/executor/decimal.rs delete mode 100644 vortex-arrow/src/executor/decimal/mod.rs delete mode 100644 vortex-arrow/src/executor/decimal/tests.rs diff --git a/vortex-arrow/src/executor/decimal.rs b/vortex-arrow/src/executor/decimal.rs new file mode 100644 index 00000000000..1e26039716c --- /dev/null +++ b/vortex-arrow/src/executor/decimal.rs @@ -0,0 +1,552 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decimal export from stored values with the logical precision and scale. +//! +//! Exporting a `Widen` child directly avoids an intermediate widening copy. Narrowing conversions +//! ignore null payloads because those bytes do not constrain the Arrow value width. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::Decimal32Array as ArrowDecimal32Array; +use arrow_array::Decimal64Array as ArrowDecimal64Array; +use arrow_array::Decimal128Array as ArrowDecimal128Array; +use arrow_array::Decimal256Array as ArrowDecimal256Array; +use arrow_buffer::i256; +use arrow_schema::DataType; +use itertools::Itertools; +use num_traits::AsPrimitive; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::Widen; +use vortex_array::arrays::widen::WidenArraySlotsExt; +use vortex_array::dtype::BigCast; +use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::DecimalType; +use vortex_array::dtype::NativeDecimalType; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::null_buffer::to_null_buffer; + +pub(super) fn to_arrow_decimal( + array: ArrayRef, + data_type: &DataType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let decimal_dtype = *array.dtype().as_decimal_opt().ok_or_else(|| { + vortex_err!( + "Cannot convert Vortex array with dtype {} to an Arrow {} array", + array.dtype(), + data_type + ) + })?; + + // A `Widen` holds the same values in a narrower child. Exporting the child copies each value + // once, into the Arrow width. Executing the `Widen` first would also copy each value into the + // canonical width of `decimal_dtype`. + let array = array.execute_until::(ctx)?; + let values = match array.as_opt::() { + Some(widen) => widen.values().clone(), + None => array, + }; + let values = values.execute::(ctx)?; + + match data_type { + DataType::Decimal32(..) => to_arrow_decimal32(values, decimal_dtype, ctx), + DataType::Decimal64(..) => to_arrow_decimal64(values, decimal_dtype, ctx), + DataType::Decimal128(..) => to_arrow_decimal128(values, decimal_dtype, ctx), + DataType::Decimal256(..) => to_arrow_decimal256(values, decimal_dtype, ctx), + _ => unreachable!("to_arrow_decimal called with non-decimal type"), + } +} + +/// Exports `values` as an Arrow `Decimal32` with the precision and scale of `decimal_dtype`. +/// +/// `values` can be the narrower child of a [`Widen`], whose own decimal dtype has a smaller +/// precision. The Arrow precision and scale therefore come from `decimal_dtype`, never from +/// `values`. +fn to_arrow_decimal32( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => values.buffer::(), + DecimalType::I64 => narrow_values(values.buffer::(), &validity)?, + DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, + DecimalType::I256 => { + narrow_values(values.buffer::(), &validity)? + } + }; + + Ok(Arc::new( + ArrowDecimal32Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal64`. +fn to_arrow_decimal64( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I64 => values.buffer::(), + DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, + DecimalType::I256 => { + narrow_values(values.buffer::(), &validity)? + } + }; + + Ok(Arc::new( + ArrowDecimal64Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal128`. +fn to_arrow_decimal128( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I64 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I128 => values.buffer::(), + DecimalType::I256 => { + narrow_values(values.buffer::(), &validity)? + } + }; + + Ok(Arc::new( + ArrowDecimal128Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal256`. +fn to_arrow_decimal256( + values: DecimalArray, + decimal_dtype: DecimalDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let validity = values + .as_ref() + .validity()? + .execute_mask(values.as_ref().len(), ctx)?; + + let buffer: Buffer = match values.values_type() { + DecimalType::I8 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I16 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I32 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I64 => { + Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) + } + DecimalType::I128 => Buffer::from_trusted_len_iter( + values + .buffer::() + .into_iter() + .map(|x| vortex_array::dtype::i256::from_i128(x).into()), + ), + DecimalType::I256 => { + Buffer::::from_byte_buffer(values.buffer_handle().clone().into_host_sync()) + } + }; + + Ok(Arc::new( + ArrowDecimal256Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) + .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, + )) +} + +/// Narrows the decimal `values` to `T`, and writes zero into every null row. +/// +/// Every valid value **must** fit in `T`, or this returns an error. A null row can hold any value, +/// including one that does not fit in `T`, so this never converts it. The zeroed null rows match +/// [`copy_to_canonical_width`], but unlike it, this checks only that each valid value fits in `T`, +/// not that it fits the precision. +/// +/// [`copy_to_canonical_width`]: vortex_array::arrays::decimal::copy_to_canonical_width +fn narrow_values( + values: Buffer, + validity: &Mask, +) -> VortexResult> { + let narrow = |value: S| { + ::from(value).ok_or_else(|| { + vortex_err!( + "every valid {} value must fit in {} to narrow, got {value}", + S::DECIMAL_TYPE, + T::DECIMAL_TYPE, + ) + }) + }; + + match validity.bit_buffer() { + AllOr::All => values + .into_iter() + .map(narrow) + .process_results(|iter| Buffer::from_trusted_len_iter(iter)), + AllOr::None => Ok(Buffer::zeroed(values.len())), + AllOr::Some(is_valid) => values + .into_iter() + .zip(is_valid.iter()) + .map(|(value, is_valid)| { + if is_valid { + narrow(value) + } else { + Ok(T::default()) + } + }) + .process_results(|iter| Buffer::from_trusted_len_iter(iter)), + } +} + +#[cfg(test)] +mod tests { + //! Decimal export equivalence and width checks for canonical arrays and `Widen` children. + + use std::sync::Arc; + use std::sync::LazyLock; + + use arrow_array::Array; + use arrow_array::Decimal32Array; + use arrow_array::Decimal64Array; + use arrow_array::Decimal128Array; + use arrow_array::Decimal256Array; + use arrow_array::StructArray as ArrowStructArray; + use arrow_buffer::i256; + use arrow_schema::DataType; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::DecimalArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::WidenArray; + use vortex_array::builders::ArrayBuilder; + use vortex_array::builders::DecimalBuilder; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::NativeDecimalType; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_error::vortex_panic; + use vortex_session::VortexSession; + + use crate::ArrowArrayExecutor; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + /// Checks the concrete Arrow array type and compares its dtype, values, and nulls. + #[track_caller] + fn assert_arrow_eq(actual: &arrow_array::ArrayRef, expected: T) { + assert!(actual.as_any().is::()); + let expected: arrow_array::ArrayRef = Arc::new(expected); + assert_eq!(actual, &expected); + } + + /// Presents `values` as a `Widen` with the logical `precision`, and the scale and nullability + /// of `values`. + fn widen(values: DecimalArray, precision: u8) -> VortexResult { + let decimal_dtype = DecimalDType::new(precision, values.decimal_dtype().scale()); + let dtype = DType::Decimal(decimal_dtype, values.dtype().nullability()); + + Ok(WidenArray::try_new(values.into_array(), dtype)?.into_array()) + } + + /// A child at the largest precision that `T` holds, with the smallest and largest values of + /// that precision in its valid rows. + /// + /// A nullable child holds `garbage` in its null rows. Arrow equality ignores null rows, so the + /// export must match the canonical export whatever the null rows hold. + fn widen_child(garbage: T, nullable: bool) -> DecimalArray { + let decimal_dtype = DecimalDType::new(T::MAX_PRECISION, 1); + let min = T::MIN_BY_PRECISION[usize::from(T::MAX_PRECISION)]; + let max = T::MAX_BY_PRECISION[usize::from(T::MAX_PRECISION)]; + + if nullable { + DecimalArray::new( + Buffer::from_iter([min, garbage, max, min, garbage]), + decimal_dtype, + Validity::from_iter([true, false, true, true, false]), + ) + } else { + DecimalArray::new( + Buffer::from_iter([min, max, max, min, max]), + decimal_dtype, + Validity::NonNullable, + ) + } + } + + #[test] + fn decimal_to_arrow() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let decimal_vortex = DecimalArray::new( + buffer![1i128, 2i128, 3i128, 4i128, 5i128], + DecimalDType::new(19, 2), + Validity::NonNullable, + ); + let arrow = decimal_vortex + .into_array() + .execute_arrow(Some(&DataType::Decimal128(19, 2)), &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal128Array::from(vec![1, 2, 3, 4, 5]).with_precision_and_scale(19, 2)?, + ); + + Ok(()) + } + + #[rstest] + #[case::i8(0i8)] + #[case::i16(0i16)] + #[case::i32(0i32)] + #[case::i64(0i64)] + #[case::i128(0i128)] + #[case::i256(vortex_array::dtype::i256::ZERO)] + fn test_to_arrow_decimal( + #[case] _decimal_type: T, + #[values( + DataType::Decimal32(2, 1), + DataType::Decimal64(2, 1), + DataType::Decimal128(2, 1), + DataType::Decimal256(2, 1) + )] + data_type: DataType, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut decimal = DecimalBuilder::new_in::( + DecimalDType::new(2, 1), + false.into(), + vortex_buffer::BufferAllocatorRef::static_ref(), + ); + decimal.append_value(10); + decimal.append_value(11); + decimal.append_value(12); + let decimal = decimal.finish(); + + let actual = decimal.execute_arrow(Some(&data_type), &mut ctx)?; + let expected: arrow_array::ArrayRef = match data_type { + DataType::Decimal32(precision, scale) => Arc::new( + Decimal32Array::from(vec![10, 11, 12]) + .with_precision_and_scale(precision, scale)?, + ), + DataType::Decimal64(precision, scale) => Arc::new( + Decimal64Array::from(vec![10, 11, 12]) + .with_precision_and_scale(precision, scale)?, + ), + DataType::Decimal128(precision, scale) => Arc::new( + Decimal128Array::from(vec![10, 11, 12]) + .with_precision_and_scale(precision, scale)?, + ), + DataType::Decimal256(precision, scale) => Arc::new( + Decimal256Array::from(vec![ + i256::from_i128(10), + i256::from_i128(11), + i256::from_i128(12), + ]) + .with_precision_and_scale(precision, scale)?, + ), + data_type => vortex_panic!("the test cases only select decimal types, got {data_type}"), + }; + assert_eq!(&actual, &expected); + + Ok(()) + } + + #[rstest] + #[case::i8_p4_decimal32(i8::MAX, 4, DataType::Decimal32(4, 1))] + #[case::i8_p4_decimal64(i8::MAX, 4, DataType::Decimal64(4, 1))] + #[case::i8_p4_decimal128(i8::MAX, 4, DataType::Decimal128(4, 1))] + #[case::i8_p4_decimal256(i8::MAX, 4, DataType::Decimal256(4, 1))] + #[case::i32_p18_decimal64(i32::MAX, 18, DataType::Decimal64(18, 1))] + #[case::i32_p18_decimal128(i32::MAX, 18, DataType::Decimal128(18, 1))] + #[case::i32_p18_decimal256(i32::MAX, 18, DataType::Decimal256(18, 1))] + #[case::i64_p38_decimal128(i64::MAX, 38, DataType::Decimal128(38, 1))] + #[case::i64_p38_decimal256(i64::MAX, 38, DataType::Decimal256(38, 1))] + #[case::i128_p76_decimal256(i128::MAX, 76, DataType::Decimal256(76, 1))] + fn test_widen_to_arrow_matches_canonical( + #[case] garbage: T, + #[case] precision: u8, + #[case] data_type: DataType, + #[values(false, true)] nullable: bool, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let array = widen(widen_child(garbage, nullable), precision)?; + let canonical = array + .clone() + .execute::(&mut ctx)? + .into_array(); + + let actual = array.execute_arrow(Some(&data_type), &mut ctx)?; + let expected = canonical.execute_arrow(Some(&data_type), &mut ctx)?; + assert_eq!(actual.data_type(), &data_type); + assert_eq!(&actual, &expected); + + Ok(()) + } + + #[test] + fn test_widen_to_arrow_uses_logical_precision() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new( + buffer![-99i8, 0, 99], + DecimalDType::new(2, 1), + Validity::NonNullable, + ); + + let arrow = widen(values, 4)?.execute_arrow(None, &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal128Array::from(vec![-99, 0, 99]).with_precision_and_scale(4, 1)?, + ); + + Ok(()) + } + + #[test] + fn test_sliced_widen_to_arrow() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new( + buffer![10i32, i32::MAX, 30, 40, 50], + DecimalDType::new(9, 2), + Validity::from_iter([true, false, true, true, true]), + ); + + let sliced = widen(values, 18)?.slice(1..4)?; + let arrow = sliced.execute_arrow(Some(&DataType::Decimal128(18, 2)), &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal128Array::from(vec![None, Some(30), Some(40)]) + .with_precision_and_scale(18, 2)?, + ); + + Ok(()) + } + + #[test] + fn test_struct_with_widen_field_to_arrow() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new( + buffer![1i32, 2, 3], + DecimalDType::new(9, 2), + Validity::NonNullable, + ); + let array = StructArray::from_fields(&[("price", widen(values, 18)?)])?.into_array(); + + let arrow = array.execute_arrow(None, &mut ctx)?; + let arrow = arrow + .as_any() + .downcast_ref::() + .vortex_expect("the requested Arrow type fixes the concrete array type"); + assert_arrow_eq( + arrow.column(0), + Decimal128Array::from(vec![1, 2, 3]).with_precision_and_scale(18, 2)?, + ); + + Ok(()) + } + + /// A child stored wider than its canonical width, with a value that does not fit in the Arrow + /// width in a null row. Some or all of the rows are null. The export narrows the child, so it + /// must skip the null rows. + #[rstest] + #[case::partially_null(Validity::from_iter([true, false]), vec![Some(1), None])] + #[case::all_null(Validity::AllInvalid, vec![None, None])] + fn test_widen_to_arrow_skips_null_rows_when_narrowing( + #[case] validity: Validity, + #[case] expected: Vec>, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = + DecimalArray::new(buffer![1i128, i128::MAX], DecimalDType::new(9, 2), validity); + + let arrow = + widen(values, 18)?.execute_arrow(Some(&DataType::Decimal64(18, 2)), &mut ctx)?; + assert_arrow_eq( + &arrow, + Decimal64Array::from(expected).with_precision_and_scale(18, 2)?, + ); + + Ok(()) + } + + #[rstest] + #[case::all_valid(Validity::NonNullable)] + #[case::partially_valid(Validity::from_iter([true, false]))] + fn test_to_arrow_rejects_valid_value_that_does_not_fit( + #[case] validity: Validity, + ) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let values = DecimalArray::new(buffer![i128::MAX, 0], DecimalDType::new(9, 2), validity); + + let err = values + .into_array() + .execute_arrow(Some(&DataType::Decimal64(9, 2)), &mut ctx) + .expect_err("a valid value outside the target width must fail"); + assert!( + err.to_string() + .contains("every valid i128 value must fit in i64 to narrow") + ); + + Ok(()) + } +} diff --git a/vortex-arrow/src/executor/decimal/mod.rs b/vortex-arrow/src/executor/decimal/mod.rs deleted file mode 100644 index 41d44fdafd7..00000000000 --- a/vortex-arrow/src/executor/decimal/mod.rs +++ /dev/null @@ -1,259 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Decimal export from stored values with the logical precision and scale. -//! -//! Exporting a `Widen` child directly avoids an intermediate widening copy. Narrowing conversions -//! ignore null payloads because those bytes do not constrain the Arrow value width. - -use std::sync::Arc; - -use arrow_array::ArrayRef as ArrowArrayRef; -use arrow_array::Decimal32Array as ArrowDecimal32Array; -use arrow_array::Decimal64Array as ArrowDecimal64Array; -use arrow_array::Decimal128Array as ArrowDecimal128Array; -use arrow_array::Decimal256Array as ArrowDecimal256Array; -use arrow_buffer::i256; -use arrow_schema::DataType; -use itertools::Itertools; -use num_traits::AsPrimitive; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::Widen; -use vortex_array::arrays::widen::WidenArraySlotsExt; -use vortex_array::dtype::BigCast; -use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::DecimalType; -use vortex_array::dtype::NativeDecimalType; -use vortex_buffer::Buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::AllOr; -use vortex_mask::Mask; - -use crate::null_buffer::to_null_buffer; - -pub(super) fn to_arrow_decimal( - array: ArrayRef, - data_type: &DataType, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let decimal_dtype = *array.dtype().as_decimal_opt().ok_or_else(|| { - vortex_err!( - "Cannot convert Vortex array with dtype {} to an Arrow {} array", - array.dtype(), - data_type - ) - })?; - - // A `Widen` holds the same values in a narrower child. Exporting the child copies each value - // once, into the Arrow width. Executing the `Widen` first would also copy each value into the - // canonical width of `decimal_dtype`. - let array = array.execute_until::(ctx)?; - let values = match array.as_opt::() { - Some(widen) => widen.values().clone(), - None => array, - }; - let values = values.execute::(ctx)?; - - match data_type { - DataType::Decimal32(..) => to_arrow_decimal32(values, decimal_dtype, ctx), - DataType::Decimal64(..) => to_arrow_decimal64(values, decimal_dtype, ctx), - DataType::Decimal128(..) => to_arrow_decimal128(values, decimal_dtype, ctx), - DataType::Decimal256(..) => to_arrow_decimal256(values, decimal_dtype, ctx), - _ => unreachable!("to_arrow_decimal called with non-decimal type"), - } -} - -/// Exports `values` as an Arrow `Decimal32` with the precision and scale of `decimal_dtype`. -/// -/// `values` can be the narrower child of a [`Widen`], whose own decimal dtype has a smaller -/// precision. The Arrow precision and scale therefore come from `decimal_dtype`, never from -/// `values`. -fn to_arrow_decimal32( - values: DecimalArray, - decimal_dtype: DecimalDType, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let validity = values - .as_ref() - .validity()? - .execute_mask(values.as_ref().len(), ctx)?; - - let buffer: Buffer = match values.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => values.buffer::(), - DecimalType::I64 => narrow_values(values.buffer::(), &validity)?, - DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, - DecimalType::I256 => { - narrow_values(values.buffer::(), &validity)? - } - }; - - Ok(Arc::new( - ArrowDecimal32Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) - .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, - )) -} - -/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal64`. -fn to_arrow_decimal64( - values: DecimalArray, - decimal_dtype: DecimalDType, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let validity = values - .as_ref() - .validity()? - .execute_mask(values.as_ref().len(), ctx)?; - - let buffer: Buffer = match values.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I64 => values.buffer::(), - DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, - DecimalType::I256 => { - narrow_values(values.buffer::(), &validity)? - } - }; - - Ok(Arc::new( - ArrowDecimal64Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) - .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, - )) -} - -/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal128`. -fn to_arrow_decimal128( - values: DecimalArray, - decimal_dtype: DecimalDType, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let validity = values - .as_ref() - .validity()? - .execute_mask(values.as_ref().len(), ctx)?; - - let buffer: Buffer = match values.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I64 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I128 => values.buffer::(), - DecimalType::I256 => { - narrow_values(values.buffer::(), &validity)? - } - }; - - Ok(Arc::new( - ArrowDecimal128Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) - .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, - )) -} - -/// Like [`to_arrow_decimal32`], but exports an Arrow `Decimal256`. -fn to_arrow_decimal256( - values: DecimalArray, - decimal_dtype: DecimalDType, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let validity = values - .as_ref() - .validity()? - .execute_mask(values.as_ref().len(), ctx)?; - - let buffer: Buffer = match values.values_type() { - DecimalType::I8 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I16 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I32 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I64 => { - Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) - } - DecimalType::I128 => Buffer::from_trusted_len_iter( - values - .buffer::() - .into_iter() - .map(|x| vortex_array::dtype::i256::from_i128(x).into()), - ), - DecimalType::I256 => { - Buffer::::from_byte_buffer(values.buffer_handle().clone().into_host_sync()) - } - }; - - Ok(Arc::new( - ArrowDecimal256Array::new(buffer.into_arrow_scalar_buffer(), to_null_buffer(validity)) - .with_precision_and_scale(decimal_dtype.precision(), decimal_dtype.scale())?, - )) -} - -/// Narrows the decimal `values` to `T`, and writes zero into every null row. -/// -/// Every valid value **must** fit in `T`, or this returns an error. A null row can hold any value, -/// including one that does not fit in `T`, so this never converts it. The zeroed null rows match -/// [`copy_to_canonical_width`], but unlike it, this checks only that each valid value fits in `T`, -/// not that it fits the precision. -/// -/// [`copy_to_canonical_width`]: vortex_array::arrays::decimal::copy_to_canonical_width -fn narrow_values( - values: Buffer, - validity: &Mask, -) -> VortexResult> { - let narrow = |value: S| { - ::from(value).ok_or_else(|| { - vortex_err!( - "every valid {} value must fit in {} to narrow, got {value}", - S::DECIMAL_TYPE, - T::DECIMAL_TYPE, - ) - }) - }; - - match validity.bit_buffer() { - AllOr::All => values - .into_iter() - .map(narrow) - .process_results(|iter| Buffer::from_trusted_len_iter(iter)), - AllOr::None => Ok(Buffer::zeroed(values.len())), - AllOr::Some(is_valid) => values - .into_iter() - .zip(is_valid.iter()) - .map(|(value, is_valid)| { - if is_valid { - narrow(value) - } else { - Ok(T::default()) - } - }) - .process_results(|iter| Buffer::from_trusted_len_iter(iter)), - } -} - -#[cfg(test)] -mod tests; diff --git a/vortex-arrow/src/executor/decimal/tests.rs b/vortex-arrow/src/executor/decimal/tests.rs deleted file mode 100644 index c67da312228..00000000000 --- a/vortex-arrow/src/executor/decimal/tests.rs +++ /dev/null @@ -1,289 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Decimal export equivalence and width checks for canonical arrays and `Widen` children. - -use std::sync::Arc; -use std::sync::LazyLock; - -use arrow_array::Array; -use arrow_array::Decimal32Array; -use arrow_array::Decimal64Array; -use arrow_array::Decimal128Array; -use arrow_array::Decimal256Array; -use arrow_array::StructArray as ArrowStructArray; -use arrow_buffer::i256; -use arrow_schema::DataType; -use rstest::rstest; -use vortex_array::ArrayRef; -use vortex_array::IntoArray; -use vortex_array::VortexSessionExecute; -use vortex_array::arrays::DecimalArray; -use vortex_array::arrays::StructArray; -use vortex_array::arrays::WidenArray; -use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::DecimalBuilder; -use vortex_array::dtype::DType; -use vortex_array::dtype::DecimalDType; -use vortex_array::dtype::NativeDecimalType; -use vortex_array::validity::Validity; -use vortex_buffer::Buffer; -use vortex_buffer::buffer; -use vortex_error::VortexExpect; -use vortex_error::VortexResult; -use vortex_error::vortex_panic; -use vortex_session::VortexSession; - -use crate::ArrowArrayExecutor; - -static SESSION: LazyLock = LazyLock::new(VortexSession::empty); - -/// Checks the concrete Arrow array type and compares its dtype, values, and nulls. -#[track_caller] -fn assert_arrow_eq(actual: &arrow_array::ArrayRef, expected: T) { - assert!(actual.as_any().is::()); - let expected: arrow_array::ArrayRef = Arc::new(expected); - assert_eq!(actual, &expected); -} - -/// Presents `values` as a `Widen` with the logical `precision`, and the scale and nullability -/// of `values`. -fn widen(values: DecimalArray, precision: u8) -> VortexResult { - let decimal_dtype = DecimalDType::new(precision, values.decimal_dtype().scale()); - let dtype = DType::Decimal(decimal_dtype, values.dtype().nullability()); - - Ok(WidenArray::try_new(values.into_array(), dtype)?.into_array()) -} - -/// A child at the largest precision that `T` holds, with the smallest and largest values of -/// that precision in its valid rows. -/// -/// A nullable child holds `garbage` in its null rows. Arrow equality ignores null rows, so the -/// export must match the canonical export whatever the null rows hold. -fn widen_child(garbage: T, nullable: bool) -> DecimalArray { - let decimal_dtype = DecimalDType::new(T::MAX_PRECISION, 1); - let min = T::MIN_BY_PRECISION[usize::from(T::MAX_PRECISION)]; - let max = T::MAX_BY_PRECISION[usize::from(T::MAX_PRECISION)]; - - if nullable { - DecimalArray::new( - Buffer::from_iter([min, garbage, max, min, garbage]), - decimal_dtype, - Validity::from_iter([true, false, true, true, false]), - ) - } else { - DecimalArray::new( - Buffer::from_iter([min, max, max, min, max]), - decimal_dtype, - Validity::NonNullable, - ) - } -} - -#[test] -fn decimal_to_arrow() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let decimal_vortex = DecimalArray::new( - buffer![1i128, 2i128, 3i128, 4i128, 5i128], - DecimalDType::new(19, 2), - Validity::NonNullable, - ); - let arrow = decimal_vortex - .into_array() - .execute_arrow(Some(&DataType::Decimal128(19, 2)), &mut ctx)?; - assert_arrow_eq( - &arrow, - Decimal128Array::from(vec![1, 2, 3, 4, 5]).with_precision_and_scale(19, 2)?, - ); - - Ok(()) -} - -#[rstest] -#[case::i8(0i8)] -#[case::i16(0i16)] -#[case::i32(0i32)] -#[case::i64(0i64)] -#[case::i128(0i128)] -#[case::i256(vortex_array::dtype::i256::ZERO)] -fn test_to_arrow_decimal( - #[case] _decimal_type: T, - #[values( - DataType::Decimal32(2, 1), - DataType::Decimal64(2, 1), - DataType::Decimal128(2, 1), - DataType::Decimal256(2, 1) - )] - data_type: DataType, -) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let mut decimal = DecimalBuilder::new_in::( - DecimalDType::new(2, 1), - false.into(), - vortex_buffer::BufferAllocatorRef::static_ref(), - ); - decimal.append_value(10); - decimal.append_value(11); - decimal.append_value(12); - let decimal = decimal.finish(); - - let actual = decimal.execute_arrow(Some(&data_type), &mut ctx)?; - let expected: arrow_array::ArrayRef = match data_type { - DataType::Decimal32(precision, scale) => Arc::new( - Decimal32Array::from(vec![10, 11, 12]).with_precision_and_scale(precision, scale)?, - ), - DataType::Decimal64(precision, scale) => Arc::new( - Decimal64Array::from(vec![10, 11, 12]).with_precision_and_scale(precision, scale)?, - ), - DataType::Decimal128(precision, scale) => Arc::new( - Decimal128Array::from(vec![10, 11, 12]).with_precision_and_scale(precision, scale)?, - ), - DataType::Decimal256(precision, scale) => Arc::new( - Decimal256Array::from(vec![ - i256::from_i128(10), - i256::from_i128(11), - i256::from_i128(12), - ]) - .with_precision_and_scale(precision, scale)?, - ), - data_type => vortex_panic!("the test cases only select decimal types, got {data_type}"), - }; - assert_eq!(&actual, &expected); - - Ok(()) -} - -#[rstest] -#[case::i8_p4_decimal32(i8::MAX, 4, DataType::Decimal32(4, 1))] -#[case::i8_p4_decimal64(i8::MAX, 4, DataType::Decimal64(4, 1))] -#[case::i8_p4_decimal128(i8::MAX, 4, DataType::Decimal128(4, 1))] -#[case::i8_p4_decimal256(i8::MAX, 4, DataType::Decimal256(4, 1))] -#[case::i32_p18_decimal64(i32::MAX, 18, DataType::Decimal64(18, 1))] -#[case::i32_p18_decimal128(i32::MAX, 18, DataType::Decimal128(18, 1))] -#[case::i32_p18_decimal256(i32::MAX, 18, DataType::Decimal256(18, 1))] -#[case::i64_p38_decimal128(i64::MAX, 38, DataType::Decimal128(38, 1))] -#[case::i64_p38_decimal256(i64::MAX, 38, DataType::Decimal256(38, 1))] -#[case::i128_p76_decimal256(i128::MAX, 76, DataType::Decimal256(76, 1))] -fn test_widen_to_arrow_matches_canonical( - #[case] garbage: T, - #[case] precision: u8, - #[case] data_type: DataType, - #[values(false, true)] nullable: bool, -) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let array = widen(widen_child(garbage, nullable), precision)?; - let canonical = array - .clone() - .execute::(&mut ctx)? - .into_array(); - - let actual = array.execute_arrow(Some(&data_type), &mut ctx)?; - let expected = canonical.execute_arrow(Some(&data_type), &mut ctx)?; - assert_eq!(actual.data_type(), &data_type); - assert_eq!(&actual, &expected); - - Ok(()) -} - -#[test] -fn test_widen_to_arrow_uses_logical_precision() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = DecimalArray::new( - buffer![-99i8, 0, 99], - DecimalDType::new(2, 1), - Validity::NonNullable, - ); - - let arrow = widen(values, 4)?.execute_arrow(None, &mut ctx)?; - assert_arrow_eq( - &arrow, - Decimal128Array::from(vec![-99, 0, 99]).with_precision_and_scale(4, 1)?, - ); - - Ok(()) -} - -#[test] -fn test_sliced_widen_to_arrow() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = DecimalArray::new( - buffer![10i32, i32::MAX, 30, 40, 50], - DecimalDType::new(9, 2), - Validity::from_iter([true, false, true, true, true]), - ); - - let sliced = widen(values, 18)?.slice(1..4)?; - let arrow = sliced.execute_arrow(Some(&DataType::Decimal128(18, 2)), &mut ctx)?; - assert_arrow_eq( - &arrow, - Decimal128Array::from(vec![None, Some(30), Some(40)]).with_precision_and_scale(18, 2)?, - ); - - Ok(()) -} - -#[test] -fn test_struct_with_widen_field_to_arrow() -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = DecimalArray::new( - buffer![1i32, 2, 3], - DecimalDType::new(9, 2), - Validity::NonNullable, - ); - let array = StructArray::from_fields(&[("price", widen(values, 18)?)])?.into_array(); - - let arrow = array.execute_arrow(None, &mut ctx)?; - let arrow = arrow - .as_any() - .downcast_ref::() - .vortex_expect("the requested Arrow type fixes the concrete array type"); - assert_arrow_eq( - arrow.column(0), - Decimal128Array::from(vec![1, 2, 3]).with_precision_and_scale(18, 2)?, - ); - - Ok(()) -} - -/// A child stored wider than its canonical width, with a value that does not fit in the Arrow -/// width in a null row. Some or all of the rows are null. The export narrows the child, so it -/// must skip the null rows. -#[rstest] -#[case::partially_null(Validity::from_iter([true, false]), vec![Some(1), None])] -#[case::all_null(Validity::AllInvalid, vec![None, None])] -fn test_widen_to_arrow_skips_null_rows_when_narrowing( - #[case] validity: Validity, - #[case] expected: Vec>, -) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = DecimalArray::new(buffer![1i128, i128::MAX], DecimalDType::new(9, 2), validity); - - let arrow = widen(values, 18)?.execute_arrow(Some(&DataType::Decimal64(18, 2)), &mut ctx)?; - assert_arrow_eq( - &arrow, - Decimal64Array::from(expected).with_precision_and_scale(18, 2)?, - ); - - Ok(()) -} - -#[rstest] -#[case::all_valid(Validity::NonNullable)] -#[case::partially_valid(Validity::from_iter([true, false]))] -fn test_to_arrow_rejects_valid_value_that_does_not_fit( - #[case] validity: Validity, -) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let values = DecimalArray::new(buffer![i128::MAX, 0], DecimalDType::new(9, 2), validity); - - let err = values - .into_array() - .execute_arrow(Some(&DataType::Decimal64(9, 2)), &mut ctx) - .expect_err("a valid value outside the target width must fail"); - assert!( - err.to_string() - .contains("every valid i128 value must fit in i64 to narrow") - ); - - Ok(()) -} From a61ffff2036b7c52dfcc4936e3a0563eeb191a01 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 3 Nov 2026 13:43:43 -0500 Subject: [PATCH 3/3] Make decimal narrowing types explicit Signed-off-by: Connor Tsui --- vortex-arrow/src/executor/decimal.rs | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/vortex-arrow/src/executor/decimal.rs b/vortex-arrow/src/executor/decimal.rs index 1e26039716c..912bdb8f7da 100644 --- a/vortex-arrow/src/executor/decimal.rs +++ b/vortex-arrow/src/executor/decimal.rs @@ -89,11 +89,12 @@ fn to_arrow_decimal32( Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) } DecimalType::I32 => values.buffer::(), - DecimalType::I64 => narrow_values(values.buffer::(), &validity)?, - DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, - DecimalType::I256 => { - narrow_values(values.buffer::(), &validity)? - } + DecimalType::I64 => narrow_values::(values.buffer::(), &validity)?, + DecimalType::I128 => narrow_values::(values.buffer::(), &validity)?, + DecimalType::I256 => narrow_values::( + values.buffer::(), + &validity, + )?, }; Ok(Arc::new( @@ -124,10 +125,11 @@ fn to_arrow_decimal64( Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) } DecimalType::I64 => values.buffer::(), - DecimalType::I128 => narrow_values(values.buffer::(), &validity)?, - DecimalType::I256 => { - narrow_values(values.buffer::(), &validity)? - } + DecimalType::I128 => narrow_values::(values.buffer::(), &validity)?, + DecimalType::I256 => narrow_values::( + values.buffer::(), + &validity, + )?, }; Ok(Arc::new( @@ -161,9 +163,10 @@ fn to_arrow_decimal128( Buffer::from_trusted_len_iter(values.buffer::().into_iter().map(|x| x.as_())) } DecimalType::I128 => values.buffer::(), - DecimalType::I256 => { - narrow_values(values.buffer::(), &validity)? - } + DecimalType::I256 => narrow_values::( + values.buffer::(), + &validity, + )?, }; Ok(Arc::new(