diff --git a/datafusion/functions-nested/src/extract.rs b/datafusion/functions-nested/src/extract.rs index 9d367f0161fdd..8f2ea1f40dcb2 100644 --- a/datafusion/functions-nested/src/extract.rs +++ b/datafusion/functions-nested/src/extract.rs @@ -23,9 +23,8 @@ use arrow::array::{ }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; -use arrow::datatypes::{ - DataType::{FixedSizeList, LargeList, LargeListView, List, ListView, Null}, - Field, +use arrow::datatypes::DataType::{ + FixedSizeList, LargeList, LargeListView, List, ListView, Null, }; use datafusion_common::cast::as_large_list_array; use datafusion_common::cast::as_list_array; @@ -622,9 +621,23 @@ where let values = array.values(); let original_data = values.to_data(); let capacity = Capacities::Array(original_data.len()); + // Carry the input's list field through to the output so that the returned + // type matches the one promised by `return_type` / `return_field_from_args`, + // including the field name, nullability and metadata. + let field = match array.data_type() { + List(field) | LargeList(field) => Arc::clone(field), + other => { + return internal_err!( + "general_array_slice got unexpected data type: {other}" + ); + } + }; + // `use_nulls` is false because we never call `try_extend_nulls`: null rows are + // emitted as empty slices. Arrow still allocates a validity buffer on its own + // if the child array has nulls. let mut mutable = - MutableArrayData::with_capacities(vec![&original_data], true, capacity); + MutableArrayData::with_capacities(vec![&original_data], false, capacity); // We have the slice syntax compatible with DuckDB v0.8.1. // The rule `adjusted_from_index` and `adjusted_to_index` follows the rule of array_slice in duckdb. @@ -638,9 +651,11 @@ where let end = offset_window[1]; let len = end - start; + // The row is null, so its contents are never observed. Emit an empty + // slice rather than a null child element: the input's list field may be + // non-nullable, in which case a null child would be invalid. if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { - mutable.try_extend_nulls(1)?; - offsets.push(offsets[row_index] + O::usize_as(1)); + offsets.push(offsets[row_index]); continue; } @@ -682,7 +697,7 @@ where let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(array.value_type(), true)), + field, OffsetBuffer::::new(offsets.into()), arrow::array::make_array(data), nulls, @@ -704,12 +719,15 @@ where let field = match array.data_type() { ListView(field) | LargeListView(field) => Arc::clone(field), other => { - return internal_err!("array_slice got unexpected data type: {}", other); + return internal_err!( + "general_list_view_array_slice got unexpected data type: {other}" + ); } }; + // See the note on `use_nulls` in `general_array_slice`. let mut mutable = - MutableArrayData::with_capacities(vec![&original_data], true, capacity); + MutableArrayData::with_capacities(vec![&original_data], false, capacity); // We must build `offsets` and `sizes` buffers manually as ListView does not enforce // monotonically increasing offsets. diff --git a/datafusion/sqllogictest/test_files/array/array_pop.slt b/datafusion/sqllogictest/test_files/array/array_pop.slt index a72e566b9e7ab..0b7ebf75f4b2b 100644 --- a/datafusion/sqllogictest/test_files/array/array_pop.slt +++ b/datafusion/sqllogictest/test_files/array/array_pop.slt @@ -318,5 +318,38 @@ select array_pop_front(arrow_cast([1, 2], 'LargeListView(Int64)')); ---- [2] +# maintains inner nullability +query ??TT +select + array_pop_front(column1), + array_pop_back(column1), + arrow_typeof(array_pop_front(column1)), + arrow_typeof(array_pop_back(column1)) +from values + (arrow_cast([], 'List(non-null Int32)')), + (arrow_cast(NULL, 'List(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'List(non-null Int32)')) +; +---- +[] [] List(non-null Int32) List(non-null Int32) +NULL NULL List(non-null Int32) List(non-null Int32) +[3, 5, -5] [1, 3, 5] List(non-null Int32) List(non-null Int32) + +query ??TT +select + array_pop_front(column1), + array_pop_back(column1), + arrow_typeof(array_pop_front(column1)), + arrow_typeof(array_pop_back(column1)) +from values + (arrow_cast([], 'LargeList(non-null Int32)')), + (arrow_cast(NULL, 'LargeList(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'LargeList(non-null Int32)')) +; +---- +[] [] LargeList(non-null Int32) LargeList(non-null Int32) +NULL NULL LargeList(non-null Int32) LargeList(non-null Int32) +[3, 5, -5] [1, 3, 5] LargeList(non-null Int32) LargeList(non-null Int32) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_slice.slt b/datafusion/sqllogictest/test_files/array/array_slice.slt index 14587a50b2266..76b81b28efc58 100644 --- a/datafusion/sqllogictest/test_files/array/array_slice.slt +++ b/datafusion/sqllogictest/test_files/array/array_slice.slt @@ -450,6 +450,37 @@ NULL NULL [1, 3, 5] +# maintains inner nullability +query ?T +select array_slice(column1, 2, 3), arrow_typeof(array_slice(column1, 2, 3)) +from values + (arrow_cast([], 'List(non-null Int32)')), + (arrow_cast(NULL, 'List(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'List(non-null Int32)')) +; +---- +[] List(non-null Int32) +NULL List(non-null Int32) +[3, 5] List(non-null Int32) + +query ?T +select array_slice(column1, 2, 3), arrow_typeof(array_slice(column1, 2, 3)) +from values + (arrow_cast([], 'LargeList(non-null Int32)')), + (arrow_cast(NULL, 'LargeList(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'LargeList(non-null Int32)')) +; +---- +[] LargeList(non-null Int32) +NULL LargeList(non-null Int32) +[3, 5] LargeList(non-null Int32) + +query ?T +select array_slice(column1, 2, 3, 2), arrow_typeof(array_slice(column1, 2, 3, 2)) +from values (arrow_cast([1, 3, 5, -5], 'List(non-null Int32)')); +---- +[3] List(non-null Int32) + # Testing with empty arguments should result in an error query error DataFusion error: Error during planning: 'array_slice' does not support zero arguments select array_slice(); diff --git a/datafusion/sqllogictest/test_files/spark/array/slice.slt b/datafusion/sqllogictest/test_files/spark/array/slice.slt index aaf4aa4909dfd..f6fb431a0769b 100644 --- a/datafusion/sqllogictest/test_files/spark/array/slice.slt +++ b/datafusion/sqllogictest/test_files/spark/array/slice.slt @@ -152,3 +152,9 @@ query ? SELECT slice(make_array(1), 3, 4) ---- [] + +# the inner field name of the input list is preserved +query ?T +SELECT slice(array(1, 2, 3, 4), 2, 2), arrow_typeof(slice(array(1, 2, 3, 4), 2, 2)); +---- +[2, 3] List(Int64, field: 'element')