diff --git a/arrow-cast/src/cast/decimal.rs b/arrow-cast/src/cast/decimal.rs index a5d704bba8d5..f7049fe18759 100644 --- a/arrow-cast/src/cast/decimal.rs +++ b/arrow-cast/src/cast/decimal.rs @@ -188,7 +188,8 @@ where // then an increase of scale by 3 will have the following effect on the representation: // [xxxxx] -> [xxxxx000], so for the cast to be infallible, the output type // needs to provide at least 8 digits precision - let is_infallible_cast = (input_precision as i8) + delta_scale <= (output_precision as i8); + let is_infallible_cast = I::MAX_PRECISION <= O::MAX_PRECISION + && (input_precision as i8) + delta_scale <= (output_precision as i8); let f_infallible = is_infallible_cast .then_some(move |x| O::Native::from_decimal(x).unwrap().mul_wrapping(mul)); Some((f_fallible, f_infallible)) @@ -259,7 +260,8 @@ where // the output type needs to have at least 3 digits of precision. // e.g. Decimal(5, 3) 99.999 to Decimal(3, 0) will result in 100: // [99999] -> [99] + 1 = [100], a cast to Decimal(2, 0) would not be possible - let is_infallible_cast = (input_precision as i8) - delta_scale < (output_precision as i8); + let is_infallible_cast = I::MAX_PRECISION <= O::MAX_PRECISION + && (input_precision as i8) - delta_scale < (output_precision as i8); let f_infallible = is_infallible_cast.then_some(move |x| f_fallible(x).unwrap()); Some((f_fallible, f_infallible)) } @@ -838,6 +840,39 @@ where mod tests { use super::*; + #[test] + fn decimal_narrowing_checks_values_and_skips_null_payloads() { + use arrow_array::{Decimal128Array, Decimal256Array}; + use arrow_buffer::NullBuffer; + + let options = CastOptions { + safe: false, + ..Default::default() + }; + for scale in [0, 2, 3] { + let input = Decimal256Array::new( + vec![i256::from_i128(12300), i256::MAX].into(), + Some(NullBuffer::from(vec![true, false])), + ) + .with_precision_and_scale(38, 2) + .unwrap(); + let output = + cast_with_options(&input, &DataType::Decimal128(38, scale), &options).unwrap(); + let expected = 123 * 10_i128.pow(scale as u32); + assert_eq!( + output.as_primitive::(), + &Decimal128Array::from(vec![Some(expected), None]) + .with_precision_and_scale(38, scale) + .unwrap() + ); + + let input = Decimal256Array::from(vec![i256::MAX]) + .with_precision_and_scale(38, 2) + .unwrap(); + assert!(cast_with_options(&input, &DataType::Decimal128(38, scale), &options).is_err()); + } + } + #[test] #[expect(deprecated)] fn test_parse_string_to_decimal_native() { diff --git a/parquet-variant-compute/src/variant_array.rs b/parquet-variant-compute/src/variant_array.rs index da051fe549d3..d536806daca4 100644 --- a/parquet-variant-compute/src/variant_array.rs +++ b/parquet-variant-compute/src/variant_array.rs @@ -27,7 +27,7 @@ use arrow::array::{ new_null_array, }; use arrow::buffer::NullBuffer; -use arrow::compute::cast; +use arrow::compute::{CastOptions, cast_with_options}; use arrow::datatypes::{ Date32Type, Decimal32Type, Decimal64Type, Decimal128Type, Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Time64MicrosecondType, @@ -313,6 +313,10 @@ pub struct VariantArray { impl VariantArray { /// Creates a new `VariantArray` from a [`StructArray`]. /// + /// Decimal fields, including Decimal256 with precision at most 38, are narrowed + /// to supported Variant decimal types. Narrowing preserves scale and nulls and + /// returns an error if a value overflows. + /// /// # Arguments /// - `inner` - The underlying [`StructArray`] that contains the variant data. /// @@ -776,6 +780,7 @@ impl ShreddedVariantFieldArray { /// or be a list, large_list, list_view or struct /// pub fn try_new(inner: &dyn Array) -> Result { + let inner = canonicalize_shredded_types(inner)?; let Some(inner_struct) = inner.as_struct_opt() else { return Err(ArrowError::InvalidArgumentError( "Invalid ShreddedVariantFieldArray: requires StructArray as input".to_string(), @@ -1230,7 +1235,14 @@ fn canonicalize_shredded_types(array: &dyn Array) -> Result { { return Ok(Arc::new(array.clone())); // bypass the unnecessary cast } - cast(array, new_type.as_ref()) + cast_with_options( + array, + new_type.as_ref(), + &CastOptions { + safe: false, + ..Default::default() + }, + ) } /// Recursively visits a data type, ensuring that it only contains data types that can legally @@ -1272,14 +1284,19 @@ fn canonicalize_and_verify_data_type_impl( // // NOTE: arrow-parquet reads widens 32- and 64-bit decimals to 128-bit, but the variant spec // requires using the narrowest decimal type for a given precision. Fix those up first. - Decimal64(p, s) | Decimal128(p, s) + Decimal64(p, s) | Decimal128(p, s) | Decimal256(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => { Cow::Owned(Decimal32(*p, *s)) } - Decimal128(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => { + Decimal128(p, s) | Decimal256(p, s) + if VariantDecimal8::is_valid_precision_and_scale(p, s) => + { Cow::Owned(Decimal64(*p, *s)) } + Decimal256(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => { + Cow::Owned(Decimal128(*p, *s)) + } Decimal32(p, s) if VariantDecimal4::is_valid_precision_and_scale(p, s) => borrow!(), Decimal64(p, s) if VariantDecimal8::is_valid_precision_and_scale(p, s) => borrow!(), Decimal128(p, s) if VariantDecimal16::is_valid_precision_and_scale(p, s) => borrow!(), @@ -1407,9 +1424,9 @@ mod test { use super::*; use arrow::array::{ BinaryArray, BinaryDictionaryBuilder, BinaryRunBuilder, BinaryViewArray, Decimal32Array, - Decimal64Array, Decimal128Array, FixedSizeBinaryArray, FixedSizeListArray, Int8Array, - Int32Array, Int64Array, LargeBinaryArray, LargeListArray, LargeListViewArray, ListArray, - ListViewArray, StringArray, Time64MicrosecondArray, + Decimal64Array, Decimal128Array, Decimal256Array, FixedSizeBinaryArray, FixedSizeListArray, + Int8Array, Int32Array, Int64Array, LargeBinaryArray, LargeListArray, LargeListViewArray, + ListArray, ListViewArray, StringArray, Time64MicrosecondArray, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow_schema::{Field, Fields}; @@ -1592,6 +1609,93 @@ mod test { .build() } + #[test] + fn decimal256_inputs_are_narrowed_without_losing_values() { + use arrow::datatypes::i256; + + for (precision, scale, expected_type) in [ + (9, 2, DataType::Decimal32(9, 2)), + (18, 2, DataType::Decimal64(18, 2)), + (38, 2, DataType::Decimal128(38, 2)), + (38, 38, DataType::Decimal128(38, 38)), + ] { + let max = 10_i128.pow(precision as u32) - 1; + let typed_value = Decimal256Array::new( + vec![ + i256::ZERO, + i256::from_i128(max), + i256::from_i128(-max), + i256::MAX, // Arbitrary payload under a null must not be converted. + i256::ZERO, + ] + .into(), + Some(NullBuffer::from(vec![true, true, true, false, true])), + ) + .with_precision_and_scale(precision, scale) + .unwrap(); + let input = make_variant_struct_with_typed_value(Arc::new(typed_value)); + let input = StructArray::new( + input.fields().clone(), + input.columns().to_vec(), + Some(NullBuffer::from(vec![true, true, true, true, false])), + ) + .slice(1, 4); + let variant = VariantArray::try_new(&input).unwrap(); + assert_eq!( + variant.typed_value_column().unwrap().data_type(), + &expected_type + ); + assert_eq!(variant.nulls(), input.nulls()); + assert!(variant.typed_value_column().unwrap().is_null(2)); + assert_eq!(variant.value(2), Variant::Null); + for (i, raw) in [max, -max].into_iter().enumerate() { + let expected = match precision { + 9 => Variant::from(VariantDecimal4::try_new(raw as i32, scale as u8).unwrap()), + 18 => Variant::from(VariantDecimal8::try_new(raw as i64, scale as u8).unwrap()), + _ => Variant::from(VariantDecimal16::try_new(raw, scale as u8).unwrap()), + }; + assert_eq!(variant.value(i), expected); + } + let unshredded = crate::unshred_variant(&variant).unwrap(); + assert_eq!( + unshredded.iter().collect::>(), + variant.iter().collect::>() + ); + } + } + + #[test] + fn decimal256_inputs_reject_invalid_precision_scale_and_overflow() { + use arrow::datatypes::i256; + + for (precision, scale) in [(39, 2), (38, -1), (0, 0), (38, 39)] { + let data_type = DataType::Decimal256(precision, scale); + let input = make_variant_struct_with_typed_value(new_null_array(&data_type, 1)); + assert!(matches!( + VariantArray::try_new(&input), + Err(ArrowError::InvalidArgumentError(_)) + )); + } + + for (precision, raw) in [ + (9, i256::from_i128(1_000_000_000)), + (18, i256::from_i128(1_000_000_000_000_000_000)), + (38, i256::from_i128(10_i128.pow(38))), + (38, i256::from_i128(-10_i128.pow(38))), + (38, i256::MAX), + (38, i256::MIN), + ] { + let typed_value = Decimal256Array::from(vec![raw]) + .with_precision_and_scale(precision, 2) + .unwrap(); + let input = make_variant_struct_with_typed_value(Arc::new(typed_value)); + assert!( + VariantArray::try_new(&input).is_err(), + "accepted {raw} at precision {precision}" + ); + } + } + #[test] fn try_new_tags_untagged_uuid_on_read() { // Simulate a foreign writer that shredded a UUID column as bare FixedSizeBinary(16), diff --git a/parquet/tests/variant_integration.rs b/parquet/tests/variant_integration.rs index a80428b9617c..7b242248c48d 100644 --- a/parquet/tests/variant_integration.rs +++ b/parquet/tests/variant_integration.rs @@ -480,3 +480,86 @@ fn variant_is_valid(batch: &RecordBatch, filename: &str) -> bool { Variant::try_new(metadata.value(0), value.value(0)).is_ok() } + +#[test] +fn test_variant_wide_physical_decimal() { + use arrow::array::{Array, AsArray}; + use arrow::datatypes::{DataType, Decimal128Type}; + use bytes::Bytes; + use parquet::data_type::{ByteArray, ByteArrayType, FixedLenByteArrayType}; + use parquet::file::writer::SerializedFileWriter; + use parquet::schema::parser::parse_message_type; + use parquet_variant::{EMPTY_VARIANT_METADATA_BYTES, VariantDecimal16}; + use std::sync::Arc; + + for width in [17, 32] { + let schema = parse_message_type(&format!( + "message test {{ required group v (VARIANT) {{ + required binary metadata; + optional binary value; + optional fixed_len_byte_array({width}) typed_value (DECIMAL(38,2)); + }} }}" + )) + .unwrap(); + let mut bytes = Vec::new(); + let mut writer = + SerializedFileWriter::new(&mut bytes, Arc::new(schema), Default::default()).unwrap(); + let mut group = writer.next_row_group().unwrap(); + let mut column = group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch( + &vec![ByteArray::from(EMPTY_VARIANT_METADATA_BYTES); 3], + None, + None, + ) + .unwrap(); + column.close().unwrap(); + let mut column = group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&[], Some(&[0, 0, 0]), None) + .unwrap(); + column.close().unwrap(); + let values = [12345_i128, -12345].map(|raw| { + let mut bytes = vec![if raw < 0 { 0xff } else { 0 }; width]; + bytes[width - 16..].copy_from_slice(&raw.to_be_bytes()); + bytes.into() + }); + let mut column = group.next_column().unwrap().unwrap(); + column + .typed::() + .write_batch(&values, Some(&[1, 1, 0]), None) + .unwrap(); + column.close().unwrap(); + group.close().unwrap(); + writer.close().unwrap(); + + let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes)) + .unwrap() + .build() + .unwrap(); + let batch = reader.next().unwrap().unwrap(); + let input = batch.column(0).as_struct(); + assert_eq!( + input.column_by_name("typed_value").unwrap().data_type(), + &DataType::Decimal256(38, 2) + ); + let variant = VariantArray::try_new(input).unwrap(); + let decimals = variant + .typed_value_column() + .unwrap() + .as_primitive::(); + assert_eq!(decimals.value(0), 12345); + assert_eq!(decimals.value(1), -12345); + assert!(decimals.is_null(2)); + for (i, raw) in [12345, -12345].into_iter().enumerate() { + assert_eq!( + variant.value(i), + Variant::from(VariantDecimal16::try_new(raw, 2).unwrap()) + ); + } + assert_eq!(variant.value(2), Variant::Null); + assert!(reader.next().is_none()); + } +}