From e3f28e6153e03b807ba636a24b75f417cd3baaa0 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Thu, 27 Aug 2026 23:18:23 -0700 Subject: [PATCH 01/13] fix: align Spark collect aggregate element nullability --- .../spark/src/function/aggregate/collect.rs | 244 +++++++++++++++++- 1 file changed, 230 insertions(+), 14 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 310bc1c890657..58d4e3edc5340 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::ArrayRef; +use arrow::array::{Array, ArrayRef}; +use arrow::compute::cast; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::utils::SingleRowListArrayBuilder; use datafusion_common::{Result, ScalarValue, internal_err}; @@ -43,7 +44,58 @@ fn empty_list_scalar(list_type: &DataType) -> Result { ); }; let empty = arrow::array::new_empty_array(field.data_type()); - Ok(SingleRowListArrayBuilder::new(empty).build_list_scalar()) + Ok(SingleRowListArrayBuilder::new(empty) + .with_field(field) + .build_list_scalar()) +} + +fn collect_type(element_type: DataType) -> DataType { + DataType::List(Arc::new(Field::new_list_field(element_type, false))) +} + +/// Rebuild an accumulator result with the aggregate's declared list field. +/// +/// The shared array aggregate accumulators use a nullable list field and can +/// derive nested fields from runtime arrays. Spark collect aggregates always +/// drop null inputs, so their element field is non-nullable. Reusing the +/// declared field also keeps nested types consistent with planning. +fn normalize_list_scalar( + value: ScalarValue, + list_type: &DataType, +) -> Result { + let DataType::List(field) = list_type else { + return internal_err!( + "collect_list/collect_set expected List return type, got {list_type:?}" + ); + }; + let ScalarValue::List(array) = value else { + return internal_err!( + "collect_list/collect_set accumulator returned a non-List value" + ); + }; + if array.len() != 1 { + return internal_err!( + "collect_list/collect_set accumulator returned {} rows, expected one", + array.len() + ); + } + if array.is_null(0) { + return Ok(ScalarValue::new_null_list( + field.data_type().clone(), + field.is_nullable(), + 1, + )); + } + + let values = array.value(0); + let values = if values.data_type() == field.data_type() { + values + } else { + cast(values.as_ref(), field.data_type())? + }; + Ok(SingleRowListArrayBuilder::new(values) + .with_field(field) + .build_list_scalar()) } // @@ -76,17 +128,14 @@ impl AggregateUDFImpl for SparkCollectList { } fn return_type(&self, arg_types: &[DataType]) -> Result { - Ok(DataType::List(Arc::new(Field::new_list_field( - arg_types[0].clone(), - true, - )))) + Ok(collect_type(arg_types[0].clone())) } fn state_fields(&self, args: StateFieldsArgs) -> Result> { Ok(vec![ Field::new_list( format_state_name(args.name, "collect_list"), - Field::new_list_field(args.input_fields[0].data_type().clone(), true), + Field::new_list_field(args.input_fields[0].data_type().clone(), false), true, ) .into(), @@ -137,17 +186,14 @@ impl AggregateUDFImpl for SparkCollectSet { } fn return_type(&self, arg_types: &[DataType]) -> Result { - Ok(DataType::List(Arc::new(Field::new_list_field( - arg_types[0].clone(), - true, - )))) + Ok(collect_type(arg_types[0].clone())) } fn state_fields(&self, args: StateFieldsArgs) -> Result> { Ok(vec![ Field::new_list( format_state_name(args.name, "collect_set"), - Field::new_list_field(args.input_fields[0].data_type().clone(), true), + Field::new_list_field(args.input_fields[0].data_type().clone(), false), true, ) .into(), @@ -192,7 +238,11 @@ impl Accumulator for NullToEmptyListAccumulator { } fn state(&mut self) -> Result> { - self.inner.state() + self.inner + .state()? + .into_iter() + .map(|value| normalize_list_scalar(value, &self.list_type)) + .collect() } fn evaluate(&mut self) -> Result { @@ -200,7 +250,7 @@ impl Accumulator for NullToEmptyListAccumulator { if result.is_null() { empty_list_scalar(&self.list_type) } else { - Ok(result) + normalize_list_scalar(result, &self.list_type) } } @@ -216,3 +266,169 @@ impl Accumulator for NullToEmptyListAccumulator { self.inner.size() + self.list_type.size() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, StructArray}; + use arrow::datatypes::Fields; + + fn list_type(element_type: DataType) -> DataType { + DataType::List(Arc::new(Field::new_list_field(element_type, false))) + } + + fn accumulator( + element_type: &DataType, + distinct: bool, + ) -> Result> { + let return_type = list_type(element_type.clone()); + if distinct { + Ok(Box::new(NullToEmptyListAccumulator::new( + DistinctArrayAggAccumulator::try_new(element_type, None, true)?, + return_type, + ))) + } else { + Ok(Box::new(NullToEmptyListAccumulator::new( + ArrayAggAccumulator::try_new(element_type, true)?, + return_type, + ))) + } + } + + fn assert_empty_list(value: &ScalarValue) { + let ScalarValue::List(array) = value else { + panic!("expected a list scalar") + }; + assert_eq!(array.value(0).len(), 0); + } + + fn assert_nested_values(value: &ScalarValue) { + let ScalarValue::List(array) = value else { + panic!("expected a list scalar") + }; + let values = array.value(0); + let structs = values + .as_any() + .downcast_ref::() + .expect("expected struct values"); + let integers = structs + .column(0) + .as_any() + .downcast_ref::() + .expect("expected Int32 struct field"); + let mut actual: Vec = integers.iter().map(Option::unwrap).collect(); + actual.sort_unstable(); + assert_eq!(actual, vec![1, 2]); + } + + #[test] + fn collect_types_have_non_nullable_elements() -> Result<()> { + let element_type = DataType::Int32; + let expected = list_type(element_type.clone()); + + for aggregate in [ + &SparkCollectList::new() as &dyn AggregateUDFImpl, + &SparkCollectSet::new() as &dyn AggregateUDFImpl, + ] { + assert_eq!( + aggregate.return_type(std::slice::from_ref(&element_type))?, + expected + ); + + let input_field = Arc::new(Field::new("input", element_type.clone(), true)); + let state_fields = aggregate.state_fields(StateFieldsArgs { + name: aggregate.name(), + input_fields: &[input_field], + return_field: Arc::new(Field::new("result", expected.clone(), false)), + ordering_fields: &[], + is_distinct: false, + })?; + assert_eq!(state_fields[0].data_type(), &expected); + } + + Ok(()) + } + + #[test] + fn empty_results_have_non_nullable_elements() -> Result<()> { + let expected = list_type(DataType::Int32); + + for aggregate in [ + &SparkCollectList::new() as &dyn AggregateUDFImpl, + &SparkCollectSet::new() as &dyn AggregateUDFImpl, + ] { + let value = aggregate.default_value(&expected)?; + assert!(!value.is_null()); + assert_eq!(value.data_type(), expected); + assert_empty_list(&value); + } + + for distinct in [false, true] { + let value = accumulator(&DataType::Int32, distinct)?.evaluate()?; + assert!(!value.is_null()); + assert_eq!(value.data_type(), expected); + assert_empty_list(&value); + } + + Ok(()) + } + + #[test] + fn accumulator_state_and_output_preserve_nested_type() -> Result<()> { + let declared_fields = + Fields::from(vec![Field::new("required", DataType::Int32, false)]); + let element_type = DataType::Struct(declared_fields.clone()); + let expected = list_type(element_type.clone()); + + // Exercise the downstream case from the issue: runtime arrays can carry + // different nested nullability than the aggregate's declared type. + let runtime_fields = + Fields::from(vec![Field::new("required", DataType::Int32, true)]); + let values = Arc::new(StructArray::new( + runtime_fields, + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2)]))], + None, + )) as ArrayRef; + + for distinct in [false, true] { + let mut partial = accumulator(&element_type, distinct)?; + partial.update_batch(std::slice::from_ref(&values))?; + + let state = partial.state()?; + assert_eq!(state[0].data_type(), expected); + assert_nested_values(&state[0]); + + let value = partial.evaluate()?; + assert_eq!(value.data_type(), expected); + assert_nested_values(&value); + + let mut final_accumulator = accumulator(&element_type, distinct)?; + final_accumulator.merge_batch(&[state[0].to_array()?])?; + let merged = final_accumulator.evaluate()?; + assert_eq!(merged.data_type(), expected); + assert_nested_values(&merged); + } + + Ok(()) + } + + #[test] + fn normalization_reuses_matching_primitive_values() -> Result<()> { + let values = Arc::new(Int32Array::from(vec![1, 2, 3])); + let values_ptr = values.values().as_ptr(); + let scalar = SingleRowListArrayBuilder::new(values).build_list_scalar(); + + let normalized = normalize_list_scalar(scalar, &list_type(DataType::Int32))?; + let ScalarValue::List(array) = normalized else { + panic!("expected a list scalar") + }; + let normalized_values = array.value(0); + let normalized_values = normalized_values + .as_any() + .downcast_ref::() + .expect("expected Int32 values"); + + assert_eq!(normalized_values.values().as_ptr(), values_ptr); + Ok(()) + } +} From ddbadd312cee1dd117c4324d6e10305993fbd0dd Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 29 Aug 2026 19:15:24 -0700 Subject: [PATCH 02/13] fix: handle encoded values in non-nullable list builder --- datafusion/common/src/utils/mod.rs | 64 +++- .../spark/src/function/aggregate/collect.rs | 285 +++++++++++++++++- .../test_files/spark/aggregate/collect.slt | 12 + 3 files changed, 350 insertions(+), 11 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index e047db39a5740..4b3ef0580f247 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -30,8 +30,8 @@ use crate::error::{ }; use crate::{Result, ScalarValue}; use arrow::array::{ - Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait, - cast::AsArray, + Array, ArrayData, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, + OffsetSizeTrait, cast::AsArray, }; use arrow::array::{ ArrowPrimitiveType, BooleanArray, Datum, GenericListArray, Int32Array, Int64Array, @@ -591,9 +591,7 @@ impl SingleRowListArrayBuilder { /// Build a single element [`ListArray`] pub fn build_list_array(self) -> ListArray { - let (field, arr) = self.into_field_and_arr(); - let offsets = OffsetBuffer::from_lengths([arr.len()]); - ListArray::new(field, offsets, arr, None) + self.build_generic_list_array() } /// Build a single element [`ListArray`] and wrap as [`ScalarValue::List`] @@ -603,9 +601,31 @@ impl SingleRowListArrayBuilder { /// Build a single element [`LargeListArray`] pub fn build_large_list_array(self) -> LargeListArray { + self.build_generic_list_array() + } + + fn build_generic_list_array( + self, + ) -> GenericListArray { let (field, arr) = self.into_field_and_arr(); let offsets = OffsetBuffer::from_lengths([arr.len()]); - LargeListArray::new(field, offsets, arr, None) + + // `is_nullable` is conservative for encoded arrays and can be true + // even when the array contains no logical nulls. In that case the + // generic constructor rejects a valid non-nullable list child. + if !field.is_nullable() && arr.is_nullable() && arr.logical_null_count() == 0 { + let data = ArrayData::builder( + GenericListArray::::DATA_TYPE_CONSTRUCTOR(field), + ) + .len(1) + .add_buffer(offsets.into_inner().into_inner()) + .add_child_data(arr.to_data()) + .build() + .expect("single-row list array should contain valid data"); + return GenericListArray::from(data); + } + + GenericListArray::new(field, offsets, arr, None) } /// Build a single element [`LargeListArray`] and wrap as [`ScalarValue::LargeList`] @@ -1503,9 +1523,9 @@ mod tests { use super::*; use crate::ScalarValue::Null; use arrow::{ - array::{Float64Array, Int32Array}, + array::{DictionaryArray, Float64Array, Int8Array, Int32Array, StringArray}, buffer::NullBuffer, - datatypes::Int32Type, + datatypes::{Int8Type, Int32Type}, }; #[cfg(feature = "sql")] use sqlparser::ast::Ident; @@ -1534,6 +1554,34 @@ mod tests { } } + #[test] + fn single_row_list_builder_handles_conservatively_nullable_values() { + let keys = Int8Array::from(vec![0, 0]); + let dictionary_values = Arc::new(StringArray::from(vec![Some("a"), None])); + let values = Arc::new( + DictionaryArray::::try_new(keys, dictionary_values).unwrap(), + ) as ArrayRef; + + assert_eq!(values.logical_null_count(), 0); + assert!(values.is_nullable()); + + let list = SingleRowListArrayBuilder::new(Arc::clone(&values)) + .with_nullable(false) + .build_list_array(); + list.to_data().validate_full().unwrap(); + + let value = ScalarValue::try_from_array(&list, 0).unwrap(); + assert_eq!(value.data_type(), list.data_type().clone()); + + let list = SingleRowListArrayBuilder::new(values) + .with_nullable(false) + .build_large_list_array(); + list.to_data().validate_full().unwrap(); + + let value = ScalarValue::try_from_array(&list, 0).unwrap(); + assert_eq!(value.data_type(), list.data_type().clone()); + } + #[test] fn test_bisect_linear_left_and_right() -> Result<()> { let arrays: Vec = vec![ diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 58d4e3edc5340..3f3bf3e38741e 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -131,6 +131,10 @@ impl AggregateUDFImpl for SparkCollectList { Ok(collect_type(arg_types[0].clone())) } + fn is_nullable(&self) -> bool { + false + } + fn state_fields(&self, args: StateFieldsArgs) -> Result> { Ok(vec![ Field::new_list( @@ -189,6 +193,10 @@ impl AggregateUDFImpl for SparkCollectSet { Ok(collect_type(arg_types[0].clone())) } + fn is_nullable(&self) -> bool { + false + } + fn state_fields(&self, args: StateFieldsArgs) -> Result> { Ok(vec![ Field::new_list( @@ -226,11 +234,29 @@ impl NullToEmptyListAccumulator { pub fn new(inner: T, list_type: DataType) -> Self { Self { inner, list_type } } + + fn normalize_input(&self, value: &ArrayRef) -> Result { + let DataType::List(field) = &self.list_type else { + return internal_err!( + "collect_list/collect_set expected List return type, got {:?}", + self.list_type + ); + }; + if value.data_type() == field.data_type() { + Ok(Arc::clone(value)) + } else { + Ok(cast(value.as_ref(), field.data_type())?) + } + } } impl Accumulator for NullToEmptyListAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.inner.update_batch(values) + let [value] = values else { + return self.inner.update_batch(values); + }; + let value = self.normalize_input(value)?; + self.inner.update_batch(&[value]) } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { @@ -270,8 +296,16 @@ impl Accumulator for NullToEmptyListAccumulator { #[cfg(test)] mod tests { use super::*; - use arrow::array::{Int32Array, StructArray}; - use arrow::datatypes::Fields; + use arrow::array::{ + DictionaryArray, Int8Array, Int16Array, Int32Array, ListArray, RunArray, + StringArray, StructArray, UnionArray, + }; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{Fields, Int8Type, Int16Type, Schema, UnionFields}; + use arrow::record_batch::RecordBatch; + use arrow::util::display::array_value_to_string; + use datafusion::prelude::SessionContext; + use datafusion_expr::AggregateUDF; fn list_type(element_type: DataType) -> DataType { DataType::List(Arc::new(Field::new_list_field(element_type, false))) @@ -321,6 +355,63 @@ mod tests { assert_eq!(actual, vec![1, 2]); } + fn assert_list_values( + value: &ScalarValue, + element_type: &DataType, + expected: &[&str], + ) -> Result<()> { + assert_eq!(value.data_type(), list_type(element_type.clone())); + let ScalarValue::List(array) = value else { + panic!("expected a list scalar") + }; + assert_list_array_values(array, element_type, expected) + } + + fn assert_list_array_values( + array: &ListArray, + element_type: &DataType, + expected: &[&str], + ) -> Result<()> { + assert_eq!(array.data_type(), &list_type(element_type.clone())); + array.to_data().validate_full()?; + + let values = array.value(0); + assert_eq!(values.logical_null_count(), 0); + let mut actual = (0..values.len()) + .map(|index| array_value_to_string(values.as_ref(), index)) + .collect::, _>>()?; + actual.sort_unstable(); + + let mut expected = expected + .iter() + .map(|value| (*value).to_string()) + .collect::>(); + expected.sort_unstable(); + assert_eq!(actual, expected); + Ok(()) + } + + fn assert_accumulator_outputs( + values: ArrayRef, + distinct: bool, + expected: &[&str], + ) -> Result<()> { + let element_type = values.data_type().clone(); + + let mut partial = accumulator(&element_type, distinct)?; + partial.update_batch(std::slice::from_ref(&values))?; + let state = partial.state()?; + assert_list_values(&state[0], &element_type, expected)?; + + let mut final_accumulator = accumulator(&element_type, distinct)?; + final_accumulator.merge_batch(&[state[0].to_array()?])?; + assert_list_values(&final_accumulator.evaluate()?, &element_type, expected)?; + + let mut single = accumulator(&element_type, distinct)?; + single.update_batch(&[values])?; + assert_list_values(&single.evaluate()?, &element_type, expected) + } + #[test] fn collect_types_have_non_nullable_elements() -> Result<()> { let element_type = DataType::Int32; @@ -330,12 +421,18 @@ mod tests { &SparkCollectList::new() as &dyn AggregateUDFImpl, &SparkCollectSet::new() as &dyn AggregateUDFImpl, ] { + assert!(!aggregate.is_nullable()); assert_eq!( aggregate.return_type(std::slice::from_ref(&element_type))?, expected ); let input_field = Arc::new(Field::new("input", element_type.clone(), true)); + let return_field = + aggregate.return_field(std::slice::from_ref(&input_field))?; + assert_eq!(return_field.data_type(), &expected); + assert!(!return_field.is_nullable()); + let state_fields = aggregate.state_fields(StateFieldsArgs { name: aggregate.name(), input_fields: &[input_field], @@ -349,6 +446,38 @@ mod tests { Ok(()) } + #[test] + fn empty_partial_state_has_non_nullable_elements() -> Result<()> { + let element_type = DataType::Int32; + let expected = list_type(element_type.clone()); + + for distinct in [false, true] { + let mut partial = accumulator(&element_type, distinct)?; + let state = partial.state()?; + assert_eq!(state.len(), 1); + assert!(state[0].is_null()); + assert_eq!(state[0].data_type(), expected); + + let ScalarValue::List(array) = &state[0] else { + panic!("expected a list scalar") + }; + let DataType::List(field) = array.data_type() else { + panic!("expected a list data type") + }; + assert_eq!(field.name(), "item"); + assert!(!field.is_nullable()); + + let mut final_accumulator = accumulator(&element_type, distinct)?; + final_accumulator.merge_batch(&[state[0].to_array()?])?; + let value = final_accumulator.evaluate()?; + assert!(!value.is_null()); + assert_eq!(value.data_type(), expected); + assert_empty_list(&value); + } + + Ok(()) + } + #[test] fn empty_results_have_non_nullable_elements() -> Result<()> { let expected = list_type(DataType::Int32); @@ -390,6 +519,12 @@ mod tests { None, )) as ArrayRef; + let scalar = + SingleRowListArrayBuilder::new(Arc::clone(&values)).build_list_scalar(); + let normalized = normalize_list_scalar(scalar, &expected)?; + assert_eq!(normalized.data_type(), expected); + assert_nested_values(&normalized); + for distinct in [false, true] { let mut partial = accumulator(&element_type, distinct)?; partial.update_batch(std::slice::from_ref(&values))?; @@ -412,6 +547,48 @@ mod tests { Ok(()) } + #[test] + fn nested_runtime_null_in_non_nullable_declared_field_is_rejected() -> Result<()> { + // Widening the output type would violate the aggregate's declared schema and + // reintroduce the AggregateExec schema mismatch this normalization prevents. + let declared_fields = + Fields::from(vec![Field::new("required", DataType::Int32, false)]); + let element_type = DataType::Struct(declared_fields); + let runtime_fields = + Fields::from(vec![Field::new("required", DataType::Int32, true)]); + let values = Arc::new(StructArray::new( + runtime_fields, + vec![Arc::new(Int32Array::from(vec![Some(1), None]))], + None, + )) as ArrayRef; + + let scalar = + SingleRowListArrayBuilder::new(Arc::clone(&values)).build_list_scalar(); + let error = + normalize_list_scalar(scalar, &list_type(element_type.clone())).unwrap_err(); + assert!( + error + .to_string() + .contains("Found unmasked nulls for non-nullable"), + "unexpected error: {error}" + ); + + for distinct in [false, true] { + let mut partial = accumulator(&element_type, distinct)?; + let error = partial + .update_batch(std::slice::from_ref(&values)) + .unwrap_err(); + assert!( + error + .to_string() + .contains("Found unmasked nulls for non-nullable"), + "unexpected error: {error}" + ); + } + + Ok(()) + } + #[test] fn normalization_reuses_matching_primitive_values() -> Result<()> { let values = Arc::new(Int32Array::from(vec![1, 2, 3])); @@ -431,4 +608,106 @@ mod tests { assert_eq!(normalized_values.values().as_ptr(), values_ptr); Ok(()) } + + #[test] + fn collect_list_handles_dictionary_with_unused_null() -> Result<()> { + let keys = Int8Array::from(vec![0, 0]); + let dictionary_values = Arc::new(StringArray::from(vec![Some("a"), None])); + let values = Arc::new(DictionaryArray::::try_new( + keys, + dictionary_values, + )?) as ArrayRef; + + assert_eq!(values.logical_null_count(), 0); + assert!(values.is_nullable()); + + assert_accumulator_outputs(Arc::clone(&values), false, &["a", "a"])?; + assert_accumulator_outputs(values, true, &["a"])?; + + Ok(()) + } + + #[test] + fn collect_aggregates_handle_sparse_union_inactive_nulls() -> Result<()> { + let fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("string", DataType::Utf8, false), + ], + )?; + let values = Arc::new(UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0_i8, 1, 0]), + None, + vec![ + Arc::new(Int32Array::from(vec![Some(1), None, Some(1)])), + Arc::new(StringArray::from(vec![None, Some("a"), None])), + ], + )?) as ArrayRef; + + assert_eq!(values.logical_null_count(), 0); + assert!(values.is_nullable()); + assert_accumulator_outputs( + Arc::clone(&values), + false, + &["{integer=1}", "{string=a}", "{integer=1}"], + )?; + assert_accumulator_outputs(values, true, &["{integer=1}", "{string=a}"])?; + + Ok(()) + } + + #[test] + fn collect_aggregates_handle_sliced_run_array_unused_null() -> Result<()> { + let run_ends = Int16Array::from(vec![2, 4]); + let run_values = StringArray::from(vec![Some("a"), None]); + let values = + Arc::new(RunArray::::try_new(&run_ends, &run_values)?.slice(0, 2)) + as ArrayRef; + + assert_eq!(values.logical_null_count(), 0); + assert!(values.is_nullable()); + assert_accumulator_outputs(Arc::clone(&values), false, &["a", "a"])?; + assert_accumulator_outputs(values, true, &["a"])?; + + Ok(()) + } + + #[tokio::test] + async fn collect_list_dictionary_sql() -> Result<()> { + let keys = Int8Array::from(vec![0, 0]); + let dictionary_values = Arc::new(StringArray::from(vec![Some("a"), None])); + let values = Arc::new(DictionaryArray::::try_new( + keys, + dictionary_values, + )?) as ArrayRef; + let element_type = values.data_type().clone(); + + let ctx = SessionContext::new(); + ctx.register_udaf(AggregateUDF::new_from_impl(SparkCollectList::new())); + ctx.register_batch( + "dictionary_input", + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "x", + element_type.clone(), + false, + )])), + vec![values], + )?, + )?; + + let batches = ctx + .sql("SELECT collect_list(x) AS values FROM dictionary_input") + .await? + .collect() + .await?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); + let value = ScalarValue::try_from_array(batches[0].column(0), 0)?; + assert_list_values(&value, &element_type, &["a", "a"])?; + + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt b/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt index c367c9cb7a6a6..3f8f01e05d301 100644 --- a/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt +++ b/datafusion/sqllogictest/test_files/spark/aggregate/collect.slt @@ -20,6 +20,12 @@ SELECT collect_list(a) FROM (VALUES (1), (2), (3)) AS t(a); ---- [1, 2, 3] +# collect_list drops null inputs, so its element field is non-nullable +query T +SELECT arrow_typeof(collect_list(a)) FROM (VALUES (1)) AS t(a); +---- +List(non-null Int64) + query ? SELECT collect_list(a) FROM (VALUES (1), (2), (2), (3), (1)) AS t(a); ---- @@ -59,6 +65,12 @@ SELECT array_sort(collect_set(a)) FROM (VALUES (1), (2), (3)) AS t(a); ---- [1, 2, 3] +# collect_set drops null inputs, so its element field is non-nullable +query T +SELECT arrow_typeof(collect_set(a)) FROM (VALUES (1)) AS t(a); +---- +List(non-null Int64) + query ? SELECT array_sort(collect_set(a)) FROM (VALUES (1), (2), (2), (3), (1)) AS t(a); ---- From 4f792b350a6a77dfd2830dbb3833a91901862b16 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sun, 13 Sep 2026 23:19:57 -0700 Subject: [PATCH 03/13] fix: normalize Spark collect retractions and retained input rows --- .../spark/src/function/aggregate/collect.rs | 84 ++++++++++++++++++- 1 file changed, 80 insertions(+), 4 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 3f3bf3e38741e..caea5b7c3d972 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, ArrayRef}; -use arrow::compute::cast; +use arrow::array::{Array, ArrayRef, UInt64Array}; +use arrow::compute::{cast, take}; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::utils::SingleRowListArrayBuilder; use datafusion_common::{Result, ScalarValue, internal_err}; @@ -245,6 +245,15 @@ impl NullToEmptyListAccumulator { if value.data_type() == field.data_type() { Ok(Arc::clone(value)) } else { + // Materialize only retained rows before narrowing nested fields. + // A slice can still reference null payload outside its logical rows. + let indices = match value.logical_nulls() { + Some(nulls) => UInt64Array::from_iter_values( + nulls.valid_indices().map(|index| index as u64), + ), + None => UInt64Array::from_iter_values(0..value.len() as u64), + }; + let value = take(value.as_ref(), &indices, None)?; Ok(cast(value.as_ref(), field.data_type())?) } } @@ -281,7 +290,11 @@ impl Accumulator for NullToEmptyListAccumulator { } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.inner.retract_batch(values) + let [value] = values else { + return self.inner.retract_batch(values); + }; + let value = self.normalize_input(value)?; + self.inner.retract_batch(&[value]) } fn supports_retract_batch(&self) -> bool { @@ -300,7 +313,7 @@ mod tests { DictionaryArray, Int8Array, Int16Array, Int32Array, ListArray, RunArray, StringArray, StructArray, UnionArray, }; - use arrow::buffer::ScalarBuffer; + use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{Fields, Int8Type, Int16Type, Schema, UnionFields}; use arrow::record_batch::RecordBatch; use arrow::util::display::array_value_to_string; @@ -589,6 +602,69 @@ mod tests { Ok(()) } + #[test] + fn retract_preserves_nested_type() -> Result<()> { + let element_type = DataType::Struct(Fields::from(vec![Field::new( + "required", + DataType::Int32, + false, + )])); + let values = Arc::new(StructArray::new( + Fields::from(vec![Field::new("required", DataType::Int32, true)]), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + None, + )) as ArrayRef; + + for distinct in [false, true] { + let mut acc = accumulator(&element_type, distinct)?; + acc.update_batch(std::slice::from_ref(&values))?; + acc.retract_batch(&[values.slice(0, 1)])?; + assert_list_values(&acc.evaluate()?, &element_type, &["{required: 2}"])?; + let state = acc.state()?; + assert_list_values(&state[0], &element_type, &["{required: 2}"])?; + let mut merged = accumulator(&element_type, distinct)?; + merged.merge_batch(&[state[0].to_array()?])?; + assert_list_values(&merged.evaluate()?, &element_type, &["{required: 2}"])?; + acc.retract_batch(&[values.slice(1, 1)])?; + assert_empty_list(&acc.evaluate()?); + } + Ok(()) + } + + #[test] + fn ignored_null_rows_do_not_narrow_nested_payload() -> Result<()> { + let element_type = + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))); + // A null list can retain a null child, even though every retained list + // satisfies the declared non-nullable item field. + let values = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1, 2])), + Arc::new(Int32Array::from(vec![None, Some(1)])), + Some(NullBuffer::from(vec![false, true])), + )) as ArrayRef; + values.to_data().validate_full()?; + + for distinct in [false, true] { + let mut acc = accumulator(&element_type, distinct)?; + acc.update_batch(std::slice::from_ref(&values))?; + assert_list_values(&acc.evaluate()?, &element_type, &["[1]"])?; + let state = acc.state()?; + assert_list_values(&state[0], &element_type, &["[1]"])?; + let mut merged = accumulator(&element_type, distinct)?; + merged.merge_batch(&[state[0].to_array()?])?; + assert_list_values(&merged.evaluate()?, &element_type, &["[1]"])?; + + acc.retract_batch(&[values.slice(0, 1)])?; + assert_list_values(&acc.evaluate()?, &element_type, &["[1]"])?; + acc.retract_batch(&[values.slice(1, 1)])?; + assert_empty_list(&acc.evaluate()?); + acc.update_batch(&[values.slice(0, 1)])?; + assert_empty_list(&acc.evaluate()?); + } + Ok(()) + } + #[test] fn normalization_reuses_matching_primitive_values() -> Result<()> { let values = Arc::new(Int32Array::from(vec![1, 2, 3])); From 735a4b30900028d96433ed8000a2573115d5b15c Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 11:49:36 -0700 Subject: [PATCH 04/13] fix: normalize Spark collect encoded values Signed-off-by: goutamadwant --- Cargo.lock | 1 + Cargo.toml | 1 + datafusion/spark/Cargo.toml | 1 + .../spark/src/function/aggregate/collect.rs | 376 +++++++++++++++++- 4 files changed, 367 insertions(+), 12 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7a68fcb5e63d7..782261411dd5e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2618,6 +2618,7 @@ name = "datafusion-spark" version = "55.0.0" dependencies = [ "arrow", + "arrow-select", "base64 0.23.1", "bigdecimal", "chrono", diff --git a/Cargo.toml b/Cargo.toml index d529f551e918c..9b4bf286fa9ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -113,6 +113,7 @@ arrow-ipc = { version = "59.2.0", default-features = false, features = [ "zstd", ] } arrow-ord = { version = "59.2.0", default-features = false } +arrow-select = { version = "59.2.0", default-features = false } arrow-schema = { version = "59.2.0", default-features = false } async-trait = "0.1.89" bigdecimal = "0.4.8" diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index f40708863d108..fc4b5b150d563 100644 --- a/datafusion/spark/Cargo.toml +++ b/datafusion/spark/Cargo.toml @@ -44,6 +44,7 @@ name = "datafusion_spark" [dependencies] arrow = { workspace = true } +arrow-select = { workspace = true } base64 = "0.23" bigdecimal = { workspace = true } chrono = { workspace = true } diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index caea5b7c3d972..fcba6b81976f1 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -15,9 +15,14 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, ArrayRef, UInt64Array}; +use arrow::array::{ + Array, ArrayRef, GenericListArray, GenericListViewArray, OffsetSizeTrait, + StructArray, UInt64Array, UnionArray, cast::AsArray, make_array, +}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::compute::{cast, take}; -use arrow::datatypes::{DataType, Field, FieldRef}; +use arrow::datatypes::{DataType, Field, FieldRef, UnionMode}; +use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion_common::utils::SingleRowListArrayBuilder; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; @@ -53,6 +58,238 @@ fn collect_type(element_type: DataType) -> DataType { DataType::List(Arc::new(Field::new_list_field(element_type, false))) } +fn identity_indices(len: usize) -> UInt64Array { + UInt64Array::from_iter_values(0..len as u64) +} + +/// Materialize only logically reachable values and align their nested type with +/// `target_type`. +/// +/// Arrow container arrays may retain unreachable null payload in null lists, +/// unused dictionary values, sparse-union children, and view backing buffers. +/// That payload is valid, but it makes `Array::is_nullable` conservative and +/// prevents the result from being embedded below Spark's non-null list field. +fn normalize_array( + value: &ArrayRef, + target_type: &DataType, + require_non_null: bool, +) -> Result { + match (value.data_type(), target_type) { + (DataType::List(_), DataType::List(field)) => normalize_list::(value, field), + (DataType::LargeList(_), DataType::LargeList(field)) => { + normalize_list::(value, field) + } + (DataType::ListView(_), DataType::ListView(field)) => { + normalize_list_view::(value, field) + } + (DataType::LargeListView(_), DataType::LargeListView(field)) => { + normalize_list_view::(value, field) + } + ( + DataType::Dictionary(source_key, _), + DataType::Dictionary(target_key, target_value), + ) if source_key == target_key => { + let compact = garbage_collect_any_dictionary(value.as_any_dictionary())?; + let dictionary = compact.as_any_dictionary(); + let values = normalize_array(dictionary.values(), target_value, true)?; + Ok(dictionary.with_values(values)) + } + ( + DataType::RunEndEncoded(source_run_ends, _), + DataType::RunEndEncoded(target_run_ends, target_value), + ) if source_run_ends.data_type() == target_run_ends.data_type() => { + // `take` rebuilds run ends from the selected logical rows and drops + // unused runs, including null runs outside a slice. + let compact = take(value.as_ref(), &identity_indices(value.len()), None)?; + let runs = compact.as_any_ree(); + let values = normalize_array(runs.values(), target_value.data_type(), true)?; + Ok(runs.with_values(values)) + } + (DataType::Struct(_), DataType::Struct(fields)) => { + let source = value.as_struct(); + let columns = fields + .iter() + .zip(source.columns()) + .map(|(field, column)| { + let column = if !field.is_nullable() { + materialize_masked_values( + column, + source.nulls(), + field.data_type(), + )? + } else { + Arc::clone(column) + }; + normalize_array(&column, field.data_type(), !field.is_nullable()) + }) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new( + fields.clone(), + columns, + source.nulls().cloned(), + )?)) + } + ( + DataType::Union(_, UnionMode::Sparse), + DataType::Union(fields, UnionMode::Sparse), + ) => normalize_sparse_union(value, fields, require_non_null), + _ => { + let value = if value.data_type() == target_type { + Arc::clone(value) + } else { + cast(value.as_ref(), target_type)? + }; + if require_non_null + && value.logical_null_count() == 0 + && value.is_nullable() + && value.to_data().nulls().is_some() + { + let data = value.to_data().into_builder().nulls(None).build()?; + Ok(make_array(data)) + } else { + Ok(value) + } + } + } +} + +fn materialize_masked_values( + value: &ArrayRef, + parent_nulls: Option<&arrow::buffer::NullBuffer>, + target_type: &DataType, +) -> Result { + let Some(parent_nulls) = parent_nulls else { + return Ok(Arc::clone(value)); + }; + let fallback = parent_nulls + .valid_indices() + .find(|index| value.is_valid(*index)); + let Some(fallback) = fallback else { + return ScalarValue::new_default(target_type)?.to_array_of_size(value.len()); + }; + let indices = UInt64Array::from_iter_values((0..value.len()).map(|index| { + if parent_nulls.is_valid(index) { + index as u64 + } else { + fallback as u64 + } + })); + Ok(take(value.as_ref(), &indices, None)?) +} + +fn normalize_list( + value: &ArrayRef, + field: &FieldRef, +) -> Result { + let source = value.as_list::(); + let source_offsets = source.value_offsets(); + let mut indices = Vec::new(); + let mut offsets = Vec::with_capacity(source.len() + 1); + offsets.push(Offset::zero()); + for index in 0..source.len() { + if source.is_valid(index) { + let start = source_offsets[index].as_usize(); + let end = source_offsets[index + 1].as_usize(); + indices.extend((start..end).map(|index| index as u64)); + } + offsets.push(Offset::from_usize(indices.len()).expect("list offset overflow")); + } + let values = take( + source.values().as_ref(), + &UInt64Array::from_iter_values(indices), + None, + )?; + let values = normalize_array(&values, field.data_type(), !field.is_nullable())?; + Ok(Arc::new(GenericListArray::::try_new( + Arc::clone(field), + OffsetBuffer::new(ScalarBuffer::from(offsets)), + values, + source.nulls().cloned(), + )?)) +} + +fn normalize_list_view( + value: &ArrayRef, + field: &FieldRef, +) -> Result { + let source = value.as_list_view::(); + let mut indices = Vec::new(); + let mut offsets = Vec::with_capacity(source.len()); + let mut sizes = Vec::with_capacity(source.len()); + for index in 0..source.len() { + offsets + .push(Offset::from_usize(indices.len()).expect("list view offset overflow")); + if source.is_valid(index) { + let start = source.value_offsets()[index].as_usize(); + let size = source.value_sizes()[index].as_usize(); + indices.extend((start..start + size).map(|index| index as u64)); + sizes.push(Offset::from_usize(size).expect("list view size overflow")); + } else { + sizes.push(Offset::zero()); + } + } + let values = take( + source.values().as_ref(), + &UInt64Array::from_iter_values(indices), + None, + )?; + let values = normalize_array(&values, field.data_type(), !field.is_nullable())?; + Ok(Arc::new(GenericListViewArray::::try_new( + Arc::clone(field), + ScalarBuffer::from(offsets), + ScalarBuffer::from(sizes), + values, + source.nulls().cloned(), + )?)) +} + +fn normalize_sparse_union( + value: &ArrayRef, + fields: &arrow::datatypes::UnionFields, + require_non_null: bool, +) -> Result { + let source = value + .as_any() + .downcast_ref::() + .expect("union array"); + let children = fields + .iter() + .map(|(type_id, field)| { + let child = source.child(type_id); + let child = if require_non_null { + let fallback = (0..source.len()).find(|index| { + source.type_id(*index) == type_id && child.is_valid(*index) + }); + match fallback { + Some(fallback) => { + let indices = UInt64Array::from_iter_values( + (0..source.len()).map(|index| { + if source.type_id(index) == type_id { + index as u64 + } else { + fallback as u64 + } + }), + ); + take(child.as_ref(), &indices, None)? + } + None => ScalarValue::new_default(field.data_type())? + .to_array_of_size(source.len())?, + } + } else { + Arc::clone(child) + }; + normalize_array(&child, field.data_type(), require_non_null) + }) + .collect::>>()?; + Ok(Arc::new(UnionArray::try_new( + fields.clone(), + source.type_ids().clone(), + None, + children, + )?)) +} + /// Rebuild an accumulator result with the aggregate's declared list field. /// /// The shared array aggregate accumulators use a nullable list field and can @@ -87,12 +324,7 @@ fn normalize_list_scalar( )); } - let values = array.value(0); - let values = if values.data_type() == field.data_type() { - values - } else { - cast(values.as_ref(), field.data_type())? - }; + let values = normalize_array(&array.value(0), field.data_type(), true)?; Ok(SingleRowListArrayBuilder::new(values) .with_field(field) .build_list_scalar()) @@ -254,7 +486,7 @@ impl NullToEmptyListAccumulator { None => UInt64Array::from_iter_values(0..value.len() as u64), }; let value = take(value.as_ref(), &indices, None)?; - Ok(cast(value.as_ref(), field.data_type())?) + normalize_array(&value, field.data_type(), true) } } } @@ -314,7 +546,7 @@ mod tests { StringArray, StructArray, UnionArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; - use arrow::datatypes::{Fields, Int8Type, Int16Type, Schema, UnionFields}; + use arrow::datatypes::{Fields, Int8Type, Int16Type, Int32Type, Schema, UnionFields}; use arrow::record_batch::RecordBatch; use arrow::util::display::array_value_to_string; use datafusion::prelude::SessionContext; @@ -415,14 +647,19 @@ mod tests { partial.update_batch(std::slice::from_ref(&values))?; let state = partial.state()?; assert_list_values(&state[0], &element_type, expected)?; + assert_list_values(&state[0].clone().compacted(), &element_type, expected)?; let mut final_accumulator = accumulator(&element_type, distinct)?; final_accumulator.merge_batch(&[state[0].to_array()?])?; - assert_list_values(&final_accumulator.evaluate()?, &element_type, expected)?; + let merged = final_accumulator.evaluate()?; + assert_list_values(&merged, &element_type, expected)?; + assert_list_values(&merged.compacted(), &element_type, expected)?; let mut single = accumulator(&element_type, distinct)?; single.update_batch(&[values])?; - assert_list_values(&single.evaluate()?, &element_type, expected) + let value = single.evaluate()?; + assert_list_values(&value, &element_type, expected)?; + assert_list_values(&value.compacted(), &element_type, expected) } #[test] @@ -665,6 +902,57 @@ mod tests { Ok(()) } + #[test] + fn nested_null_lists_do_not_narrow_unreachable_payload() -> Result<()> { + let element_type = DataType::List(Arc::new(Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))), + true, + ))); + let runtime_inner = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Int32, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1, 2])), + Arc::new(Int32Array::from(vec![None, Some(1)])), + Some(NullBuffer::from(vec![false, true])), + )); + let values = Arc::new(ListArray::new( + Arc::new(Field::new_list_field( + runtime_inner.data_type().clone(), + true, + )), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2])), + runtime_inner, + None, + )) as ArrayRef; + values.to_data().validate_full()?; + + for distinct in [false, true] { + let mut partial = accumulator(&element_type, distinct)?; + partial.update_batch(std::slice::from_ref(&values))?; + let state = partial.state()?; + assert_eq!(state[0].data_type(), list_type(element_type.clone())); + + let ScalarValue::List(array) = &state[0] else { + panic!("expected collect state") + }; + let collected_values = array.value(0); + let collected = collected_values.as_list::(); + let inner_values = collected.value(0); + let inner = inner_values.as_list::(); + assert!(inner.is_null(0)); + assert_eq!(inner.value(1).as_primitive::().value(0), 1); + + let mut merged = accumulator(&element_type, distinct)?; + merged.merge_batch(&[state[0].to_array()?])?; + assert_eq!( + merged.evaluate()?.data_type(), + list_type(element_type.clone()) + ); + partial.retract_batch(std::slice::from_ref(&values))?; + assert_empty_list(&partial.evaluate()?); + } + Ok(()) + } + #[test] fn normalization_reuses_matching_primitive_values() -> Result<()> { let values = Arc::new(Int32Array::from(vec![1, 2, 3])); @@ -786,4 +1074,68 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn collect_list_dictionary_grouped_and_window_sql() -> Result<()> { + let keys = Int8Array::from(vec![0, 0, 1, 1]); + let dictionary_values = + Arc::new(StringArray::from(vec![Some("a"), Some("b"), None])); + let values = Arc::new(DictionaryArray::::try_new( + keys, + dictionary_values, + )?) as ArrayRef; + let element_type = values.data_type().clone(); + + let ctx = SessionContext::new(); + ctx.register_udaf(AggregateUDF::new_from_impl(SparkCollectList::new())); + ctx.register_batch( + "dictionary_input", + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int32, false), + Field::new("id", DataType::Int32, false), + Field::new("x", element_type.clone(), false), + ])), + vec![ + Arc::new(Int32Array::from(vec![0, 0, 1, 1])), + Arc::new(Int32Array::from(vec![0, 1, 2, 3])), + values, + ], + )?, + )?; + + let grouped = ctx + .sql( + "SELECT g, collect_list(x) AS values \ + FROM dictionary_input GROUP BY g ORDER BY g", + ) + .await? + .collect() + .await?; + assert_eq!(grouped.iter().map(RecordBatch::num_rows).sum::(), 2); + let grouped = + arrow::compute::concat_batches(&grouped[0].schema(), grouped.iter())?; + for (row, expected) in [["a", "a"], ["b", "b"]].iter().enumerate() { + let value = ScalarValue::try_from_array(grouped.column(1), row)?; + assert_list_values(&value, &element_type, expected)?; + } + + let window = ctx + .sql( + "SELECT id, collect_list(x) OVER (\ + ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW\ + ) AS values FROM dictionary_input ORDER BY id", + ) + .await? + .collect() + .await?; + let window = arrow::compute::concat_batches(&window[0].schema(), window.iter())?; + let expected = [vec!["a"], vec!["a", "a"], vec!["a", "b"], vec!["b", "b"]]; + for (row, expected) in expected.iter().enumerate() { + let value = ScalarValue::try_from_array(window.column(1), row)?; + assert_list_values(&value, &element_type, expected)?; + } + + Ok(()) + } } From d7d26bda13ccc17a794c94807ed9d3087bd2c69d Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 12:16:39 -0700 Subject: [PATCH 05/13] fix: bound collect encoded normalization Signed-off-by: goutamadwant --- .../spark/src/function/aggregate/collect.rs | 630 ++++++++++++++++-- 1 file changed, 578 insertions(+), 52 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index fcba6b81976f1..3a7ef41e57255 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -16,21 +16,23 @@ // under the License. use arrow::array::{ - Array, ArrayRef, GenericListArray, GenericListViewArray, OffsetSizeTrait, - StructArray, UInt64Array, UnionArray, cast::AsArray, make_array, + Array, ArrayData, ArrayRef, GenericListArray, GenericListViewArray, OffsetSizeTrait, + PrimitiveArray, RunArray, StructArray, UInt64Array, UnionArray, cast::AsArray, + downcast_run_array, make_array, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::compute::{cast, take}; -use arrow::datatypes::{DataType, Field, FieldRef, UnionMode}; +use arrow::datatypes::{DataType, Field, FieldRef, RunEndIndexType, UnionMode}; use arrow_select::dictionary::garbage_collect_any_dictionary; use datafusion_common::utils::SingleRowListArrayBuilder; -use datafusion_common::{Result, ScalarValue, internal_err}; +use datafusion_common::{Result, ScalarValue, internal_datafusion_err, internal_err}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility}; use datafusion_functions_aggregate::array_agg::{ ArrayAggAccumulator, DistinctArrayAggAccumulator, }; +use std::collections::BTreeSet; use std::sync::Arc; // Spark implementation of collect_list/collect_set aggregate function. @@ -58,10 +60,6 @@ fn collect_type(element_type: DataType) -> DataType { DataType::List(Arc::new(Field::new_list_field(element_type, false))) } -fn identity_indices(len: usize) -> UInt64Array { - UInt64Array::from_iter_values(0..len as u64) -} - /// Materialize only logically reachable values and align their nested type with /// `target_type`. /// @@ -74,6 +72,9 @@ fn normalize_array( target_type: &DataType, require_non_null: bool, ) -> Result { + if value.data_type() == target_type && (!require_non_null || !value.is_nullable()) { + return Ok(Arc::clone(value)); + } match (value.data_type(), target_type) { (DataType::List(_), DataType::List(field)) => normalize_list::(value, field), (DataType::LargeList(_), DataType::LargeList(field)) => { @@ -92,18 +93,25 @@ fn normalize_array( let compact = garbage_collect_any_dictionary(value.as_any_dictionary())?; let dictionary = compact.as_any_dictionary(); let values = normalize_array(dictionary.values(), target_value, true)?; - Ok(dictionary.with_values(values)) + let dictionary = dictionary.with_values(values); + if dictionary.null_count() == 0 && dictionary.to_data().nulls().is_some() { + let data = dictionary.to_data().into_builder().nulls(None).build()?; + Ok(make_array(data)) + } else { + Ok(dictionary) + } } ( DataType::RunEndEncoded(source_run_ends, _), DataType::RunEndEncoded(target_run_ends, target_value), ) if source_run_ends.data_type() == target_run_ends.data_type() => { - // `take` rebuilds run ends from the selected logical rows and drops - // unused runs, including null runs outside a slice. - let compact = take(value.as_ref(), &identity_indices(value.len()), None)?; - let runs = compact.as_any_ree(); - let values = normalize_array(runs.values(), target_value.data_type(), true)?; - Ok(runs.with_values(values)) + // Work at the physical run count, not the potentially enormous + // logical length represented by those runs. + let value = value.as_ref(); + arrow::array::downcast_run_array! { + value => normalize_run_end(value, target_type, target_value.data_type()), + _ => internal_err!("collect_list/collect_set expected a run-end encoded array"), + } } (DataType::Struct(_), DataType::Struct(fields)) => { let source = value.as_struct(); @@ -130,9 +138,13 @@ fn normalize_array( )?)) } ( - DataType::Union(_, UnionMode::Sparse), + DataType::Union(source_fields, UnionMode::Sparse), DataType::Union(fields, UnionMode::Sparse), - ) => normalize_sparse_union(value, fields, require_non_null), + ) => normalize_sparse_union(value, source_fields, fields, require_non_null), + ( + DataType::Union(source_fields, UnionMode::Dense), + DataType::Union(fields, UnionMode::Dense), + ) => normalize_dense_union(value, source_fields, fields, require_non_null), _ => { let value = if value.data_type() == target_type { Arc::clone(value) @@ -161,11 +173,18 @@ fn materialize_masked_values( let Some(parent_nulls) = parent_nulls else { return Ok(Arc::clone(value)); }; - let fallback = parent_nulls - .valid_indices() + let mut valid_parent_indices = parent_nulls.valid_indices(); + let Some(first_valid_parent) = valid_parent_indices.next() else { + return ScalarValue::new_default(target_type)?.to_array_of_size(value.len()); + }; + let fallback = std::iter::once(first_valid_parent) + .chain(valid_parent_indices) .find(|index| value.is_valid(*index)); let Some(fallback) = fallback else { - return ScalarValue::new_default(target_type)?.to_array_of_size(value.len()); + // The nulls are logically visible because the parent rows are valid. + // Preserve them so the declared non-null field validation rejects the + // input instead of silently replacing user data with defaults. + return Ok(Arc::clone(value)); }; let indices = UInt64Array::from_iter_values((0..value.len()).map(|index| { if parent_nulls.is_valid(index) { @@ -177,6 +196,25 @@ fn materialize_masked_values( Ok(take(value.as_ref(), &indices, None)?) } +fn normalize_run_end( + source: &RunArray, + target_type: &DataType, + target_value_type: &DataType, +) -> Result { + let run_ends = + PrimitiveArray::::from_iter_values(source.run_ends().sliced_values()); + let values = source.values_slice(); + let values = normalize_array(&values, target_value_type, true)?; + + let data = ArrayData::builder(target_type.clone()) + .len(source.len()) + .add_child_data(run_ends.to_data()) + .add_child_data(values.to_data()) + .build()?; + data.validate_full()?; + Ok(make_array(data)) +} + fn normalize_list( value: &ArrayRef, field: &FieldRef, @@ -192,7 +230,9 @@ fn normalize_list( let end = source_offsets[index + 1].as_usize(); indices.extend((start..end).map(|index| index as u64)); } - offsets.push(Offset::from_usize(indices.len()).expect("list offset overflow")); + offsets.push(Offset::from_usize(indices.len()).ok_or_else(|| { + internal_datafusion_err!("list offset exceeds target offset type") + })?); } let values = take( source.values().as_ref(), @@ -213,27 +253,125 @@ fn normalize_list_view( field: &FieldRef, ) -> Result { let source = value.as_list_view::(); - let mut indices = Vec::new(); - let mut offsets = Vec::with_capacity(source.len()); - let mut sizes = Vec::with_capacity(source.len()); + let mut ranges = Vec::new(); + ranges.try_reserve(source.len()).map_err(|error| { + internal_datafusion_err!("failed to reserve list-view ranges: {error}") + })?; for index in 0..source.len() { - offsets - .push(Offset::from_usize(indices.len()).expect("list view offset overflow")); if source.is_valid(index) { let start = source.value_offsets()[index].as_usize(); let size = source.value_sizes()[index].as_usize(); - indices.extend((start..start + size).map(|index| index as u64)); - sizes.push(Offset::from_usize(size).expect("list view size overflow")); - } else { - sizes.push(Offset::zero()); + let end = start.checked_add(size).ok_or_else(|| { + internal_datafusion_err!("list-view offset plus size overflow") + })?; + if end > source.values().len() { + return internal_err!( + "list-view range {start}..{end} exceeds child length {}", + source.values().len() + ); + } + if start != end { + ranges.push((start, end)); + } + } + } + + // Copy the union of referenced ranges once. Overlapping views therefore + // remain bounded by the original backing array instead of sum(view_sizes). + ranges.sort_unstable_by_key(|(start, _)| *start); + let mut merged: Vec<(usize, usize, usize)> = Vec::new(); + merged.try_reserve(ranges.len()).map_err(|error| { + internal_datafusion_err!("failed to reserve merged list-view ranges: {error}") + })?; + for (start, end) in ranges { + match merged.last_mut() { + Some((_, merged_end, _)) if start <= *merged_end => { + *merged_end = (*merged_end).max(end); + } + _ => { + let new_start = match merged.last() { + Some((previous_start, previous_end, previous_new_start)) => { + previous_new_start + .checked_add(previous_end - previous_start) + .ok_or_else(|| { + internal_datafusion_err!( + "compacted list-view length overflow" + ) + })? + } + None => 0, + }; + merged.push((start, end, new_start)); + } } } + + let retained_len = match merged.last() { + Some((start, end, new_start)) => { + new_start.checked_add(end - start).ok_or_else(|| { + internal_datafusion_err!("compacted list-view length overflow") + })? + } + None => 0, + }; + let mut indices = Vec::new(); + indices.try_reserve(retained_len).map_err(|error| { + internal_datafusion_err!("failed to reserve compacted list-view values: {error}") + })?; + for (start, end, _) in &merged { + indices.extend((*start..*end).map(|index| index as u64)); + } + let values = take( source.values().as_ref(), &UInt64Array::from_iter_values(indices), None, )?; let values = normalize_array(&values, field.data_type(), !field.is_nullable())?; + + let mut offsets = Vec::new(); + let mut sizes = Vec::new(); + offsets.try_reserve(source.len()).map_err(|error| { + internal_datafusion_err!("failed to reserve list-view offsets: {error}") + })?; + sizes.try_reserve(source.len()).map_err(|error| { + internal_datafusion_err!("failed to reserve list-view sizes: {error}") + })?; + for index in 0..source.len() { + if source.is_null(index) { + offsets.push(Offset::zero()); + sizes.push(Offset::zero()); + continue; + } + let start = source.value_offsets()[index].as_usize(); + let size = source.value_sizes()[index].as_usize(); + if size == 0 { + offsets.push(Offset::zero()); + sizes.push(Offset::zero()); + continue; + } + let merged_index = merged.partition_point(|(_, end, _)| *end <= start); + let &(range_start, range_end, new_start) = + merged.get(merged_index).ok_or_else(|| { + internal_datafusion_err!("list-view range was not retained") + })?; + let end = start.checked_add(size).ok_or_else(|| { + internal_datafusion_err!("list-view offset plus size overflow") + })?; + if start < range_start || end > range_end { + return internal_err!("list-view range was not retained contiguously"); + } + let offset = new_start + .checked_add(start - range_start) + .ok_or_else(|| internal_datafusion_err!("list-view offset overflow"))?; + offsets.push(Offset::from_usize(offset).ok_or_else(|| { + internal_datafusion_err!("list-view offset exceeds target offset type") + })?); + sizes.push(Offset::from_usize(size).ok_or_else(|| { + internal_datafusion_err!("list-view size exceeds target offset type") + })?); + } + Ok(Arc::new(GenericListViewArray::::try_new( Arc::clone(field), ScalarBuffer::from(offsets), @@ -245,51 +383,203 @@ fn normalize_list_view( fn normalize_sparse_union( value: &ArrayRef, - fields: &arrow::datatypes::UnionFields, + source_fields: &arrow::datatypes::UnionFields, + target_fields: &arrow::datatypes::UnionFields, require_non_null: bool, ) -> Result { let source = value .as_any() .downcast_ref::() - .expect("union array"); - let children = fields + .ok_or_else(|| internal_datafusion_err!("expected sparse union array"))?; + if source_fields.len() != target_fields.len() { + return internal_err!( + "cannot normalize sparse union with {} runtime fields to {} declared fields", + source_fields.len(), + target_fields.len() + ); + } + let field_mapping = source_fields + .iter() + .zip(target_fields.iter()) + .collect::>(); + let type_ids = source + .type_ids() .iter() - .map(|(type_id, field)| { - let child = source.child(type_id); + .map(|source_type_id| { + field_mapping + .iter() + .find_map(|((source_id, _), (target_id, _))| { + (*source_id == *source_type_id).then_some(*target_id) + }) + .ok_or_else(|| { + internal_datafusion_err!( + "sparse union contains unknown runtime type id {source_type_id}" + ) + }) + }) + .collect::>>()?; + let children = field_mapping + .iter() + .map(|((source_type_id, _), (_, target_field))| { + let child = source.child(*source_type_id); let child = if require_non_null { - let fallback = (0..source.len()).find(|index| { - source.type_id(*index) == type_id && child.is_valid(*index) - }); - match fallback { - Some(fallback) => { + let first_active = + (0..source.len()).find(|index| source.type_id(*index) == *source_type_id); + match first_active { + None => ScalarValue::new_default(target_field.data_type())? + .to_array_of_size(source.len())?, + Some(first_active) => { + if (first_active..source.len()).any(|index| { + source.type_id(index) == *source_type_id + && child.is_null(index) + }) { + return internal_err!( + "Found unmasked nulls for non-nullable sparse union field {}", + target_field.name() + ); + } let indices = UInt64Array::from_iter_values( (0..source.len()).map(|index| { - if source.type_id(index) == type_id { + if source.type_id(index) == *source_type_id { index as u64 } else { - fallback as u64 + first_active as u64 } }), ); take(child.as_ref(), &indices, None)? } - None => ScalarValue::new_default(field.data_type())? - .to_array_of_size(source.len())?, } } else { Arc::clone(child) }; - normalize_array(&child, field.data_type(), require_non_null) + normalize_array(&child, target_field.data_type(), require_non_null) }) .collect::>>()?; Ok(Arc::new(UnionArray::try_new( - fields.clone(), - source.type_ids().clone(), + target_fields.clone(), + ScalarBuffer::from(type_ids), None, children, )?)) } +fn normalize_dense_union( + value: &ArrayRef, + source_fields: &arrow::datatypes::UnionFields, + target_fields: &arrow::datatypes::UnionFields, + require_non_null: bool, +) -> Result { + let source = value + .as_any() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("expected dense union array"))?; + if source_fields.len() != target_fields.len() { + return internal_err!( + "cannot normalize dense union with {} runtime fields to {} declared fields", + source_fields.len(), + target_fields.len() + ); + } + let field_mapping = source_fields + .iter() + .zip(target_fields.iter()) + .collect::>(); + + let mut referenced_offsets = vec![BTreeSet::new(); field_mapping.len()]; + for index in 0..source.len() { + let source_type_id = source.type_id(index); + let field_index = field_mapping + .iter() + .position(|((field_type_id, _), _)| *field_type_id == source_type_id) + .ok_or_else(|| { + internal_datafusion_err!( + "dense union contains unknown runtime type id {source_type_id}" + ) + })?; + referenced_offsets[field_index].insert(source.value_offset(index)); + } + + let mut children = Vec::new(); + let mut compacted_offsets = Vec::new(); + children.try_reserve(field_mapping.len()).map_err(|error| { + internal_datafusion_err!("failed to reserve dense union children: {error}") + })?; + compacted_offsets + .try_reserve(field_mapping.len()) + .map_err(|error| { + internal_datafusion_err!("failed to reserve dense union offsets: {error}") + })?; + for (((source_type_id, _), (_, target_field)), referenced) in + field_mapping.iter().zip(referenced_offsets) + { + let child = source.child(*source_type_id); + if let Some(offset) = referenced.last() + && *offset >= child.len() + { + return internal_err!( + "dense union offset {offset} exceeds child length {}", + child.len() + ); + } + let referenced = referenced.into_iter().collect::>(); + let indices = + UInt64Array::from_iter_values(referenced.iter().map(|offset| *offset as u64)); + let child = take(child.as_ref(), &indices, None)?; + let child = normalize_array(&child, target_field.data_type(), require_non_null)?; + if require_non_null && child.logical_null_count() != 0 { + return internal_err!( + "Found unmasked nulls for non-nullable dense union field {}", + target_field.name() + ); + } + children.push(child); + compacted_offsets.push(referenced); + } + + let mut type_ids = Vec::new(); + let mut offsets = Vec::new(); + type_ids.try_reserve(source.len()).map_err(|error| { + internal_datafusion_err!("failed to reserve dense union type ids: {error}") + })?; + offsets.try_reserve(source.len()).map_err(|error| { + internal_datafusion_err!("failed to reserve dense union offsets: {error}") + })?; + for index in 0..source.len() { + let source_type_id = source.type_id(index); + let field_index = field_mapping + .iter() + .position(|((field_type_id, _), _)| *field_type_id == source_type_id) + .ok_or_else(|| { + internal_datafusion_err!( + "dense union contains unknown runtime type id {source_type_id}" + ) + })?; + let target_type_id = field_mapping[field_index].1.0; + let source_offset = source.value_offset(index); + let compacted_offset = compacted_offsets[field_index] + .binary_search(&source_offset) + .map_err(|_| { + internal_datafusion_err!( + "dense union offset {source_offset} was not retained" + ) + })?; + type_ids.push(target_type_id); + offsets.push( + i32::try_from(compacted_offset).map_err(|_| { + internal_datafusion_err!("dense union offset exceeds i32") + })?, + ); + } + + Ok(Arc::new(UnionArray::try_new( + target_fields.clone(), + ScalarBuffer::from(type_ids), + Some(ScalarBuffer::from(offsets)), + children, + )?)) +} + /// Rebuild an accumulator result with the aggregate's declared list field. /// /// The shared array aggregate accumulators use a nullable list field and can @@ -542,11 +832,13 @@ impl Accumulator for NullToEmptyListAccumulator { mod tests { use super::*; use arrow::array::{ - DictionaryArray, Int8Array, Int16Array, Int32Array, ListArray, RunArray, - StringArray, StructArray, UnionArray, + DictionaryArray, Int8Array, Int16Array, Int32Array, Int64Array, ListArray, + ListViewArray, RunArray, StringArray, StructArray, UnionArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; - use arrow::datatypes::{Fields, Int8Type, Int16Type, Int32Type, Schema, UnionFields}; + use arrow::datatypes::{ + Fields, Int8Type, Int16Type, Int32Type, Int64Type, Schema, UnionFields, + }; use arrow::record_batch::RecordBatch; use arrow::util::display::array_value_to_string; use datafusion::prelude::SessionContext; @@ -839,6 +1131,29 @@ mod tests { Ok(()) } + #[test] + fn all_valid_struct_rows_with_only_null_required_values_are_rejected() -> Result<()> { + let declared_fields = + Fields::from(vec![Field::new("required", DataType::Int32, false)]); + let element_type = DataType::Struct(declared_fields); + let runtime_fields = + Fields::from(vec![Field::new("required", DataType::Int32, true)]); + let values = Arc::new(StructArray::new( + runtime_fields, + vec![Arc::new(Int32Array::from(vec![None, None]))], + None, + )) as ArrayRef; + + let error = normalize_array(&values, &element_type, true).unwrap_err(); + assert!( + error + .to_string() + .contains("Found unmasked nulls for non-nullable"), + "unexpected error: {error}" + ); + Ok(()) + } + #[test] fn retract_preserves_nested_type() -> Result<()> { let element_type = DataType::Struct(Fields::from(vec![Field::new( @@ -975,7 +1290,10 @@ mod tests { #[test] fn collect_list_handles_dictionary_with_unused_null() -> Result<()> { - let keys = Int8Array::from(vec![0, 0]); + let keys = Int8Array::new( + ScalarBuffer::from(vec![0_i8, 0]), + Some(NullBuffer::new_valid(2)), + ); let dictionary_values = Arc::new(StringArray::from(vec![Some("a"), None])); let values = Arc::new(DictionaryArray::::try_new( keys, @@ -984,6 +1302,7 @@ mod tests { assert_eq!(values.logical_null_count(), 0); assert!(values.is_nullable()); + assert!(values.nulls().is_some()); assert_accumulator_outputs(Arc::clone(&values), false, &["a", "a"])?; assert_accumulator_outputs(values, true, &["a"])?; @@ -1022,6 +1341,118 @@ mod tests { Ok(()) } + #[test] + fn collect_aggregates_remap_sparse_union_type_ids() -> Result<()> { + let runtime_fields = UnionFields::try_new( + vec![4, 9], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("string", DataType::Utf8, false), + ], + )?; + let declared_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("string", DataType::Utf8, false), + ], + )?; + let element_type = DataType::Union(declared_fields, UnionMode::Sparse); + let values = Arc::new(UnionArray::try_new( + runtime_fields, + ScalarBuffer::from(vec![4_i8, 9, 4]), + None, + vec![ + Arc::new(Int32Array::from(vec![Some(1), None, Some(1)])), + Arc::new(StringArray::from(vec![None, Some("a"), None])), + ], + )?) as ArrayRef; + + for distinct in [false, true] { + let expected = if distinct { + vec!["{integer=1}", "{string=a}"] + } else { + vec!["{integer=1}", "{string=a}", "{integer=1}"] + }; + let mut partial = accumulator(&element_type, distinct)?; + partial.update_batch(std::slice::from_ref(&values))?; + let state = partial.state()?; + assert_list_values(&state[0], &element_type, &expected)?; + + let mut merged = accumulator(&element_type, distinct)?; + merged.merge_batch(&[state[0].to_array()?])?; + assert_list_values(&merged.evaluate()?, &element_type, &expected)?; + } + Ok(()) + } + + #[test] + fn active_sparse_union_nulls_are_not_replaced_with_defaults() -> Result<()> { + let runtime_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, true), + Field::new("string", DataType::Utf8, true), + ], + )?; + let declared_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("string", DataType::Utf8, false), + ], + )?; + let element_type = DataType::Union(declared_fields, UnionMode::Sparse); + let values = Arc::new(UnionArray::try_new( + runtime_fields, + ScalarBuffer::from(vec![0_i8, 0]), + None, + vec![ + Arc::new(Int32Array::from(vec![None, None])), + Arc::new(StringArray::from(vec![None::<&str>, None])), + ], + )?) as ArrayRef; + + let error = normalize_array(&values, &element_type, true).unwrap_err(); + assert!( + error + .to_string() + .contains("Found unmasked nulls for non-nullable"), + "unexpected error: {error}" + ); + Ok(()) + } + + #[test] + fn collect_aggregates_compact_dense_union_children() -> Result<()> { + let fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("string", DataType::Utf8, false), + ], + )?; + let values = Arc::new(UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0_i8, 1, 0]), + Some(ScalarBuffer::from(vec![0_i32, 0, 0])), + vec![ + Arc::new(Int32Array::from(vec![Some(1), None])), + Arc::new(StringArray::from(vec![Some("a"), None])), + ], + )?) as ArrayRef; + + assert_eq!(values.logical_null_count(), 0); + assert!(values.is_nullable()); + assert_accumulator_outputs( + Arc::clone(&values), + false, + &["{integer=1}", "{string=a}", "{integer=1}"], + )?; + assert_accumulator_outputs(values, true, &["{integer=1}", "{string=a}"])?; + Ok(()) + } + #[test] fn collect_aggregates_handle_sliced_run_array_unused_null() -> Result<()> { let run_ends = Int16Array::from(vec![2, 4]); @@ -1038,6 +1469,44 @@ mod tests { Ok(()) } + #[test] + fn run_end_normalization_is_bounded_by_physical_runs() -> Result<()> { + let logical_len = i32::MAX as i64; + let run_ends = Int64Array::from(vec![logical_len]); + let values = StringArray::from(vec!["a"]); + let values = + Arc::new(RunArray::::try_new(&run_ends, &values)?) as ArrayRef; + let data_type = values.data_type().clone(); + + let normalized = normalize_array(&values, &data_type, true)?; + assert_eq!(normalized.len(), logical_len as usize); + assert_eq!(normalized.as_run::().values().len(), 1); + Ok(()) + } + + #[test] + fn list_view_normalization_copies_overlapping_backing_once() -> Result<()> { + const VIEW_COUNT: usize = 8_192; + const VALUE_COUNT: usize = 1_024; + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let values = Arc::new(ListViewArray::new( + Arc::clone(&field), + ScalarBuffer::from(vec![0_i32; VIEW_COUNT]), + ScalarBuffer::from(vec![VALUE_COUNT as i32; VIEW_COUNT]), + Arc::new(Int32Array::from_iter_values(0..VALUE_COUNT as i32)), + None, + )) as ArrayRef; + let data_type = values.data_type().clone(); + + let normalized = normalize_array(&values, &data_type, true)?; + let normalized = normalized.as_list_view::(); + assert_eq!(normalized.len(), VIEW_COUNT); + assert_eq!(normalized.values().len(), VALUE_COUNT); + assert_eq!(normalized.value_offsets()[VIEW_COUNT - 1], 0); + assert_eq!(normalized.value_sizes()[VIEW_COUNT - 1], VALUE_COUNT as i32); + Ok(()) + } + #[tokio::test] async fn collect_list_dictionary_sql() -> Result<()> { let keys = Int8Array::from(vec![0, 0]); @@ -1077,7 +1546,10 @@ mod tests { #[tokio::test] async fn collect_list_dictionary_grouped_and_window_sql() -> Result<()> { - let keys = Int8Array::from(vec![0, 0, 1, 1]); + let keys = Int8Array::new( + ScalarBuffer::from(vec![0_i8, 0, 1, 1]), + Some(NullBuffer::new_valid(4)), + ); let dictionary_values = Arc::new(StringArray::from(vec![Some("a"), Some("b"), None])); let values = Arc::new(DictionaryArray::::try_new( @@ -1085,6 +1557,7 @@ mod tests { dictionary_values, )?) as ArrayRef; let element_type = values.data_type().clone(); + assert!(values.nulls().is_some()); let ctx = SessionContext::new(); ctx.register_udaf(AggregateUDF::new_from_impl(SparkCollectList::new())); @@ -1138,4 +1611,57 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn collect_list_dense_union_grouped_sql() -> Result<()> { + let fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("string", DataType::Utf8, false), + ], + )?; + let values = Arc::new(UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0_i8, 1, 0]), + Some(ScalarBuffer::from(vec![0_i32, 0, 0])), + vec![ + Arc::new(Int32Array::from(vec![Some(1), None])), + Arc::new(StringArray::from(vec![Some("a"), None])), + ], + )?) as ArrayRef; + let element_type = values.data_type().clone(); + + let ctx = SessionContext::new(); + ctx.register_udaf(AggregateUDF::new_from_impl(SparkCollectList::new())); + ctx.register_batch( + "dense_union_input", + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int32, false), + Field::new("x", element_type.clone(), false), + ])), + vec![Arc::new(Int32Array::from(vec![0, 0, 1])), values], + )?, + )?; + + let grouped = ctx + .sql( + "SELECT g, collect_list(x) AS values \ + FROM dense_union_input GROUP BY g ORDER BY g", + ) + .await? + .collect() + .await?; + let grouped = + arrow::compute::concat_batches(&grouped[0].schema(), grouped.iter())?; + for (row, expected) in [vec!["{integer=1}", "{string=a}"], vec!["{integer=1}"]] + .iter() + .enumerate() + { + let value = ScalarValue::try_from_array(grouped.column(1), row)?; + assert_list_values(&value, &element_type, expected)?; + } + Ok(()) + } } From eb0a76d94dd4f7d210f44ba9d136925d6d0f04e4 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 12:38:17 -0700 Subject: [PATCH 06/13] fix: harden collect encoded normalization Signed-off-by: goutamadwant --- .../spark/src/function/aggregate/collect.rs | 110 +++++++++++++++++- 1 file changed, 109 insertions(+), 1 deletion(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 3a7ef41e57255..75fb0046bea13 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -453,7 +453,15 @@ fn normalize_sparse_union( } else { Arc::clone(child) }; - normalize_array(&child, target_field.data_type(), require_non_null) + let child = + normalize_array(&child, target_field.data_type(), require_non_null)?; + if require_non_null && child.logical_null_count() != 0 { + return internal_err!( + "Found unmasked nulls for non-nullable sparse union field {}", + target_field.name() + ); + } + Ok(child) }) .collect::>>()?; Ok(Arc::new(UnionArray::try_new( @@ -766,6 +774,17 @@ impl NullToEmptyListAccumulator { }; if value.data_type() == field.data_type() { Ok(Arc::clone(value)) + } else if matches!( + (value.data_type(), field.data_type()), + ( + DataType::RunEndEncoded(source_run_ends, _), + DataType::RunEndEncoded(target_run_ends, _) + ) if source_run_ends.data_type() == target_run_ends.data_type() + ) { + // Normalizing a run-end encoded array through `take` requires an + // index for every logical row. Dispatch it directly so work stays + // bounded by the physical run count. + normalize_array(value, field.data_type(), false) } else { // Materialize only retained rows before narrowing nested fields. // A slice can still reference null payload outside its logical rows. @@ -1423,6 +1442,62 @@ mod tests { Ok(()) } + fn assert_active_sparse_union_logical_null_is_rejected( + active_child: ArrayRef, + ) -> Result<()> { + assert!(active_child.is_valid(0)); + assert_eq!(active_child.logical_null_count(), 1); + + let runtime_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("logical", active_child.data_type().clone(), true), + Field::new("integer", DataType::Int32, true), + ], + )?; + let declared_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("logical", active_child.data_type().clone(), false), + Field::new("integer", DataType::Int32, false), + ], + )?; + let element_type = DataType::Union(declared_fields, UnionMode::Sparse); + let values = Arc::new(UnionArray::try_new( + runtime_fields, + ScalarBuffer::from(vec![0_i8]), + None, + vec![active_child, Arc::new(Int32Array::from(vec![Some(1)]))], + )?) as ArrayRef; + + let error = normalize_array(&values, &element_type, true).unwrap_err(); + assert!( + error.to_string().contains( + "Found unmasked nulls for non-nullable sparse union field logical" + ), + "unexpected error: {error}" + ); + Ok(()) + } + + #[test] + fn active_sparse_union_dictionary_logical_null_is_rejected() -> Result<()> { + let keys = Int8Array::from(vec![0]); + let values = Arc::new(StringArray::from(vec![None::<&str>])); + let dictionary = + Arc::new(DictionaryArray::::try_new(keys, values)?) as ArrayRef; + assert_active_sparse_union_logical_null_is_rejected(dictionary) + } + + #[test] + fn active_sparse_union_run_end_logical_null_is_rejected() -> Result<()> { + let run_ends = Int16Array::from(vec![1]); + let values = StringArray::from(vec![None::<&str>]); + let run = + Arc::new(RunArray::::try_new(&run_ends, &values)?) as ArrayRef; + assert_active_sparse_union_logical_null_is_rejected(run) + } + #[test] fn collect_aggregates_compact_dense_union_children() -> Result<()> { let fields = UnionFields::try_new( @@ -1484,6 +1559,39 @@ mod tests { Ok(()) } + #[test] + fn collect_input_run_end_normalization_is_bounded_by_physical_runs() -> Result<()> { + let logical_len = i32::MAX as i64; + let run_ends = Int64Array::from(vec![logical_len]); + let values = StringArray::from(vec!["a"]); + let values = + Arc::new(RunArray::::try_new(&run_ends, &values)?) as ArrayRef; + let DataType::RunEndEncoded(run_ends_field, values_field) = values.data_type() + else { + panic!("expected run-end encoded values") + }; + let target_type = DataType::RunEndEncoded( + Arc::clone(run_ends_field), + Arc::new(Field::new( + values_field.name(), + DataType::LargeUtf8, + values_field.is_nullable(), + )), + ); + let accumulator = NullToEmptyListAccumulator::new( + ArrayAggAccumulator::try_new(&target_type, true)?, + list_type(target_type.clone()), + ); + + let normalized = accumulator.normalize_input(&values)?; + assert_eq!(normalized.data_type(), &target_type); + assert_eq!(normalized.len(), logical_len as usize); + let normalized = normalized.as_run::(); + assert_eq!(normalized.values().len(), 1); + assert_eq!(normalized.values().data_type(), &DataType::LargeUtf8); + Ok(()) + } + #[test] fn list_view_normalization_copies_overlapping_backing_once() -> Result<()> { const VIEW_COUNT: usize = 8_192; From 3cf08706bb08826a5cb8af4422e6575a169eecea Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 12:48:59 -0700 Subject: [PATCH 07/13] fix: handle inactive sparse union variants --- datafusion/common/src/scalar/mod.rs | 32 ++++++-- .../spark/src/function/aggregate/collect.rs | 80 ++++++++++++++++--- 2 files changed, 93 insertions(+), 19 deletions(-) diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index fb08ef284280e..c146705245a86 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -5008,12 +5008,32 @@ impl ScalarValue { macro_rules! gc_list { ($field:expr, $offset_type:ty, $array_type:ty) => {{ let list = array.as_list::<$offset_type>(); - Arc::new(<$array_type>::new( - Arc::clone($field), - list.offsets().clone(), - ScalarValue::compact_view_buffers(Arc::clone(list.values())), - list.nulls().cloned(), - )) as ArrayRef + let offsets = list.offsets().clone(); + let values = ScalarValue::compact_view_buffers(Arc::clone(list.values())); + if !$field.is_nullable() + && values.is_nullable() + && values.logical_null_count() == 0 + { + let data = ArrayData::builder( + GenericListArray::<$offset_type>::DATA_TYPE_CONSTRUCTOR( + Arc::clone($field), + ), + ) + .len(list.len()) + .nulls(list.nulls().cloned()) + .add_buffer(offsets.into_inner().into_inner()) + .add_child_data(values.to_data()) + .build() + .expect("compacted list array should contain valid data"); + Arc::new(GenericListArray::<$offset_type>::from(data)) as ArrayRef + } else { + Arc::new(<$array_type>::new( + Arc::clone($field), + offsets, + values, + list.nulls().cloned(), + )) as ArrayRef + } }}; } // Macro for the i32/i64-offset list-view pair (ListView / LargeListView). diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 75fb0046bea13..405f9381121b9 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -422,9 +422,12 @@ fn normalize_sparse_union( .iter() .map(|((source_type_id, _), (_, target_field))| { let child = source.child(*source_type_id); + let first_active = if require_non_null { + (0..source.len()).find(|index| source.type_id(*index) == *source_type_id) + } else { + None + }; let child = if require_non_null { - let first_active = - (0..source.len()).find(|index| source.type_id(*index) == *source_type_id); match first_active { None => ScalarValue::new_default(target_field.data_type())? .to_array_of_size(source.len())?, @@ -455,7 +458,10 @@ fn normalize_sparse_union( }; let child = normalize_array(&child, target_field.data_type(), require_non_null)?; - if require_non_null && child.logical_null_count() != 0 { + if require_non_null + && first_active.is_some() + && child.logical_null_count() != 0 + { return internal_err!( "Found unmasked nulls for non-nullable sparse union field {}", target_field.name() @@ -852,7 +858,7 @@ mod tests { use super::*; use arrow::array::{ DictionaryArray, Int8Array, Int16Array, Int32Array, Int64Array, ListArray, - ListViewArray, RunArray, StringArray, StructArray, UnionArray, + ListViewArray, NullArray, RunArray, StringArray, StructArray, UnionArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ @@ -953,24 +959,32 @@ mod tests { expected: &[&str], ) -> Result<()> { let element_type = values.data_type().clone(); + assert_accumulator_outputs_with_type(values, &element_type, distinct, expected) + } - let mut partial = accumulator(&element_type, distinct)?; + fn assert_accumulator_outputs_with_type( + values: ArrayRef, + element_type: &DataType, + distinct: bool, + expected: &[&str], + ) -> Result<()> { + let mut partial = accumulator(element_type, distinct)?; partial.update_batch(std::slice::from_ref(&values))?; let state = partial.state()?; - assert_list_values(&state[0], &element_type, expected)?; - assert_list_values(&state[0].clone().compacted(), &element_type, expected)?; + assert_list_values(&state[0], element_type, expected)?; + assert_list_values(&state[0].clone().compacted(), element_type, expected)?; - let mut final_accumulator = accumulator(&element_type, distinct)?; + let mut final_accumulator = accumulator(element_type, distinct)?; final_accumulator.merge_batch(&[state[0].to_array()?])?; let merged = final_accumulator.evaluate()?; - assert_list_values(&merged, &element_type, expected)?; - assert_list_values(&merged.compacted(), &element_type, expected)?; + assert_list_values(&merged, element_type, expected)?; + assert_list_values(&merged.compacted(), element_type, expected)?; - let mut single = accumulator(&element_type, distinct)?; + let mut single = accumulator(element_type, distinct)?; single.update_batch(&[values])?; let value = single.evaluate()?; - assert_list_values(&value, &element_type, expected)?; - assert_list_values(&value.compacted(), &element_type, expected) + assert_list_values(&value, element_type, expected)?; + assert_list_values(&value.compacted(), element_type, expected) } #[test] @@ -1360,6 +1374,46 @@ mod tests { Ok(()) } + #[test] + fn collect_aggregates_handle_inactive_sparse_union_null_variant() -> Result<()> { + let runtime_fields = UnionFields::try_new( + vec![4, 9], + vec![ + Field::new("null", DataType::Null, true), + Field::new("integer", DataType::Int32, false), + ], + )?; + let declared_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("null", DataType::Null, false), + Field::new("integer", DataType::Int32, false), + ], + )?; + let element_type = DataType::Union(declared_fields, UnionMode::Sparse); + let values = Arc::new(UnionArray::try_new( + runtime_fields, + ScalarBuffer::from(vec![9_i8, 9]), + None, + vec![ + Arc::new(NullArray::new(2)), + Arc::new(Int32Array::from(vec![1, 2])), + ], + )?) as ArrayRef; + + assert_eq!(values.logical_null_count(), 0); + assert!(values.is_nullable()); + for distinct in [false, true] { + assert_accumulator_outputs_with_type( + Arc::clone(&values), + &element_type, + distinct, + &["{integer=1}", "{integer=2}"], + )?; + } + Ok(()) + } + #[test] fn collect_aggregates_remap_sparse_union_type_ids() -> Result<()> { let runtime_fields = UnionFields::try_new( From cdba3e3331b405383f1653a753fea299b184c94e Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 12:57:44 -0700 Subject: [PATCH 08/13] fix: preserve nullable run values during collect normalization Signed-off-by: goutamadwant --- .../spark/src/function/aggregate/collect.rs | 82 ++++++++++++++++++- 1 file changed, 80 insertions(+), 2 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 405f9381121b9..8fb3a2b8e22c5 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -109,7 +109,12 @@ fn normalize_array( // logical length represented by those runs. let value = value.as_ref(); arrow::array::downcast_run_array! { - value => normalize_run_end(value, target_type, target_value.data_type()), + value => normalize_run_end( + value, + target_type, + target_value.data_type(), + require_non_null, + ), _ => internal_err!("collect_list/collect_set expected a run-end encoded array"), } } @@ -200,11 +205,12 @@ fn normalize_run_end( source: &RunArray, target_type: &DataType, target_value_type: &DataType, + require_non_null: bool, ) -> Result { let run_ends = PrimitiveArray::::from_iter_values(source.run_ends().sliced_values()); let values = source.values_slice(); - let values = normalize_array(&values, target_value_type, true)?; + let values = normalize_array(&values, target_value_type, require_non_null)?; let data = ArrayData::builder(target_type.clone()) .len(source.len()) @@ -1646,6 +1652,78 @@ mod tests { Ok(()) } + #[test] + fn collect_input_run_end_normalization_preserves_null_runs() -> Result<()> { + let runtime_fields = UnionFields::try_new( + vec![4, 9], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("null", DataType::Null, true), + ], + )?; + let declared_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("null", DataType::Null, false), + ], + )?; + let runtime_values = UnionArray::try_new( + runtime_fields, + ScalarBuffer::from(vec![4_i8, 9]), + None, + vec![ + Arc::new(Int32Array::from(vec![Some(1), None])), + Arc::new(NullArray::new(2)), + ], + )?; + assert_eq!(runtime_values.logical_null_count(), 1); + + let run_ends = Int16Array::from(vec![2, 4]); + let values = Arc::new(RunArray::::try_new(&run_ends, &runtime_values)?) + as ArrayRef; + assert_eq!(values.logical_null_count(), 2); + + let DataType::RunEndEncoded(run_ends_field, values_field) = values.data_type() + else { + panic!("expected run-end encoded values") + }; + let element_type = DataType::RunEndEncoded( + Arc::clone(run_ends_field), + Arc::new(Field::new( + values_field.name(), + DataType::Union(declared_fields, UnionMode::Sparse), + true, + )), + ); + let mut accumulator = NullToEmptyListAccumulator::new( + ArrayAggAccumulator::try_new(&element_type, true)?, + list_type(element_type.clone()), + ); + + let normalized = accumulator.normalize_input(&values)?; + assert_eq!(normalized.data_type(), &element_type); + assert_eq!(normalized.len(), 4); + assert_eq!(normalized.as_run::().values().len(), 2); + + accumulator.update_batch(std::slice::from_ref(&values))?; + assert_list_values( + &accumulator.evaluate()?, + &element_type, + &["{integer=1}", "{integer=1}"], + )?; + + accumulator.retract_batch(&[values.slice(2, 2)])?; + assert_list_values( + &accumulator.evaluate()?, + &element_type, + &["{integer=1}", "{integer=1}"], + )?; + accumulator.retract_batch(&[values.slice(0, 2)])?; + assert_empty_list(&accumulator.evaluate()?); + Ok(()) + } + #[test] fn list_view_normalization_copies_overlapping_backing_once() -> Result<()> { const VIEW_COUNT: usize = 8_192; From 985de9887153aa7c600b4c8d729554739aa4a86e Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 13:04:21 -0700 Subject: [PATCH 09/13] fix: preserve nullable dictionary values during collect normalization Signed-off-by: goutamadwant --- .../spark/src/function/aggregate/collect.rs | 82 ++++++++++++++++++- 1 file changed, 81 insertions(+), 1 deletion(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 8fb3a2b8e22c5..9cef5f5e0f6fc 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -92,7 +92,8 @@ fn normalize_array( ) if source_key == target_key => { let compact = garbage_collect_any_dictionary(value.as_any_dictionary())?; let dictionary = compact.as_any_dictionary(); - let values = normalize_array(dictionary.values(), target_value, true)?; + let values = + normalize_array(dictionary.values(), target_value, require_non_null)?; let dictionary = dictionary.with_values(values); if dictionary.null_count() == 0 && dictionary.to_data().nulls().is_some() { let data = dictionary.to_data().into_builder().nulls(None).build()?; @@ -1349,6 +1350,85 @@ mod tests { Ok(()) } + #[test] + fn nullable_nested_dictionary_preserves_logical_null_union_value() -> Result<()> { + let runtime_union_fields = UnionFields::try_new( + vec![4, 9], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("null", DataType::Null, true), + ], + )?; + let declared_union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("null", DataType::Null, false), + ], + )?; + let runtime_union = Arc::new(UnionArray::try_new( + runtime_union_fields, + ScalarBuffer::from(vec![4_i8, 9]), + None, + vec![ + Arc::new(Int32Array::from(vec![Some(1), None])), + Arc::new(NullArray::new(2)), + ], + )?) as ArrayRef; + assert_eq!(runtime_union.logical_null_count(), 1); + + let dictionary = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0, 1]), + runtime_union, + )?) as ArrayRef; + assert_eq!(dictionary.logical_null_count(), 1); + + let runtime_fields = Fields::from(vec![Field::new( + "optional", + dictionary.data_type().clone(), + true, + )]); + let values = Arc::new(StructArray::try_new( + runtime_fields, + vec![dictionary], + None, + )?) as ArrayRef; + let target_dictionary_type = DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Union(declared_union_fields, UnionMode::Sparse)), + ); + let element_type = DataType::Struct(Fields::from(vec![Field::new( + "optional", + target_dictionary_type.clone(), + true, + )])); + + let assert_output = |value: ScalarValue, + expected_len: usize, + expected_null_count: usize| + -> Result<()> { + let ScalarValue::List(array) = value else { + panic!("expected a list scalar") + }; + let collected = array.value(0); + assert_eq!(collected.len(), expected_len); + let dictionary = collected.as_struct().column(0); + assert_eq!(dictionary.data_type(), &target_dictionary_type); + assert_eq!(dictionary.logical_null_count(), expected_null_count); + Ok(()) + }; + + let mut accumulator = accumulator(&element_type, false)?; + accumulator.update_batch(std::slice::from_ref(&values))?; + assert_output(accumulator.evaluate()?, 2, 1)?; + + accumulator.retract_batch(&[values.slice(0, 1)])?; + assert_output(accumulator.evaluate()?, 1, 1)?; + accumulator.retract_batch(&[values.slice(1, 1)])?; + assert_empty_list(&accumulator.evaluate()?); + Ok(()) + } + #[test] fn collect_aggregates_handle_sparse_union_inactive_nulls() -> Result<()> { let fields = UnionFields::try_new( From ae8c9c021347a9989ecde9f527911936d3335064 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 13:22:05 -0700 Subject: [PATCH 10/13] fix: compact exact encoded collect values Signed-off-by: goutamadwant --- .../spark/src/function/aggregate/collect.rs | 252 ++++++++++++++++-- 1 file changed, 223 insertions(+), 29 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 9cef5f5e0f6fc..5a732d4206e52 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -60,6 +60,50 @@ fn collect_type(element_type: DataType) -> DataType { DataType::List(Arc::new(Field::new_list_field(element_type, false))) } +/// Whether an exact-type array still needs normalization to discard physical +/// storage that is not part of its logical value. +fn has_compactable_storage(value: &ArrayRef) -> bool { + match value.data_type() { + DataType::Dictionary(_, _) + | DataType::RunEndEncoded(_, _) + | DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Union(_, _) => true, + DataType::List(_) => { + let list = value.as_list::(); + list.null_count() != 0 + || list + .value_offsets() + .first() + .is_some_and(|offset| *offset != 0) + || list + .value_offsets() + .last() + .is_some_and(|offset| *offset as usize != list.values().len()) + || has_compactable_storage(list.values()) + } + DataType::LargeList(_) => { + let list = value.as_list::(); + list.null_count() != 0 + || list + .value_offsets() + .first() + .is_some_and(|offset| *offset != 0) + || list + .value_offsets() + .last() + .is_some_and(|offset| *offset as usize != list.values().len()) + || has_compactable_storage(list.values()) + } + DataType::Struct(_) => value + .as_struct() + .columns() + .iter() + .any(has_compactable_storage), + _ => false, + } +} + /// Materialize only logically reachable values and align their nested type with /// `target_type`. /// @@ -72,7 +116,10 @@ fn normalize_array( target_type: &DataType, require_non_null: bool, ) -> Result { - if value.data_type() == target_type && (!require_non_null || !value.is_nullable()) { + if value.data_type() == target_type + && !has_compactable_storage(value) + && (!require_non_null || !value.is_nullable()) + { return Ok(Arc::clone(value)); } match (value.data_type(), target_type) { @@ -429,39 +476,33 @@ fn normalize_sparse_union( .iter() .map(|((source_type_id, _), (_, target_field))| { let child = source.child(*source_type_id); - let first_active = if require_non_null { - (0..source.len()).find(|index| source.type_id(*index) == *source_type_id) - } else { - None - }; - let child = if require_non_null { - match first_active { - None => ScalarValue::new_default(target_field.data_type())? - .to_array_of_size(source.len())?, - Some(first_active) => { - if (first_active..source.len()).any(|index| { + let first_active = + (0..source.len()).find(|index| source.type_id(*index) == *source_type_id); + let child = match first_active { + None => ScalarValue::new_default(target_field.data_type())? + .to_array_of_size(source.len())?, + Some(first_active) => { + if require_non_null + && (first_active..source.len()).any(|index| { source.type_id(index) == *source_type_id && child.is_null(index) - }) { - return internal_err!( - "Found unmasked nulls for non-nullable sparse union field {}", - target_field.name() - ); - } - let indices = UInt64Array::from_iter_values( - (0..source.len()).map(|index| { - if source.type_id(index) == *source_type_id { - index as u64 - } else { - first_active as u64 - } - }), + }) + { + return internal_err!( + "Found unmasked nulls for non-nullable sparse union field {}", + target_field.name() ); - take(child.as_ref(), &indices, None)? } + let indices = + UInt64Array::from_iter_values((0..source.len()).map(|index| { + if source.type_id(index) == *source_type_id { + index as u64 + } else { + first_active as u64 + } + })); + take(child.as_ref(), &indices, None)? } - } else { - Arc::clone(child) }; let child = normalize_array(&child, target_field.data_type(), require_non_null)?; @@ -1582,6 +1623,88 @@ mod tests { Ok(()) } + #[test] + fn exact_type_unions_reject_active_nulls() -> Result<()> { + let fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, false), + Field::new("string", DataType::Utf8, false), + ], + )?; + let children = || -> Vec { + vec![ + Arc::new(Int32Array::from(vec![None])), + Arc::new(StringArray::from(vec![Some("unused")])), + ] + }; + + let sparse = Arc::new(UnionArray::try_new( + fields.clone(), + ScalarBuffer::from(vec![0_i8]), + None, + children(), + )?) as ArrayRef; + assert!(sparse.is_nullable()); + assert_eq!(sparse.logical_null_count(), 1); + let error = normalize_array(&sparse, sparse.data_type(), true).unwrap_err(); + assert!( + error + .to_string() + .contains("non-nullable sparse union field integer"), + "unexpected error: {error}" + ); + + let dense = Arc::new(UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0_i8]), + Some(ScalarBuffer::from(vec![0_i32])), + children(), + )?) as ArrayRef; + assert!(dense.is_nullable()); + assert_eq!(dense.logical_null_count(), 1); + let error = normalize_array(&dense, dense.data_type(), true).unwrap_err(); + assert!( + error + .to_string() + .contains("non-nullable dense union field integer"), + "unexpected error: {error}" + ); + Ok(()) + } + + #[test] + fn exact_type_encoded_wrappers_reject_active_union_nulls() -> Result<()> { + let union = || -> Result { + let fields = UnionFields::try_new( + vec![0], + vec![Field::new("integer", DataType::Int32, false)], + )?; + Ok(Arc::new(UnionArray::try_new( + fields, + ScalarBuffer::from(vec![0_i8]), + None, + vec![Arc::new(Int32Array::from(vec![None]))], + )?) as ArrayRef) + }; + + let dictionary = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0]), + union()?, + )?) as ArrayRef; + assert!(dictionary.is_nullable()); + assert!(normalize_array(&dictionary, dictionary.data_type(), true).is_err()); + + let run_ends = Int16Array::from(vec![1]); + let run = Arc::new(RunArray::::try_new( + &run_ends, + union()?.as_ref(), + )?) as ArrayRef; + assert!(run.is_nullable()); + assert!(normalize_array(&run, run.data_type(), true).is_err()); + Ok(()) + } + fn assert_active_sparse_union_logical_null_is_rejected( active_child: ArrayRef, ) -> Result<()> { @@ -1827,6 +1950,77 @@ mod tests { Ok(()) } + #[test] + fn exact_list_view_discards_unreferenced_backing() -> Result<()> { + const VALUE_COUNT: usize = 8_192; + const RETAINED_INDEX: usize = 4_096; + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let values = Arc::new(ListViewArray::new( + Arc::clone(&field), + ScalarBuffer::from(vec![RETAINED_INDEX as i32]), + ScalarBuffer::from(vec![1_i32]), + Arc::new(Int32Array::from_iter_values(0..VALUE_COUNT as i32)), + None, + )) as ArrayRef; + let data_type = values.data_type().clone(); + + let normalized = normalize_array(&values, &data_type, true)?; + let normalized = normalized.as_list_view::(); + assert_eq!(normalized.values().len(), 1); + assert_eq!(normalized.value_offsets(), &[0]); + assert_eq!(normalized.value_sizes(), &[1]); + assert_eq!( + normalized.values().as_primitive::().value(0), + RETAINED_INDEX as i32 + ); + Ok(()) + } + + #[test] + fn nullable_nested_sparse_union_scrubs_inactive_payload() -> Result<()> { + let union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("integer", DataType::Int32, true), + Field::new("string", DataType::Utf8, true), + ], + )?; + let union = Arc::new(UnionArray::try_new( + union_fields, + ScalarBuffer::from(vec![0_i8, 1, 0]), + None, + vec![ + Arc::new(Int32Array::from(vec![Some(1), Some(999), Some(2)])), + Arc::new(StringArray::from(vec![ + Some("inactive-secret-before"), + Some("visible"), + Some("inactive-secret-after"), + ])), + ], + )?) as ArrayRef; + let fields = Fields::from(vec![Field::new( + "optional", + union.data_type().clone(), + true, + )]); + let value = + Arc::new(StructArray::try_new(fields, vec![union], None)?) as ArrayRef; + let data_type = value.data_type().clone(); + + let normalized = normalize_array(&value, &data_type, true)?; + let union = normalized + .as_struct() + .column(0) + .as_any() + .downcast_ref::() + .expect("expected union"); + let strings = union.child(1).as_string::(); + assert_eq!(strings.value(0), "visible"); + assert_eq!(strings.value(1), "visible"); + assert_eq!(strings.value(2), "visible"); + Ok(()) + } + #[tokio::test] async fn collect_list_dictionary_sql() -> Result<()> { let keys = Int8Array::from(vec![0, 0]); From bfa35cc71b32d30044aa64cf9cc2baaabfee55f3 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 13:42:43 -0700 Subject: [PATCH 11/13] fix: compact nested collect wrappers Signed-off-by: goutamadwant --- .../spark/src/function/aggregate/collect.rs | 299 +++++++++++++++++- 1 file changed, 294 insertions(+), 5 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 5a732d4206e52..ef088aff3ecf5 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -16,9 +16,9 @@ // under the License. use arrow::array::{ - Array, ArrayData, ArrayRef, GenericListArray, GenericListViewArray, OffsetSizeTrait, - PrimitiveArray, RunArray, StructArray, UInt64Array, UnionArray, cast::AsArray, - downcast_run_array, make_array, + Array, ArrayData, ArrayRef, FixedSizeListArray, GenericListArray, + GenericListViewArray, MapArray, OffsetSizeTrait, PrimitiveArray, RunArray, + StructArray, UInt64Array, UnionArray, cast::AsArray, downcast_run_array, make_array, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::compute::{cast, take}; @@ -95,11 +95,35 @@ fn has_compactable_storage(value: &ArrayRef) -> bool { .is_some_and(|offset| *offset as usize != list.values().len()) || has_compactable_storage(list.values()) } + DataType::FixedSizeList(_, size) => { + let list = value.as_fixed_size_list(); + let has_exact_child_length = usize::try_from(*size) + .ok() + .and_then(|size| list.len().checked_mul(size)) + .is_some_and(|length| length == list.values().len()); + list.null_count() != 0 + || list.values().offset() != 0 + || !has_exact_child_length + || has_compactable_storage(list.values()) + } DataType::Struct(_) => value .as_struct() .columns() .iter() .any(has_compactable_storage), + DataType::Map(_, _) => { + let map = value.as_map(); + map.null_count() != 0 + || map + .value_offsets() + .first() + .is_some_and(|offset| *offset != 0) + || map + .value_offsets() + .last() + .is_some_and(|offset| *offset as usize != map.entries().len()) + || map.entries().columns().iter().any(has_compactable_storage) + } _ => false, } } @@ -133,6 +157,12 @@ fn normalize_array( (DataType::LargeListView(_), DataType::LargeListView(field)) => { normalize_list_view::(value, field) } + ( + DataType::FixedSizeList(_, source_size), + DataType::FixedSizeList(field, target_size), + ) if source_size == target_size => { + normalize_fixed_size_list(value, field, *target_size) + } ( DataType::Dictionary(source_key, _), DataType::Dictionary(target_key, target_value), @@ -190,6 +220,9 @@ fn normalize_array( source.nulls().cloned(), )?)) } + (DataType::Map(_, _), DataType::Map(field, ordered)) => { + normalize_map(value, field, *ordered) + } ( DataType::Union(source_fields, UnionMode::Sparse), DataType::Union(fields, UnionMode::Sparse), @@ -435,6 +468,139 @@ fn normalize_list_view( )?)) } +fn normalize_fixed_size_list( + value: &ArrayRef, + field: &FieldRef, + size: i32, +) -> Result { + let source = value.as_fixed_size_list(); + let size = usize::try_from(size).map_err(|_| { + internal_datafusion_err!("fixed-size-list size cannot be negative") + })?; + let values_len = source.len().checked_mul(size).ok_or_else(|| { + internal_datafusion_err!("fixed-size-list child length overflow") + })?; + if source.values().len() != values_len { + return internal_err!( + "fixed-size-list has {} child values, expected {values_len}", + source.values().len() + ); + } + + let values = if source.null_count() == 0 && source.values().offset() == 0 { + Arc::clone(source.values()) + } else if let Some(fallback_row) = + (0..source.len()).find(|index| source.is_valid(*index)) + { + let mut indices = Vec::new(); + indices.try_reserve(values_len).map_err(|error| { + internal_datafusion_err!( + "failed to reserve fixed-size-list child indices: {error}" + ) + })?; + for row in 0..source.len() { + let source_row = if source.is_valid(row) { + row + } else { + fallback_row + }; + let start = source_row.checked_mul(size).ok_or_else(|| { + internal_datafusion_err!("fixed-size-list child offset overflow") + })?; + for offset in 0..size { + let index = start.checked_add(offset).ok_or_else(|| { + internal_datafusion_err!("fixed-size-list child offset overflow") + })?; + indices.push(u64::try_from(index).map_err(|_| { + internal_datafusion_err!("fixed-size-list child offset exceeds u64") + })?); + } + } + take( + source.values().as_ref(), + &UInt64Array::from_iter_values(indices), + None, + )? + } else { + ScalarValue::new_default(field.data_type())?.to_array_of_size(values_len)? + }; + let values = normalize_array(&values, field.data_type(), !field.is_nullable())?; + Ok(Arc::new(FixedSizeListArray::try_new_with_length( + Arc::clone(field), + i32::try_from(size) + .map_err(|_| internal_datafusion_err!("fixed-size-list size exceeds i32"))?, + values, + source.nulls().cloned(), + source.len(), + )?)) +} + +fn normalize_map(value: &ArrayRef, field: &FieldRef, ordered: bool) -> Result { + let source = value.as_map(); + let source_offsets = source.value_offsets(); + let requires_entry_compaction = source.null_count() != 0 + || source_offsets.first().is_some_and(|offset| *offset != 0) + || source_offsets + .last() + .is_some_and(|offset| *offset as usize != source.entries().len()); + + let (offsets, entries) = if requires_entry_compaction { + let retained_len = (0..source.len()) + .filter(|index| source.is_valid(*index)) + .try_fold(0_usize, |retained, index| { + let start = source_offsets[index] as usize; + let end = source_offsets[index + 1] as usize; + retained.checked_add(end - start).ok_or_else(|| { + internal_datafusion_err!("compacted map length overflow") + }) + })?; + let mut indices = Vec::new(); + indices.try_reserve(retained_len).map_err(|error| { + internal_datafusion_err!( + "failed to reserve compacted map entry indices: {error}" + ) + })?; + let mut offsets = Vec::new(); + let offsets_len = source.len().checked_add(1).ok_or_else(|| { + internal_datafusion_err!("compacted map offset length overflow") + })?; + offsets.try_reserve(offsets_len).map_err(|error| { + internal_datafusion_err!("failed to reserve compacted map offsets: {error}") + })?; + offsets.push(0_i32); + for index in 0..source.len() { + if source.is_valid(index) { + let start = source_offsets[index] as usize; + let end = source_offsets[index + 1] as usize; + indices.extend((start..end).map(|index| index as u64)); + } + offsets.push(i32::try_from(indices.len()).map_err(|_| { + internal_datafusion_err!("compacted map offset exceeds i32") + })?); + } + let entries = take( + source.entries(), + &UInt64Array::from_iter_values(indices), + None, + )?; + (OffsetBuffer::new(ScalarBuffer::from(offsets)), entries) + } else { + ( + source.offsets().clone(), + Arc::new(source.entries().clone()) as ArrayRef, + ) + }; + + let entries = normalize_array(&entries, field.data_type(), true)?; + Ok(Arc::new(MapArray::try_new( + Arc::clone(field), + offsets, + entries.as_struct().clone(), + source.nulls().cloned(), + ordered, + )?)) +} + fn normalize_sparse_union( value: &ArrayRef, source_fields: &arrow::datatypes::UnionFields, @@ -905,8 +1071,9 @@ impl Accumulator for NullToEmptyListAccumulator { mod tests { use super::*; use arrow::array::{ - DictionaryArray, Int8Array, Int16Array, Int32Array, Int64Array, ListArray, - ListViewArray, NullArray, RunArray, StringArray, StructArray, UnionArray, + DictionaryArray, FixedSizeListArray, Int8Array, Int16Array, Int32Array, + Int64Array, ListArray, ListViewArray, MapArray, NullArray, RunArray, StringArray, + StructArray, UnionArray, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ @@ -916,6 +1083,7 @@ mod tests { use arrow::util::display::array_value_to_string; use datafusion::prelude::SessionContext; use datafusion_expr::AggregateUDF; + use std::collections::HashMap; fn list_type(element_type: DataType) -> DataType { DataType::List(Arc::new(Field::new_list_field(element_type, false))) @@ -2058,6 +2226,127 @@ mod tests { Ok(()) } + #[tokio::test] + async fn collect_list_map_discards_unreferenced_entries_sql() -> Result<()> { + let key_field = Arc::new(Field::new("key", DataType::Int32, false)); + let value_field = Arc::new(Field::new("value", DataType::Utf8, false)); + let entries = StructArray::try_new( + Fields::from(vec![Arc::clone(&key_field), Arc::clone(&value_field)]), + vec![ + Arc::new(Int32Array::from(vec![0, 1, 2])), + Arc::new(StringArray::from(vec![ + "secret-before", + "visible", + "secret-after", + ])), + ], + None, + )?; + let entry_field = Arc::new( + Field::new("entries", entries.data_type().clone(), false).with_metadata( + HashMap::from([("test-metadata".to_string(), "preserved".to_string())]), + ), + ); + let values = Arc::new(MapArray::try_new( + Arc::clone(&entry_field), + OffsetBuffer::new(ScalarBuffer::from(vec![1_i32, 2_i32])), + entries, + None, + true, + )?) as ArrayRef; + let element_type = values.data_type().clone(); + + let ctx = SessionContext::new(); + ctx.register_udaf(AggregateUDF::new_from_impl(SparkCollectList::new())); + ctx.register_batch( + "map_input", + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "x", + element_type.clone(), + false, + )])), + vec![values], + )?, + )?; + + let batches = ctx + .sql("SELECT collect_list(x) AS values FROM map_input") + .await? + .collect() + .await?; + let list = batches[0].column(0).as_list::(); + let values = list.value(0); + let map = values.as_map(); + assert_eq!(map.data_type(), &element_type); + assert_eq!(map.value_offsets(), &[0, 1]); + assert_eq!(map.entries().len(), 1); + assert_eq!( + map.values().as_string::().iter().collect::>(), + vec![Some("visible")] + ); + Ok(()) + } + + #[tokio::test] + async fn collect_list_sliced_fixed_size_list_compacts_nested_dictionary_sql() + -> Result<()> { + let dictionary = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0, 1, 2]), + Arc::new(StringArray::from(vec![ + "secret-before", + "visible", + "secret-after", + ])), + )?) as ArrayRef; + let item_field = Arc::new( + Field::new_list_field(dictionary.data_type().clone(), false).with_metadata( + HashMap::from([("test-metadata".to_string(), "preserved".to_string())]), + ), + ); + let values = + FixedSizeListArray::try_new(Arc::clone(&item_field), 1, dictionary, None)? + .slice(1, 1); + assert_eq!(values.values().len(), 1); + let values = Arc::new(values) as ArrayRef; + let element_type = values.data_type().clone(); + + let ctx = SessionContext::new(); + ctx.register_udaf(AggregateUDF::new_from_impl(SparkCollectList::new())); + ctx.register_batch( + "fixed_size_list_input", + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "x", + element_type.clone(), + false, + )])), + vec![values], + )?, + )?; + + let batches = ctx + .sql("SELECT collect_list(x) AS values FROM fixed_size_list_input") + .await? + .collect() + .await?; + let list = batches[0].column(0).as_list::(); + let values = list.value(0); + let fixed = values.as_fixed_size_list(); + assert_eq!(fixed.data_type(), &element_type); + assert_eq!(fixed.values().len(), 1); + let dictionary = fixed + .values() + .as_any() + .downcast_ref::>() + .expect("expected Int8 dictionary"); + assert_eq!(dictionary.keys().offset(), 0); + let dictionary = dictionary.values().as_string::(); + assert_eq!(dictionary.len(), 1); + assert_eq!(dictionary.value(0), "visible"); + Ok(()) + } + #[tokio::test] async fn collect_list_dictionary_grouped_and_window_sql() -> Result<()> { let keys = Int8Array::new( From d65cf5f3c23ba22a1195f2db2bab93ce3d3fa2a1 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 14:12:17 -0700 Subject: [PATCH 12/13] fix: preserve fixed-size-list default shape Signed-off-by: goutamadwant --- datafusion/common/src/scalar/mod.rs | 23 +++++--- .../spark/src/function/aggregate/collect.rs | 55 +++++++++++++++++++ 2 files changed, 70 insertions(+), 8 deletions(-) diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index 40ba1c2b1e32e..7d49307a5a6da 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -1724,14 +1724,21 @@ impl ScalarValue { ScalarValue::new_list(&[], field.data_type(), field.is_nullable()); Ok(ScalarValue::List(list)) } - DataType::FixedSizeList(field, _size) => { - let empty_arr = new_empty_array(field.data_type()); - let values = Arc::new( - SingleRowListArrayBuilder::new(empty_arr) - .with_field(field) - .build_fixed_size_list_array(0), - ); - Ok(ScalarValue::FixedSizeList(values)) + DataType::FixedSizeList(field, size) => { + let list_size = size.to_usize().ok_or_else(|| { + _internal_datafusion_err!("FixedSizeList size cannot be negative") + })?; + let values = ScalarValue::new_default(field.data_type())? + .to_array_of_size(list_size)?; + Ok(ScalarValue::FixedSizeList(Arc::new( + FixedSizeListArray::try_new_with_length( + Arc::clone(field), + *size, + values, + None, + 1, + )?, + ))) } DataType::LargeList(field) => { let list = ScalarValue::new_large_list(&[], field.data_type()); diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index ef088aff3ecf5..01d7483fa44d0 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -2347,6 +2347,61 @@ mod tests { Ok(()) } + #[tokio::test] + async fn collect_list_preserves_null_nested_fixed_size_list_sql() -> Result<()> { + let inner_field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let inner = Arc::new(FixedSizeListArray::try_new( + inner_field, + 2, + Arc::new(Int32Array::from(vec![1, 2])), + None, + )?) as ArrayRef; + let outer_field = + Arc::new(Field::new_list_field(inner.data_type().clone(), false)); + let outer = Arc::new(FixedSizeListArray::try_new( + outer_field, + 1, + inner, + Some(NullBuffer::new_null(1)), + )?) as ArrayRef; + let element_fields = Fields::from(vec![Field::new( + "optional", + outer.data_type().clone(), + true, + )]); + let element_type = DataType::Struct(element_fields.clone()); + let values = Arc::new(StructArray::try_new(element_fields, vec![outer], None)?) + as ArrayRef; + values.to_data().validate_full()?; + + let ctx = SessionContext::new(); + ctx.register_udaf(AggregateUDF::new_from_impl(SparkCollectList::new())); + ctx.register_batch( + "nested_fixed_size_list_input", + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "x", + element_type.clone(), + false, + )])), + vec![values], + )?, + )?; + + let batches = ctx + .sql("SELECT collect_list(x) AS values FROM nested_fixed_size_list_input") + .await? + .collect() + .await?; + let list = batches[0].column(0).as_list::(); + assert_eq!(list.data_type(), &list_type(element_type)); + let collected = list.value(0); + let outer = collected.as_struct().column(0).as_fixed_size_list(); + assert!(outer.is_null(0)); + outer.to_data().validate_full()?; + Ok(()) + } + #[tokio::test] async fn collect_list_dictionary_grouped_and_window_sql() -> Result<()> { let keys = Int8Array::new( From f83f0f4a9c1aa3519a048da29a873aae6befc1b1 Mon Sep 17 00:00:00 2001 From: goutamadwant Date: Sat, 19 Sep 2026 14:28:48 -0700 Subject: [PATCH 13/13] fix: harden Spark collect normalization Signed-off-by: goutamadwant --- .../spark/src/function/aggregate/collect.rs | 166 ++++++++++++++---- 1 file changed, 131 insertions(+), 35 deletions(-) diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 01d7483fa44d0..4e35936f78e9d 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -106,11 +106,11 @@ fn has_compactable_storage(value: &ArrayRef) -> bool { || !has_exact_child_length || has_compactable_storage(list.values()) } - DataType::Struct(_) => value - .as_struct() - .columns() - .iter() - .any(has_compactable_storage), + DataType::Struct(_) => { + let structure = value.as_struct(); + structure.null_count() != 0 + || structure.columns().iter().any(has_compactable_storage) + } DataType::Map(_, _) => { let map = value.as_map(); map.null_count() != 0 @@ -146,22 +146,24 @@ fn normalize_array( { return Ok(Arc::clone(value)); } - match (value.data_type(), target_type) { - (DataType::List(_), DataType::List(field)) => normalize_list::(value, field), + let normalized = match (value.data_type(), target_type) { + (DataType::List(_), DataType::List(field)) => { + normalize_list::(value, field)? + } (DataType::LargeList(_), DataType::LargeList(field)) => { - normalize_list::(value, field) + normalize_list::(value, field)? } (DataType::ListView(_), DataType::ListView(field)) => { - normalize_list_view::(value, field) + normalize_list_view::(value, field)? } (DataType::LargeListView(_), DataType::LargeListView(field)) => { - normalize_list_view::(value, field) + normalize_list_view::(value, field)? } ( DataType::FixedSizeList(_, source_size), DataType::FixedSizeList(field, target_size), ) if source_size == target_size => { - normalize_fixed_size_list(value, field, *target_size) + normalize_fixed_size_list(value, field, *target_size)? } ( DataType::Dictionary(source_key, _), @@ -174,9 +176,9 @@ fn normalize_array( let dictionary = dictionary.with_values(values); if dictionary.null_count() == 0 && dictionary.to_data().nulls().is_some() { let data = dictionary.to_data().into_builder().nulls(None).build()?; - Ok(make_array(data)) + make_array(data) } else { - Ok(dictionary) + dictionary } } ( @@ -194,7 +196,7 @@ fn normalize_array( require_non_null, ), _ => internal_err!("collect_list/collect_set expected a run-end encoded array"), - } + }? } (DataType::Struct(_), DataType::Struct(fields)) => { let source = value.as_struct(); @@ -202,35 +204,32 @@ fn normalize_array( .iter() .zip(source.columns()) .map(|(field, column)| { - let column = if !field.is_nullable() { - materialize_masked_values( - column, - source.nulls(), - field.data_type(), - )? - } else { - Arc::clone(column) - }; + let column = materialize_masked_values( + column, + source.nulls(), + field.data_type(), + !field.is_nullable(), + )?; normalize_array(&column, field.data_type(), !field.is_nullable()) }) .collect::>>()?; - Ok(Arc::new(StructArray::try_new( + Arc::new(StructArray::try_new( fields.clone(), columns, source.nulls().cloned(), - )?)) + )?) as ArrayRef } (DataType::Map(_, _), DataType::Map(field, ordered)) => { - normalize_map(value, field, *ordered) + normalize_map(value, field, *ordered)? } ( DataType::Union(source_fields, UnionMode::Sparse), DataType::Union(fields, UnionMode::Sparse), - ) => normalize_sparse_union(value, source_fields, fields, require_non_null), + ) => normalize_sparse_union(value, source_fields, fields, require_non_null)?, ( DataType::Union(source_fields, UnionMode::Dense), DataType::Union(fields, UnionMode::Dense), - ) => normalize_dense_union(value, source_fields, fields, require_non_null), + ) => normalize_dense_union(value, source_fields, fields, require_non_null)?, _ => { let value = if value.data_type() == target_type { Arc::clone(value) @@ -243,18 +242,20 @@ fn normalize_array( && value.to_data().nulls().is_some() { let data = value.to_data().into_builder().nulls(None).build()?; - Ok(make_array(data)) + make_array(data) } else { - Ok(value) + value } } - } + }; + Ok(normalized) } fn materialize_masked_values( value: &ArrayRef, parent_nulls: Option<&arrow::buffer::NullBuffer>, target_type: &DataType, + require_non_null: bool, ) -> Result { let Some(parent_nulls) = parent_nulls else { return Ok(Arc::clone(value)); @@ -263,9 +264,13 @@ fn materialize_masked_values( let Some(first_valid_parent) = valid_parent_indices.next() else { return ScalarValue::new_default(target_type)?.to_array_of_size(value.len()); }; - let fallback = std::iter::once(first_valid_parent) - .chain(valid_parent_indices) - .find(|index| value.is_valid(*index)); + let fallback = if require_non_null { + std::iter::once(first_valid_parent) + .chain(valid_parent_indices) + .find(|index| value.is_valid(*index)) + } else { + Some(first_valid_parent) + }; let Some(fallback) = fallback else { // The nulls are logically visible because the parent rows are valid. // Preserve them so the declared non-null field validation rejects the @@ -843,6 +848,11 @@ fn normalize_list_scalar( } let values = normalize_array(&array.value(0), field.data_type(), true)?; + if values.logical_null_count() != 0 { + return internal_err!( + "Found unmasked nulls for non-nullable collect_list/collect_set element" + ); + } Ok(SingleRowListArrayBuilder::new(values) .with_field(field) .build_list_scalar()) @@ -992,7 +1002,7 @@ impl NullToEmptyListAccumulator { self.list_type ); }; - if value.data_type() == field.data_type() { + if value.data_type() == field.data_type() && !has_compactable_storage(value) { Ok(Arc::clone(value)) } else if matches!( (value.data_type(), field.data_type()), @@ -1559,6 +1569,52 @@ mod tests { Ok(()) } + #[test] + fn exact_type_input_discards_unreferenced_dictionary_values() -> Result<()> { + let values = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![1]), + Arc::new(StringArray::from(vec![ + "secret-before", + "visible", + "secret-after", + ])), + )?) as ArrayRef; + let element_type = values.data_type().clone(); + let accumulator = NullToEmptyListAccumulator::new( + ArrayAggAccumulator::try_new(&element_type, true)?, + list_type(element_type.clone()), + ); + + let normalized = accumulator.normalize_input(&values)?; + assert_eq!(normalized.data_type(), &element_type); + let dictionary = normalized + .as_any() + .downcast_ref::>() + .expect("expected Int8 dictionary"); + assert_eq!(dictionary.values().len(), 1); + assert_eq!(dictionary.values().as_string::().value(0), "visible"); + Ok(()) + } + + #[test] + fn required_encoded_logical_null_returns_error() -> Result<()> { + let values = Arc::new(DictionaryArray::::try_new( + Int8Array::from(vec![0]), + Arc::new(StringArray::from(vec![None::<&str>])), + )?) as ArrayRef; + let element_type = values.data_type().clone(); + let scalar = SingleRowListArrayBuilder::new(values).build_list_scalar(); + + let error = normalize_list_scalar(scalar, &list_type(element_type)).unwrap_err(); + assert!( + error + .to_string() + .contains("Found unmasked nulls for non-nullable"), + "unexpected error: {error}" + ); + Ok(()) + } + #[test] fn nullable_nested_dictionary_preserves_logical_null_union_value() -> Result<()> { let runtime_union_fields = UnionFields::try_new( @@ -2020,6 +2076,16 @@ mod tests { let normalized = normalized.as_run::(); assert_eq!(normalized.values().len(), 1); assert_eq!(normalized.values().data_type(), &DataType::LargeUtf8); + + let exact_type = values.data_type().clone(); + let accumulator = NullToEmptyListAccumulator::new( + ArrayAggAccumulator::try_new(&exact_type, true)?, + list_type(exact_type.clone()), + ); + let normalized = accumulator.normalize_input(&values)?; + assert_eq!(normalized.data_type(), &exact_type); + assert_eq!(normalized.len(), logical_len as usize); + assert_eq!(normalized.as_run::().values().len(), 1); Ok(()) } @@ -2189,6 +2255,36 @@ mod tests { Ok(()) } + #[test] + fn nullable_struct_child_scrubs_parent_masked_payload() -> Result<()> { + let inner_fields = + Fields::from(vec![Field::new("optional_value", DataType::Utf8, true)]); + let inner = Arc::new(StructArray::try_new( + inner_fields, + vec![Arc::new(StringArray::from(vec![ + "masked-secret", + "visible", + ]))], + Some(NullBuffer::from(vec![false, true])), + )?) as ArrayRef; + let outer_fields = Fields::from(vec![Field::new( + "optional_struct", + inner.data_type().clone(), + true, + )]); + let value = + Arc::new(StructArray::try_new(outer_fields, vec![inner], None)?) as ArrayRef; + let data_type = value.data_type().clone(); + + let normalized = normalize_array(&value, &data_type, true)?; + let structure = normalized.as_struct().column(0).as_struct(); + assert!(structure.is_null(0)); + let child = structure.column(0).as_string::(); + assert_eq!(child.value(0), "visible"); + assert_eq!(child.value(1), "visible"); + Ok(()) + } + #[tokio::test] async fn collect_list_dictionary_sql() -> Result<()> { let keys = Int8Array::from(vec![0, 0]);