diff --git a/datafusion/functions-nested/src/map_entries.rs b/datafusion/functions-nested/src/map_entries.rs index f862d03f62395..7d8ebec7bde38 100644 --- a/datafusion/functions-nested/src/map_entries.rs +++ b/datafusion/functions-nested/src/map_entries.rs @@ -5,7 +5,7 @@ //! [`ScalarUDFImpl`] definitions for map_entries function. use crate::utils::{get_map_entry_field, make_scalar_function}; -use arrow::array::{Array, ArrayRef, ListArray}; +use arrow::array::{Array, ArrayRef, ListArray, StructArray}; use arrow::datatypes::{DataType, Field, Fields}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, cast::as_map_array, exec_err}; @@ -110,16 +110,20 @@ fn map_entries_inner(args: &[ArrayRef]) -> Result { _ => return exec_err!("Argument for map_entries should be a map"), }; + // The input's entry fields can differ from the `key` and `value` fields of + // the result in name, nullability or metadata (for example, a Parquet map + // with non-nullable values), so rebuild the entries with the result fields. + let fields = Fields::from(vec![ + Field::new("key", map_array.key_type().clone(), false), + Field::new("value", map_array.value_type().clone(), true), + ]); + let (_, columns, nulls) = map_array.entries().clone().into_parts(); + let entries = StructArray::try_new(fields.clone(), columns, nulls)?; + Ok(Arc::new(ListArray::new( - Arc::new(Field::new_list_field( - DataType::Struct(Fields::from(vec![ - Field::new("key", map_array.key_type().clone(), false), - Field::new("value", map_array.value_type().clone(), true), - ])), - false, - )), + Arc::new(Field::new_list_field(DataType::Struct(fields), false)), map_array.offsets().clone(), - Arc::new(map_array.entries().clone()), + Arc::new(entries), map_array.nulls().cloned(), ))) } diff --git a/datafusion/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index 85129a9b94c66..f96c1c453a87c 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/datafusion/sqllogictest/test_files/map.slt @@ -80,6 +80,13 @@ SELECT strings['not_found'], arrow_field(strings)['nullable'], arrow_field(strin ---- NULL false true +# map_entries on a map whose values are not nullable +query ? +SELECT map_entries(ints) FROM data ORDER BY ints['bytes'] DESC LIMIT 2; +---- +[{key: bytes, value: 49960}] +[{key: bytes, value: 49689}] + # Select non existent key, expect NULL for each row query I SELECT COUNT(CASE WHEN strings['not_found'] IS NULL THEN 1 ELSE 0 END) FROM data; @@ -853,6 +860,23 @@ SELECT map_entries(column1) from map_array_table_1; [{key: 7, value: [1, NULL, 3]}, {key: 8, value: [9, NULL, 6]}, {key: 9, value: [7, 8, 9]}] NULL +# map_entries on maps whose entry fields differ from the `key` and `value` +# fields of its result: non-nullable values, or other field names +query ? +SELECT map_entries(arrow_cast(MAP {'a': 1, 'b': 2}, 'Map("entries": non-null Struct("key": non-null Utf8, "value": non-null Int64), unsorted)')); +---- +[{key: a, value: 1}, {key: b, value: 2}] + +query ? +SELECT map_entries(arrow_cast(MAP {'a': 1, 'b': NULL}, 'Map("entries": non-null Struct("keys": non-null Utf8, "values": Int64), unsorted)')); +---- +[{key: a, value: 1}, {key: b, value: NULL}] + +query T +SELECT arrow_typeof(map_entries(arrow_cast(MAP {'a': 1}, 'Map("entries": non-null Struct("keys": non-null Utf8, "values": non-null Int64), unsorted)'))); +---- +List(non-null Struct("key": non-null Utf8, "value": Int64)) + # Tests for map_keys query ?