diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index 924620a930869..60da66383a556 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -3309,7 +3309,8 @@ impl ScalarValue { let values = if values.is_empty() { new_empty_array(data_type) } else { - Self::iter_to_array(values.iter().cloned()).unwrap() + let arr = Self::iter_to_array(values.iter().cloned()).unwrap(); + cast_with_options(&arr, data_type, &DEFAULT_CAST_OPTIONS).unwrap() }; Arc::new( SingleRowListArrayBuilder::new(values) @@ -3371,7 +3372,8 @@ impl ScalarValue { let values = if values.len() == 0 { new_empty_array(data_type) } else { - Self::iter_to_array(values).unwrap() + let arr = Self::iter_to_array(values).unwrap(); + cast_with_options(&arr, data_type, &DEFAULT_CAST_OPTIONS).unwrap() }; Arc::new( SingleRowListArrayBuilder::new(values) @@ -3414,7 +3416,8 @@ impl ScalarValue { let values = if values.is_empty() { new_empty_array(data_type) } else { - Self::iter_to_array(values.iter().cloned()).unwrap() + let arr = Self::iter_to_array(values.iter().cloned()).unwrap(); + cast_with_options(&arr, data_type, &DEFAULT_CAST_OPTIONS).unwrap() }; Arc::new(SingleRowListArrayBuilder::new(values).build_large_list_array()) } @@ -11544,4 +11547,28 @@ mod tests { run_tests::(); run_tests::(); } + + #[test] + fn test_new_list_nested_nullability_mismatch_issue_24022() { + // requested element type: Struct(n: Int32 nullable=true) + let requested_element_type = + DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int32, true)])); + + // inferred from concrete values: Struct(n: Int32 nullable=false) + let inferred_field = Field::new("n", DataType::Int32, false); + + let value = ScalarValue::Struct(Arc::new(StructArray::from(vec![( + Arc::new(inferred_field), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]))); + + let list = ScalarValue::new_list(&[value], &requested_element_type, true); + assert_eq!( + list.data_type(), + &DataType::List(Arc::new(Field::new_list_field( + requested_element_type, + true + ))) + ); + } } diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index cfacd771968c2..e15ddcd825d2f 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -1741,15 +1741,17 @@ mod tests { acc2.update_batch(&[data(["b", "c", "a"])])?; acc1 = merge(acc1, acc2)?; - assert_eq!(acc1.size(), 282); + assert_eq!(acc1.size(), 166); Ok(()) } #[test] fn does_not_over_account_memory_distinct() -> Result<()> { - let (mut acc1, mut acc2) = ArrayAggAccumulatorBuilder::string() - .distinct() - .build_two()?; + let (mut acc1, mut acc2) = ArrayAggAccumulatorBuilder::new(DataType::List( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + )) + .distinct() + .build_two()?; acc1.update_batch(&[string_list_data([ vec!["a", "b", "c"], @@ -1765,9 +1767,11 @@ mod tests { #[test] fn does_not_over_account_memory_ordered() -> Result<()> { - let mut acc = ArrayAggAccumulatorBuilder::string() - .order_by_col("col", SortOptions::new(false, false)) - .build()?; + let mut acc = ArrayAggAccumulatorBuilder::new(DataType::List(Arc::new( + Field::new_list_field(DataType::Utf8, true), + ))) + .order_by_col("col", SortOptions::new(false, false)) + .build()?; acc.update_batch(&[string_list_data([ vec!["a", "b", "c"], @@ -1781,6 +1785,122 @@ mod tests { Ok(()) } + #[test] + fn ordered_aggregate_nested_nullability_mismatch_issue_24022() -> Result<()> { + use arrow::array::{Int32Array, Int64Array, StructArray}; + use datafusion_physical_expr::expressions::Column; + + let requested_element_type = + DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int32, true)])); + let inferred_field = Field::new("n", DataType::Int32, false); + + let ordering_dtype = DataType::Int64; + let schema = Schema::new(vec![ + Field::new("val", requested_element_type.clone(), true), + Field::new("ord", DataType::Int64, true), + ]); + let ord_expr = Arc::new( + Column::new_with_schema("ord", &schema).expect("column not in schema"), + ) as Arc; + + let asc_opts = SortOptions { + descending: false, + nulls_first: false, + }; + let asc_ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::clone(&ord_expr), + asc_opts, + )]) + .unwrap(); + + let mut acc = OrderSensitiveArrayAggAccumulator::try_new( + &requested_element_type, + std::slice::from_ref(&ordering_dtype), + asc_ordering, + /*is_input_pre_ordered=*/ true, + /*reverse=*/ false, + /*ignore_nulls=*/ false, + )?; + + let value_arr = Arc::new(StructArray::from(vec![( + Arc::new(inferred_field), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )])) as ArrayRef; + + let ord_arr = Arc::new(Int64Array::from(vec![0i64])) as ArrayRef; + + acc.update_batch(&[value_arr, ord_arr])?; + + let evaluated = acc.evaluate()?; + + if let ScalarValue::List(arr) = evaluated { + assert_eq!( + arr.data_type(), + &DataType::List(Arc::new(Field::new_list_field( + requested_element_type.clone(), + true + ))) + ); + + let expected_struct_array = StructArray::from(vec![( + Arc::new(Field::new("n", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]); + let expected_array = Arc::new(expected_struct_array) as ArrayRef; + assert_eq!(&arr.value(0), &expected_array); + } else { + panic!("Expected ScalarValue::List"); + } + + Ok(()) + } + + #[test] + fn distinct_aggregate_nested_nullability_mismatch_issue_24022() -> Result<()> { + use arrow::array::{Int32Array, StructArray}; + use datafusion_common::ScalarValue; + + let requested_element_type = + DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int32, true)])); + let inferred_field = Field::new("n", DataType::Int32, false); + + let mut acc = DistinctArrayAggAccumulator::try_new( + &requested_element_type, + None, + /*ignore_nulls=*/ false, + )?; + + let value_arr = Arc::new(StructArray::from(vec![( + Arc::new(inferred_field), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )])) as ArrayRef; + + acc.update_batch(&[value_arr])?; + + let evaluated = acc.evaluate()?; + + if let ScalarValue::List(arr) = evaluated { + assert_eq!( + arr.data_type(), + &DataType::List(Arc::new(Field::new_list_field( + requested_element_type.clone(), + true + ))) + ); + + let expected_struct_array = StructArray::from(vec![( + Arc::new(Field::new("n", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]); + let expected_array = Arc::new(expected_struct_array) as ArrayRef; + assert_eq!(&arr.value(0), &expected_array); + } else { + panic!("Expected ScalarValue::List"); + } + + Ok(()) + } + // Reproduces the bug where `state()` emits reversed values but non-reversed // orderings when the optimizer sets is_input_pre_ordered=true + reverse=true // (DESC aggregate with ASC pre-sorted input). The partial states are fed into @@ -1904,15 +2024,19 @@ mod tests { fn new(data_type: DataType) -> Self { Self { - return_field: Field::new("f", data_type.clone(), true).into(), + return_field: Field::new( + "f", + DataType::List(Arc::new(Field::new_list_field( + data_type.clone(), + true, + ))), + true, + ) + .into(), distinct: false, order_bys: vec![], schema: Schema { - fields: Fields::from(vec![Field::new( - "col", - DataType::new_list(data_type, true), - true, - )]), + fields: Fields::from(vec![Field::new("col", data_type, true)]), metadata: Default::default(), }, }