feat: support map_entries builtin function (#16557)

Co-authored-by: Matt Butrovich <[email protected]>

Author: comphead

datafusion/common/src/utils/mod.rs

@@ -950,7 +950,7 @@ pub fn get_available_parallelism() -> usize {
         .get()
 }
 
-/// Converts a collection of function arguments into an fixed-size array of length N
+/// Converts a collection of function arguments into a fixed-size array of length N
 /// producing a reasonable error message in case of unexpected number of arguments.
 ///
 /// # Example

datafusion/functions-nested/src/lib.rs

@@ -50,6 +50,7 @@ pub mod flatten;
 pub mod length;
 pub mod make_array;
 pub mod map;
+pub mod map_entries;
 pub mod map_extract;
 pub mod map_keys;
 pub mod map_values;
@@ -95,6 +96,7 @@ pub mod expr_fn {
     pub use super::flatten::flatten;
     pub use super::length::array_length;
     pub use super::make_array::make_array;
+    pub use super::map_entries::map_entries;
     pub use super::map_extract::map_extract;
     pub use super::map_keys::map_keys;
     pub use super::map_values::map_values;
@@ -163,6 +165,7 @@ pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
         replace::array_replace_all_udf(),
         replace::array_replace_udf(),
         map::map_udf(),
+        map_entries::map_entries_udf(),
         map_extract::map_extract_udf(),
         map_keys::map_keys_udf(),
         map_values::map_values_udf(),

datafusion/functions-nested/src/map_entries.rs

@@ -0,0 +1,146 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`ScalarUDFImpl`] definitions for map_entries function.
+
+use crate::utils::{get_map_entry_field, make_scalar_function};
+use arrow::array::{Array, ArrayRef, ListArray};
+use arrow::datatypes::{DataType, Field, Fields};
+use datafusion_common::utils::take_function_args;
+use datafusion_common::{cast::as_map_array, exec_err, Result};
+use datafusion_expr::{
+    ArrayFunctionSignature, ColumnarValue, Documentation, ScalarUDFImpl, Signature,
+    TypeSignature, Volatility,
+};
+use datafusion_macros::user_doc;
+use std::any::Any;
+use std::sync::Arc;
+
+make_udf_expr_and_func!(
+    MapEntriesFunc,
+    map_entries,
+    map,
+    "Return a list of all entries in the map.",
+    map_entries_udf
+);
+
+#[user_doc(
+    doc_section(label = "Map Functions"),
+    description = "Returns a list of all entries in the map.",
+    syntax_example = "map_entries(map)",
+    sql_example = r#"```sql
+SELECT map_entries(MAP {'a': 1, 'b': NULL, 'c': 3});
+----
+[{'key': a, 'value': 1}, {'key': b, 'value': NULL}, {'key': c, 'value': 3}]
+
+SELECT map_entries(map([100, 5], [42, 43]));
+----
+[{'key': 100, 'value': 42}, {'key': 5, 'value': 43}]
+```"#,
+    argument(
+        name = "map",
+        description = "Map expression. Can be a constant, column, or function, and any combination of map operators."
+    )
+)]
+#[derive(Debug)]
+pub struct MapEntriesFunc {
+    signature: Signature,
+}
+
+impl Default for MapEntriesFunc {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl MapEntriesFunc {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::new(
+                TypeSignature::ArraySignature(ArrayFunctionSignature::MapArray),
+                Volatility::Immutable,
+            ),
+        }
+    }
+}
+
+impl ScalarUDFImpl for MapEntriesFunc {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "map_entries"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        let [map_type] = take_function_args(self.name(), arg_types)?;
+        let map_fields = get_map_entry_field(map_type)?;
+        Ok(DataType::List(Arc::new(Field::new_list_field(
+            DataType::Struct(Fields::from(vec![
+                Field::new(
+                    "key",
+                    map_fields.first().unwrap().data_type().clone(),
+                    false,
+                ),
+                Field::new(
+                    "value",
+                    map_fields.get(1).unwrap().data_type().clone(),
+                    true,
+                ),
+            ])),
+            false,
+        ))))
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: datafusion_expr::ScalarFunctionArgs,
+    ) -> Result<ColumnarValue> {
+        make_scalar_function(map_entries_inner)(&args.args)
+    }
+
+    fn documentation(&self) -> Option<&Documentation> {
+        self.doc()
+    }
+}
+
+fn map_entries_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let [map_arg] = take_function_args("map_entries", args)?;
+
+    let map_array = match map_arg.data_type() {
+        DataType::Map(_, _) => as_map_array(&map_arg)?,
+        _ => return exec_err!("Argument for map_entries should be a map"),
+    };
+
+    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,
+        )),
+        map_array.offsets().clone(),
+        Arc::new(map_array.entries().clone()),
+        map_array.nulls().cloned(),
+    )))
+}

datafusion/sqllogictest/test_files/map.slt

@@ -651,6 +651,57 @@ select map_extract(column1, 1), map_extract(column1, 5), map_extract(column1, 7)
 [NULL] [[4, NULL, 6]] [NULL]
 [NULL] [NULL] [[1, NULL, 3]]
 
+# Tests for map_entries
+
+query ?
+SELECT map_entries(MAP { 'a': 1, 'b': 3 });
+----
+[{key: a, value: 1}, {key: b, value: 3}]
+
+query error DataFusion error: Arrow error: Cast error: Cannot cast string 'a' to value of Int64 type
+SELECT map_entries(MAP { 'a': 1, 2: 3 });
+
+query ?
+SELECT map_entries(MAP {'a':1, 'b':2, 'c':3 }) FROM t;
+----
+[{key: a, value: 1}, {key: b, value: 2}, {key: c, value: 3}]
+[{key: a, value: 1}, {key: b, value: 2}, {key: c, value: 3}]
+[{key: a, value: 1}, {key: b, value: 2}, {key: c, value: 3}]
+
+query ?
+SELECT map_entries(Map{column1: column2, column3: column4}) FROM t;
+----
+[{key: a, value: 1}, {key: k1, value: 10}]
+[{key: b, value: 2}, {key: k3, value: 30}]
+[{key: d, value: 4}, {key: k5, value: 50}]
+
+query ?
+SELECT map_entries(map(column5, column6)) FROM t;
+----
+[{key: k1, value: 1}, {key: k2, value: 2}]
+[{key: k3, value: 3}]
+[{key: k5, value: 5}]
+
+query ?
+SELECT map_entries(map(column8, column9)) FROM t;
+----
+[{key: [1, 2, 3], value: a}]
+[{key: [4], value: b}]
+[{key: [1, 2], value: c}]
+
+query ?
+SELECT map_entries(Map{});
+----
+[]
+
+query ?
+SELECT map_entries(column1) from map_array_table_1;
+----
+[{key: 1, value: [1, NULL, 3]}, {key: 2, value: [4, NULL, 6]}, {key: 3, value: [7, 8, 9]}]
+[{key: 4, value: [1, NULL, 3]}, {key: 5, value: [4, NULL, 6]}, {key: 6, value: [7, 8, 9]}]
+[{key: 7, value: [1, NULL, 3]}, {key: 8, value: [9, NULL, 6]}, {key: 9, value: [7, 8, 9]}]
+NULL
+
 # Tests for map_keys
 
 query ?

docs/source/user-guide/sql/scalar_functions.md

@@ -4105,6 +4105,7 @@ select struct(a as field_a, b) from t;
 
 - [element_at](#element_at)
 - [map](#map)
+- [map_entries](#map_entries)
 - [map_extract](#map_extract)
 - [map_keys](#map_keys)
 - [map_values](#map_values)
@@ -4162,6 +4163,30 @@ SELECT MAKE_MAP(['key1', 'key2'], ['value1', null]);
 {key1: value1, key2: }
 ```
 
+### `map_entries`
+
+Returns a list of all entries in the map.
+
+```sql
+map_entries(map)
+```
+
+#### Arguments
+
+- **map**: Map expression. Can be a constant, column, or function, and any combination of map operators.
+
+#### Example
+
+```sql
+SELECT map_entries(MAP {'a': 1, 'b': NULL, 'c': 3});
+----
+[{'key': a, 'value': 1}, {'key': b, 'value': NULL}, {'key': c, 'value': 3}]
+
+SELECT map_entries(map([100, 5], [42, 43]));
+----
+[{'key': 100, 'value': 42}, {'key': 5, 'value': 43}]
+```
+
 ### `map_extract`
 
 Returns a list containing the value for the given key or an empty list if the key is not present in the map.