feat(spark): implement Spark map function map_from_entries

datafusion/spark/src/function/map/map_from_entries.rs

@@ -0,0 +1,110 @@
+// 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.
+
+use std::any::Any;
+
+use crate::function::map::utils::{
+    get_element_type, get_list_offsets, get_list_values,
+    map_from_keys_values_offsets_nulls, map_type_from_key_value_types,
+};
+use arrow::array::{Array, ArrayRef, StructArray};
+use arrow::datatypes::DataType;
+use datafusion_common::utils::take_function_args;
+use datafusion_common::{exec_err, Result};
+use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
+use datafusion_functions::utils::make_scalar_function;
+
+/// Spark-compatible `map_from_entries` expression
+/// <https://spark.apache.org/docs/latest/api/sql/index.html#map_from_entries>
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct MapFromEntries {
+    signature: Signature,
+}
+
+impl Default for MapFromEntries {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl MapFromEntries {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::array(Volatility::Immutable),
+        }
+    }
+}
+
+impl ScalarUDFImpl for MapFromEntries {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "map_from_entries"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        let [entries_type] = take_function_args("map_from_entries", arg_types)?;
+        let entries_element_type = get_element_type(entries_type)?;
+        let (keys_type, values_type) = match entries_element_type {
+            DataType::Struct(fields) if fields.len() == 2 => {
+                Ok((fields[0].data_type(), fields[1].data_type()))
+            }
+            wrong_type => exec_err!(
+                "map_from_entries: expected array<struct<key, value>>, got {:?}",
+                wrong_type
+            ),
+        }?;
+        Ok(map_type_from_key_value_types(keys_type, values_type))
+    }
+
+    fn invoke_with_args(
+        &self,
+        args: datafusion_expr::ScalarFunctionArgs,
+    ) -> Result<ColumnarValue> {
+        make_scalar_function(map_from_entries_inner, vec![])(&args.args)
+    }
+}
+
+fn map_from_entries_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let [entries] = take_function_args("map_from_entries", args)?;
+    let entries_offsets = get_list_offsets(entries)?;
+    let entries_values = get_list_values(entries)?;
+
+    let (flat_keys, flat_values) =
+        match entries_values.as_any().downcast_ref::<StructArray>() {
+            Some(a) => Ok((a.column(0), a.column(1))),
+            None => exec_err!(
+                "map_from_entries: expected array<struct<key, value>>, got {:?}",
+                entries_values.data_type()
+            ),
+        }?;
+
+    map_from_keys_values_offsets_nulls(
+        flat_keys,
+        flat_values,
+        &entries_offsets,
+        &entries_offsets,
+        entries.nulls(),
+        entries.nulls(),
+    )
+}

datafusion/spark/src/function/map/mod.rs

@@ -16,13 +16,15 @@
 // under the License.
 
 pub mod map_from_arrays;
+pub mod map_from_entries;
 mod utils;
 
 use datafusion_expr::ScalarUDF;
 use datafusion_functions::make_udf_function;
 use std::sync::Arc;
 
 make_udf_function!(map_from_arrays::MapFromArrays, map_from_arrays);
+make_udf_function!(map_from_entries::MapFromEntries, map_from_entries);
 
 pub mod expr_fn {
     use datafusion_functions::export_functions;
@@ -32,8 +34,14 @@ pub mod expr_fn {
         "Creates a map from arrays of keys and values.",
         keys values
     ));
+
+    export_functions!((
+        map_from_entries,
+        "Creates a map from array<struct<key, value>>.",
+        arg1
+    ));
 }
 
 pub fn functions() -> Vec<Arc<ScalarUDF>> {
-    vec![map_from_arrays()]
+    vec![map_from_arrays(), map_from_entries()]
 }

datafusion/sqllogictest/test_files/spark/map/map_from_entries.slt

@@ -0,0 +1,119 @@
+# 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.
+
+# Spark doctests
+query ?
+SELECT map_from_entries(array[struct(1, 'a'), struct(2, 'b')]);
+----
+{1: a, 2: b}
+
+query ?
+SELECT map_from_entries(array[struct(1, cast(null as string)), struct(2, 'b')]);
+----
+{1: NULL, 2: b}
+
+query ?
+SELECT map_from_entries(data) 
+from values 
+    (array[struct(1, 'a'), struct(2, 'b')]), 
+    (array[struct(3, 'c')]) 
+as tab(data);
+----
+{1: a, 2: b}
+{3: c}
+
+# Tests with NULL and empty input structarrays
+query ?
+SELECT map_from_entries(data) 
+from values 
+    (cast(array[] as array<struct<int, string>>)), 
+    (cast(NULL as array<struct<int, string>>)) 
+as tab(data);
+----
+{}
+NULL
+
+#Test with multiple rows: good, empty and nullable
+query ?
+SELECT map_from_entries(data) 
+from values 
+    (NULL), 
+    (array[
+        struct(1 as a, 'b' as b), 
+        struct(2 as a, cast(NULL as string) as b), 
+        struct(3 as a, 'd' as b)
+    ]), 
+    (array[]),
+    (NULL) 
+as tab(data);
+----
+NULL
+{1: b, 2: NULL, 3: d}
+{}
+NULL
+
+# Test with complex types
+query ?
+SELECT map_from_entries(array[
+    struct(array('a', 'b'), struct(1, 2, 3)), 
+    struct(array('c', 'd'), struct(4, 5, 6))
+]);
+----
+{[a, b]: {c0: 1, c1: 2, c2: 3}, [c, d]: {c0: 4, c1: 5, c2: 6}}
+
+# Test with nested function calls
+query ?
+SELECT
+    map_from_entries(
+        array[
+            struct(
+                'outer_key1', 
+                -- value for outer_key1: a map itself
+                map_from_entries(
+                    array[
+                        struct('inner_a', 1), 
+                        struct('inner_b', 2)
+                    ]
+                )
+            ),
+            struct(
+                'outer_key2',
+                -- value for outer_key2: another map
+                map_from_entries(
+                    array[
+                        struct('inner_x', 10), 
+                        struct('inner_y', 20), 
+                        struct('inner_z', 30)
+                    ]
+                )
+            )
+        ]
+    ) AS nested_map;
+----
+{outer_key1: {inner_a: 1, inner_b: 2}, outer_key2: {inner_x: 10, inner_y: 20, inner_z: 30}}
+
+# Test with duplicate keys
+query ?
+SELECT map_from_entries(array(
+    struct(true, 'a'), 
+    struct(false, 'b'), 
+    struct(true, 'c'),
+    struct(false, cast(NULL as string)), 
+    struct(true, 'd')
+));
+----
+{false: NULL, true: d}