diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index 3ae23d3b7415c..60da10c28ffd9 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -2861,17 +2861,14 @@ impl ScalarValue { ($ARRAY_TY:ident, $SCALAR_TY:ident, $TZ:expr) => {{ { let array = scalars - .map(|sv| { - if let ScalarValue::$SCALAR_TY(v, _) = sv { - Ok(v) - } else { - _exec_err!( - "Inconsistent types in ScalarValue::iter_to_array. \ - Expected {:?}, got {:?}", - data_type, - sv - ) - } + .map(|sv| match sv { + ScalarValue::$SCALAR_TY(v, tz) if &tz == $TZ => Ok(v), + sv => _exec_err!( + "Inconsistent types in ScalarValue::iter_to_array. \ + Expected {:?}, got {:?}", + data_type, + sv + ), }) .collect::>()?; Arc::new(array.with_timezone_opt($TZ.clone())) @@ -3161,15 +3158,16 @@ impl ScalarValue { } DataType::FixedSizeBinary(size) => { let array = scalars - .map(|sv| { - if let ScalarValue::FixedSizeBinary(_, v) = sv { + .map(|sv| match sv { + ScalarValue::FixedSizeBinary(inner_size, v) + if inner_size == *size => + { Ok(v) - } else { - _exec_err!( - "Inconsistent types in ScalarValue::iter_to_array. \ - Expected {data_type}, got {sv:?}" - ) } + sv => _exec_err!( + "Inconsistent types in ScalarValue::iter_to_array. \ + Expected {data_type}, got {sv:?}" + ), }) .collect::>>()?; let array = FixedSizeBinaryArray::try_from_sparse_iter_with_size( @@ -3220,10 +3218,16 @@ impl ScalarValue { let array = scalars .into_iter() .map(|element: ScalarValue| match element { - ScalarValue::Decimal32(v1, _, _) => Ok(v1), - s => { - _internal_err!("Expected ScalarValue::Null element. Received {s:?}") + ScalarValue::Decimal32(value, inner_precision, inner_scale) + if inner_precision == precision && inner_scale == scale => + { + Ok(value) } + scalar => _exec_err!( + "Inconsistent types in ScalarValue::iter_to_array. Expected {:?}, got {:?}", + DataType::Decimal32(precision, scale), + scalar + ), }) .collect::>()? .with_precision_and_scale(precision, scale)?; @@ -3238,10 +3242,16 @@ impl ScalarValue { let array = scalars .into_iter() .map(|element: ScalarValue| match element { - ScalarValue::Decimal64(v1, _, _) => Ok(v1), - s => { - _internal_err!("Expected ScalarValue::Null element. Received {s:?}") + ScalarValue::Decimal64(value, inner_precision, inner_scale) + if inner_precision == precision && inner_scale == scale => + { + Ok(value) } + scalar => _exec_err!( + "Inconsistent types in ScalarValue::iter_to_array. Expected {:?}, got {:?}", + DataType::Decimal64(precision, scale), + scalar + ), }) .collect::>()? .with_precision_and_scale(precision, scale)?; @@ -3256,10 +3266,16 @@ impl ScalarValue { let array = scalars .into_iter() .map(|element: ScalarValue| match element { - ScalarValue::Decimal128(v1, _, _) => Ok(v1), - s => { - _internal_err!("Expected ScalarValue::Null element. Received {s:?}") + ScalarValue::Decimal128(value, inner_precision, inner_scale) + if inner_precision == precision && inner_scale == scale => + { + Ok(value) } + scalar => _exec_err!( + "Inconsistent types in ScalarValue::iter_to_array. Expected {:?}, got {:?}", + DataType::Decimal128(precision, scale), + scalar + ), }) .collect::>()? .with_precision_and_scale(precision, scale)?; @@ -3274,12 +3290,16 @@ impl ScalarValue { let array = scalars .into_iter() .map(|element: ScalarValue| match element { - ScalarValue::Decimal256(v1, _, _) => Ok(v1), - s => { - _internal_err!( - "Expected ScalarValue::Decimal256 element. Received {s:?}" - ) + ScalarValue::Decimal256(value, inner_precision, inner_scale) + if inner_precision == precision && inner_scale == scale => + { + Ok(value) } + scalar => _exec_err!( + "Inconsistent types in ScalarValue::iter_to_array. Expected {:?}, got {:?}", + DataType::Decimal256(precision, scale), + scalar + ), }) .collect::>()? .with_precision_and_scale(precision, scale)?; diff --git a/datafusion/core/tests/fuzz_cases/record_batch_generator.rs b/datafusion/core/tests/fuzz_cases/record_batch_generator.rs index 22b145f5095a7..12d1e28346d42 100644 --- a/datafusion/core/tests/fuzz_cases/record_batch_generator.rs +++ b/datafusion/core/tests/fuzz_cases/record_batch_generator.rs @@ -61,10 +61,13 @@ pub fn get_supported_types_columns(rng_seed: u64) -> Vec { ColumnDescr::new("time32_ms", DataType::Time32(TimeUnit::Millisecond)), ColumnDescr::new("time64_us", DataType::Time64(TimeUnit::Microsecond)), ColumnDescr::new("time64_ns", DataType::Time64(TimeUnit::Nanosecond)), - ColumnDescr::new("timestamp_s", DataType::Timestamp(TimeUnit::Second, None)), + ColumnDescr::new( + "timestamp_s", + DataType::Timestamp(TimeUnit::Second, Some(Arc::from("UTC"))), + ), ColumnDescr::new( "timestamp_ms", - DataType::Timestamp(TimeUnit::Millisecond, None), + DataType::Timestamp(TimeUnit::Millisecond, Some(Arc::from("+05:30"))), ), ColumnDescr::new( "timestamp_us", @@ -142,6 +145,9 @@ pub fn get_supported_types_columns(rng_seed: u64) -> Vec { ColumnDescr::new("binary", DataType::Binary), ColumnDescr::new("large_binary", DataType::LargeBinary), ColumnDescr::new("binaryview", DataType::BinaryView), + ColumnDescr::new("fixed_binary_1", DataType::FixedSizeBinary(1)), + ColumnDescr::new("fixed_binary_8", DataType::FixedSizeBinary(8)), + ColumnDescr::new("fixed_binary_32", DataType::FixedSizeBinary(32)), ColumnDescr::new( "dictionary_utf8_low", DataType::Dictionary(Box::new(DataType::UInt64), Box::new(DataType::Utf8)), @@ -249,6 +255,27 @@ macro_rules! generate_primitive_array { }}; } +macro_rules! generate_timestamp_array { + ($SELF:ident, $NUM_ROWS:ident, $MAX_NUM_DISTINCT:expr, $NULL_PCT:ident, $BATCH_GEN_RNG:ident, $ARRAY_GEN_RNG:ident, $ARROW_TYPE:ident, $TIMEZONE:ident) => {{ + let array = generate_primitive_array!( + $SELF, + $NUM_ROWS, + $MAX_NUM_DISTINCT, + $NULL_PCT, + $BATCH_GEN_RNG, + $ARRAY_GEN_RNG, + $ARROW_TYPE + ); + let array = array + .as_any() + .downcast_ref::>() + .unwrap() + .clone() + .with_timezone_opt($TIMEZONE.clone()); + Arc::new(array) as ArrayRef + }}; +} + macro_rules! generate_dict { ($SELF:ident, $NUM_ROWS:ident, $MAX_NUM_DISTINCT:expr, $NULL_PCT:ident, $BATCH_GEN_RNG:ident, $ARRAY_GEN_RNG:ident, $ARROW_TYPE:ident, $VALUES: ident) => {{ debug_assert_eq!($VALUES.len(), $MAX_NUM_DISTINCT); @@ -617,48 +644,52 @@ impl RecordBatchGenerator { DurationNanosecondType ) } - DataType::Timestamp(TimeUnit::Second, None) => { - generate_primitive_array!( + DataType::Timestamp(TimeUnit::Second, ref timezone) => { + generate_timestamp_array!( self, num_rows, max_num_distinct, null_pct, batch_gen_rng, array_gen_rng, - TimestampSecondType + TimestampSecondType, + timezone ) } - DataType::Timestamp(TimeUnit::Millisecond, None) => { - generate_primitive_array!( + DataType::Timestamp(TimeUnit::Millisecond, ref timezone) => { + generate_timestamp_array!( self, num_rows, max_num_distinct, null_pct, batch_gen_rng, array_gen_rng, - TimestampMillisecondType + TimestampMillisecondType, + timezone ) } - DataType::Timestamp(TimeUnit::Microsecond, None) => { - generate_primitive_array!( + DataType::Timestamp(TimeUnit::Microsecond, ref timezone) => { + generate_timestamp_array!( self, num_rows, max_num_distinct, null_pct, batch_gen_rng, array_gen_rng, - TimestampMicrosecondType + TimestampMicrosecondType, + timezone ) } - DataType::Timestamp(TimeUnit::Nanosecond, None) => { - generate_primitive_array!( + DataType::Timestamp(TimeUnit::Nanosecond, ref timezone) => { + generate_timestamp_array!( self, num_rows, max_num_distinct, null_pct, batch_gen_rng, array_gen_rng, - TimestampNanosecondType + TimestampNanosecondType, + timezone ) } DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { @@ -697,6 +728,17 @@ impl RecordBatchGenerator { _ => unreachable!(), } } + DataType::FixedSizeBinary(width) => { + let mut generator = BinaryArrayGenerator { + max_len: usize::try_from(width) + .expect("fixed-size binary width must be nonnegative"), + num_binaries: num_rows, + num_distinct_binaries: max_num_distinct, + null_pct, + rng: array_gen_rng, + }; + generator.gen_fixed_size_binary() + } DataType::Decimal32(precision, scale) => { generate_decimal_array!( self, diff --git a/datafusion/core/tests/fuzz_cases/scalar_value_fuzz.rs b/datafusion/core/tests/fuzz_cases/scalar_value_fuzz.rs index 040a86ea3778c..3ac6d89e1d50c 100644 --- a/datafusion/core/tests/fuzz_cases/scalar_value_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/scalar_value_fuzz.rs @@ -16,6 +16,10 @@ // under the License. use arrow::array::ArrayRef; +use arrow_schema::{ + DECIMAL32_MAX_PRECISION, DECIMAL64_MAX_PRECISION, DECIMAL128_MAX_PRECISION, + DECIMAL256_MAX_PRECISION, +}; use datafusion_common::ScalarValue; use rand::random; @@ -63,6 +67,40 @@ fn scalar_value_iter_to_array_roundtrip() { }); } +#[test] +fn scalar_value_iter_to_array_rejects_mixed_parameterized_types() { + let mut mutation_counts = [0; MutationKind::COUNT]; + + for_each_array(|context, array| { + if array.len() < 2 { + return; + } + + let scalars = extract_scalars(array, context); + for (kind, mutated) in type_parameter_mutations(&scalars[1]) { + assert_ne!(mutated.data_type(), scalars[0].data_type()); + let mut mutated_scalars = scalars.clone(); + mutated_scalars[1] = mutated; + + let error = ScalarValue::iter_to_array(mutated_scalars).unwrap_err(); + assert!( + error + .to_string() + .contains("Inconsistent types in ScalarValue::iter_to_array"), + "iter_to_array returned an unexpected error for {kind:?} mutation of {context}: {error}" + ); + mutation_counts[kind as usize] += 1; + } + }); + + for kind in MutationKind::ALL { + assert!( + mutation_counts[kind as usize] > 0, + "no {kind:?} mutations were tested" + ); + } +} + #[test] fn scalar_value_to_array_roundtrip() { for_each_array(|context, array| { @@ -89,7 +127,125 @@ fn scalar_value_to_array_roundtrip() { }); } -fn for_each_array(check: impl Fn(&str, &ArrayRef)) { +#[derive(Debug, Clone, Copy)] +enum MutationKind { + DecimalPrecision, + DecimalScale, + TimestampTimezone, + FixedSizeBinaryWidth, +} + +impl MutationKind { + const ALL: [Self; 4] = [ + Self::DecimalPrecision, + Self::DecimalScale, + Self::TimestampTimezone, + Self::FixedSizeBinaryWidth, + ]; + const COUNT: usize = Self::ALL.len(); +} + +fn type_parameter_mutations(scalar: &ScalarValue) -> Vec<(MutationKind, ScalarValue)> { + use ScalarValue::*; + + macro_rules! decimal_mutations { + ($CONSTRUCTOR:ident, $VALUE:ident, $PRECISION:ident, $SCALE:ident, $MAX:expr) => {{ + let mut mutations = vec![( + MutationKind::DecimalScale, + $CONSTRUCTOR(*$VALUE, *$PRECISION, different_scale(*$SCALE)), + )]; + if let Some(precision) = different_precision(*$PRECISION, *$SCALE, $MAX) { + mutations.push(( + MutationKind::DecimalPrecision, + $CONSTRUCTOR(*$VALUE, precision, *$SCALE), + )); + } + mutations + }}; + } + + match scalar { + Decimal32(value, precision, scale) => decimal_mutations!( + Decimal32, + value, + precision, + scale, + DECIMAL32_MAX_PRECISION + ), + Decimal64(value, precision, scale) => decimal_mutations!( + Decimal64, + value, + precision, + scale, + DECIMAL64_MAX_PRECISION + ), + Decimal128(value, precision, scale) => decimal_mutations!( + Decimal128, + value, + precision, + scale, + DECIMAL128_MAX_PRECISION + ), + Decimal256(value, precision, scale) => decimal_mutations!( + Decimal256, + value, + precision, + scale, + DECIMAL256_MAX_PRECISION + ), + TimestampSecond(value, timezone) => vec![( + MutationKind::TimestampTimezone, + TimestampSecond(*value, toggled_timezone(timezone)), + )], + TimestampMillisecond(value, timezone) => vec![( + MutationKind::TimestampTimezone, + TimestampMillisecond(*value, toggled_timezone(timezone)), + )], + TimestampMicrosecond(value, timezone) => vec![( + MutationKind::TimestampTimezone, + TimestampMicrosecond(*value, toggled_timezone(timezone)), + )], + TimestampNanosecond(value, timezone) => vec![( + MutationKind::TimestampTimezone, + TimestampNanosecond(*value, toggled_timezone(timezone)), + )], + FixedSizeBinary(width, _) => vec![( + MutationKind::FixedSizeBinaryWidth, + FixedSizeBinary(width + 1, None), + )], + _ => vec![], + } +} + +fn different_precision(precision: u8, scale: i8, max_precision: u8) -> Option { + if precision < max_precision { + Some(precision + 1) + } else if precision > 1 && scale <= (precision - 1) as i8 { + Some(precision - 1) + } else { + None + } +} + +fn different_scale(scale: i8) -> i8 { + if scale == i8::MIN { + scale + 1 + } else { + scale - 1 + } +} + +fn toggled_timezone( + timezone: &Option>, +) -> Option> { + if timezone.is_some() { + None + } else { + Some(std::sync::Arc::from("UTC")) + } +} + +fn for_each_array(mut check: impl FnMut(&str, &ArrayRef)) { for _ in 0..NUM_SEEDS { let seed = random(); let columns = get_supported_types_columns(seed); @@ -99,7 +255,7 @@ fn for_each_array(check: impl Fn(&str, &ArrayRef)) { .unwrap_or_else(|e| panic!("failed to generate batch for seed {seed}: {e}")); for (field, array) in batch.schema().fields().iter().zip(batch.columns()) { - let check_array = |array: &ArrayRef| { + let mut check_array = |array: &ArrayRef| { let context = format!( "seed={seed}, column={}, type={}, len={}, offset={}", field.name(), diff --git a/test-utils/src/array_gen/binary.rs b/test-utils/src/array_gen/binary.rs index ab0530a9ab4e4..90019374a6ee2 100644 --- a/test-utils/src/array_gen/binary.rs +++ b/test-utils/src/array_gen/binary.rs @@ -16,7 +16,8 @@ // under the License. use arrow::array::{ - ArrayRef, BinaryViewArray, GenericBinaryArray, OffsetSizeTrait, UInt32Array, + ArrayRef, BinaryViewArray, FixedSizeBinaryArray, GenericBinaryArray, OffsetSizeTrait, + UInt32Array, }; use arrow::compute; use rand::Rng; @@ -43,20 +44,20 @@ impl BinaryArrayGenerator { .map(|_| Some(random_binary(&mut self.rng, self.max_len))) .collect(); - // Pick num_binaries randomly from the distinct binary table - let indices: UInt32Array = (0..self.num_binaries) - .map(|_| { - if self.rng.random::() < self.null_pct { - None - } else if self.num_distinct_binaries > 1 { - let range = 0..(self.num_distinct_binaries as u32); - Some(self.rng.random_range(range)) - } else { - Some(0) - } - }) - .collect(); + let indices = self.gen_indices(); + compute::take(&distinct_binaries, &indices, None).unwrap() + } + /// Creates a FixedSizeBinaryArray with random binary data. + pub fn gen_fixed_size_binary(&mut self) -> ArrayRef { + let width = self.max_len; + let distinct_binaries = FixedSizeBinaryArray::try_from_iter( + (0..self.num_distinct_binaries) + .map(|_| (0..width).map(|_| self.rng.random()).collect::>()), + ) + .unwrap(); + + let indices = self.gen_indices(); compute::take(&distinct_binaries, &indices, None).unwrap() } @@ -66,7 +67,13 @@ impl BinaryArrayGenerator { .map(|_| Some(random_binary(&mut self.rng, self.max_len))) .collect(); - let indices: UInt32Array = (0..self.num_binaries) + let indices = self.gen_indices(); + compute::take(&distinct_binary_views, &indices, None).unwrap() + } + + /// Generates nullable indices into the distinct binary values. + fn gen_indices(&mut self) -> UInt32Array { + (0..self.num_binaries) .map(|_| { if self.rng.random::() < self.null_pct { None @@ -77,9 +84,7 @@ impl BinaryArrayGenerator { Some(0) } }) - .collect(); - - compute::take(&distinct_binary_views, &indices, None).unwrap() + .collect() } }