Support map_keys & map_values for MAP type

datafusion/common/src/utils/mod.rs

@@ -35,7 +35,7 @@ use arrow_array::{
     Array, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait,
     RecordBatchOptions,
 };
-use arrow_schema::{DataType, Fields};
+use arrow_schema::DataType;
 use sqlparser::ast::Ident;
 use sqlparser::dialect::GenericDialect;
 use sqlparser::parser::Parser;
@@ -754,21 +754,6 @@ pub fn combine_limit(
     (combined_skip, combined_fetch)
 }
 
-pub fn get_map_entry_field(data_type: &DataType) -> Result<&Fields> {
-    match data_type {
-        DataType::Map(field, _) => {
-            let field_data_type = field.data_type();
-            match field_data_type {
-                DataType::Struct(fields) => Ok(fields),
-                _ => {
-                    _internal_err!("Expected a Struct type, got {:?}", field_data_type)
-                }
-            }
-        }
-        _ => _internal_err!("Expected a Map type, got {:?}", data_type),
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use crate::ScalarValue::Null;

datafusion/functions-nested/src/lib.rs

@@ -43,6 +43,7 @@ pub mod length;
 pub mod make_array;
 pub mod map;
 pub mod map_extract;
+mod map_keys;
 pub mod planner;
 pub mod position;
 pub mod range;
@@ -83,6 +84,7 @@ pub mod expr_fn {
     pub use super::length::array_length;
     pub use super::make_array::make_array;
     pub use super::map_extract::map_extract;
+    pub use super::map_keys::map_keys;
     pub use super::position::array_position;
     pub use super::position::array_positions;
     pub use super::range::gen_series;
@@ -146,6 +148,7 @@ pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
         replace::array_replace_udf(),
         map::map_udf(),
         map_extract::map_extract_udf(),
+        map_keys::map_keys_udf(),
     ]
 }
 

datafusion/functions-nested/src/map_extract.rs

@@ -24,15 +24,14 @@ use arrow::datatypes::DataType;
 use arrow_array::{Array, MapArray};
 use arrow_buffer::OffsetBuffer;
 use arrow_schema::Field;
-use datafusion_common::utils::get_map_entry_field;
 
 use datafusion_common::{cast::as_map_array, exec_err, Result};
 use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
 use std::any::Any;
 use std::sync::Arc;
 use std::vec;
 
-use crate::utils::make_scalar_function;
+use crate::utils::{get_map_entry_field, make_scalar_function};
 
 // Create static instances of ScalarUDFs for each function
 make_udf_expr_and_func!(

datafusion/functions-nested/src/map_keys.rs

@@ -0,0 +1,106 @@
+// 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_keys function.
+
+use crate::utils::{get_map_entry_field, make_scalar_function};
+use arrow_array::{Array, ArrayRef, ListArray};
+use arrow_schema::{DataType, Field};
+use datafusion_common::{cast::as_map_array, exec_err, Result};
+use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
+use std::any::Any;
+use std::sync::Arc;
+
+make_udf_expr_and_func!(
+    MapKeysFunc,
+    map_keys,
+    map,
+    "Return a list of all keys in the map.",
+    map_keys_udf
+);
+
+#[derive(Debug)]
+pub(crate) struct MapKeysFunc {
+    signature: Signature,
+}
+
+impl MapKeysFunc {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::user_defined(Volatility::Immutable),
+        }
+    }
+}
+
+impl ScalarUDFImpl for MapKeysFunc {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "map_keys"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result<DataType> {
+        if arg_types.len() != 1 {
+            return exec_err!("map_keys expects single argument");
+        }
+        let map_type = &arg_types[0];
+        let map_fields = get_map_entry_field(map_type)?;
+        Ok(DataType::List(Arc::new(Field::new(
+            "item",
+            map_fields.first().unwrap().data_type().clone(),
+            false,
+        ))))
+    }
+
+    fn invoke(&self, args: &[ColumnarValue]) -> datafusion_common::Result<ColumnarValue> {
+        make_scalar_function(map_keys_inner)(args)
+    }
+
+    fn coerce_types(
+        &self,
+        arg_types: &[DataType],
+    ) -> datafusion_common::Result<Vec<DataType>> {
+        if arg_types.len() != 1 {
+            return exec_err!("map_keys expects single argument");
+        }
+        Ok(vec![arg_types[0].clone()])
+    }
+}
+
+fn map_keys_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+    if args.len() != 1 {
+        return exec_err!("map_keys expects single argument");
+    }
+
+    let map_array = match args[0].data_type() {
+        DataType::Map(_, _) => as_map_array(&args[0])?,
+        _ => return exec_err!("Argument for map_extract should be a map"),
+    };
+
+    Ok(Arc::new(ListArray::new(
+        Arc::new(Field::new("item", map_array.key_type().clone(), false)),
+        map_array.offsets().clone(),
+        Arc::clone(map_array.keys()),
+        None,
+    )))
+}

datafusion/functions-nested/src/utils.rs

@@ -26,11 +26,12 @@ use arrow_array::{
     UInt32Array,
 };
 use arrow_buffer::OffsetBuffer;
-use arrow_schema::Field;
+use arrow_schema::{Field, Fields};
 use datafusion_common::cast::{as_large_list_array, as_list_array};
 use datafusion_common::{exec_err, plan_err, Result, ScalarValue};
 
 use core::any::type_name;
+use datafusion_common::error::_internal_err;
 use datafusion_common::DataFusionError;
 use datafusion_expr::{ColumnarValue, ScalarFunctionImplementation};
 
@@ -253,6 +254,21 @@ pub(crate) fn compute_array_dims(
     }
 }
 
+pub(crate) fn get_map_entry_field(data_type: &DataType) -> Result<&Fields> {
+    match data_type {
+        DataType::Map(field, _) => {
+            let field_data_type = field.data_type();
+            match field_data_type {
+                DataType::Struct(fields) => Ok(fields),
+                _ => {
+                    _internal_err!("Expected a Struct type, got {:?}", field_data_type)
+                }
+            }
+        }
+        _ => _internal_err!("Expected a Map type, got {:?}", data_type),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

datafusion/sqllogictest/test_files/map.slt

@@ -568,8 +568,55 @@ select map_extract(column1, 1), map_extract(column1, 5), map_extract(column1, 7)
 [] [[4, , 6]] []
 [] [] [[1, , 3]]
 
+# Tests for map_keys
+
+query ?
+SELECT map_keys(MAP { 'a': 1, 2: 3 });
+----
+[a, 2]
+
+query ?
+SELECT map_keys(MAP {'a':1, 'b':2, 'c':3 }) FROM t;
+----
+[a, b, c]
+[a, b, c]
+[a, b, c]
+
+query ?
+SELECT map_keys(Map{column1: column2, column3: column4}) FROM t;
+----
+[a, k1]
+[b, k3]
+[d, k5]
+
+query ?
+SELECT map_keys(map(column5, column6)) FROM t;
+----
+[k1, k2]
+[k3]
+[k5]
+
+query ?
+SELECT map_keys(map(column8, column9)) FROM t;
+----
+[[1, 2, 3]]
+[[4]]
+[[1, 2]]
+
+query ?
+SELECT map_keys(Map{});
+----
+[]
+
+query ?
+SELECT map_keys(column1) from map_array_table_1;
+----
+[1, 2, 3]
+[4, 5, 6]
+[7, 8, 9]
+
 statement ok
 drop table map_array_table_1;
 
 statement ok
-drop table map_array_table_2;
+drop table map_array_table_2;