feat(spark): implement Spark map function map_from_entries

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

@@ -0,0 +1,133 @@
+// 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, NullBufferBuilder, StructArray};
+use arrow::buffer::NullBuffer;
+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()
+            ),
+        }?;
+
+    let entries_with_nulls = entries_values.nulls().and_then(|entries_inner_nulls| {
+        let mut builder = NullBufferBuilder::new_with_len(0);
+        let mut cur_offset = entries_offsets
+            .first()
+            .map(|offset| *offset as usize)
+            .unwrap_or(0);
+
+        for next_offset in entries_offsets.iter().skip(1) {
+            let num_entries = *next_offset as usize - cur_offset;
+            builder.append(
+                entries_inner_nulls
+                    .slice(cur_offset, num_entries)
+                    .null_count()
+                    == 0,
+            );
+            cur_offset = *next_offset as usize;
+        }
+        builder.finish()
+    });
+
+    let res_nulls = NullBuffer::union(entries.nulls(), entries_with_nulls.as_ref());
+
+    map_from_keys_values_offsets_nulls(
+        flat_keys,
+        flat_values,
+        &entries_offsets,
+        &entries_offsets,
+        None,
+        res_nulls.as_ref(),
+    )
+}

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/spark/src/function/map/utils.rs

@@ -157,8 +157,15 @@ fn map_deduplicate_keys(
     let offsets_len = keys_offsets.len();
     let mut new_offsets = Vec::with_capacity(offsets_len);
 
-    let mut cur_keys_offset = 0;
-    let mut cur_values_offset = 0;
+    let mut cur_keys_offset = keys_offsets
+        .first()
+        .map(|offset| *offset as usize)
+        .unwrap_or(0);
+    let mut cur_values_offset = values_offsets
+        .first()
+        .map(|offset| *offset as usize)
+        .unwrap_or(0);
+
     let mut new_last_offset = 0;
     new_offsets.push(new_last_offset);
 
@@ -176,36 +183,38 @@ fn map_deduplicate_keys(
         let mut keys_mask_one = [false].repeat(num_keys_entries);
         let mut values_mask_one = [false].repeat(num_values_entries);
 
-        if num_keys_entries != num_values_entries {
-            let key_is_valid = keys_nulls.is_none_or(|buf| buf.is_valid(row_idx));
-            let value_is_valid = values_nulls.is_none_or(|buf| buf.is_valid(row_idx));
-            if key_is_valid && value_is_valid {
+        let key_is_valid = keys_nulls.is_none_or(|buf| buf.is_valid(row_idx));
+        let value_is_valid = values_nulls.is_none_or(|buf| buf.is_valid(row_idx));
+
+        if key_is_valid && value_is_valid {
+            if num_keys_entries != num_values_entries {
                 return exec_err!("map_deduplicate_keys: keys and values lists in the same row must have equal lengths");
+            } else if num_keys_entries != 0 {
+                let mut seen_keys = HashSet::new();
+
+                for cur_entry_idx in (0..num_keys_entries).rev() {
+                    let key = ScalarValue::try_from_array(
+                        &flat_keys,
+                        cur_keys_offset + cur_entry_idx,
+                    )?
+                    .compacted();
+                    if seen_keys.contains(&key) {
+                        // TODO: implement configuration and logic for spark.sql.mapKeyDedupPolicy=EXCEPTION (this is default spark-config)
+                        // exec_err!("invalid argument: duplicate keys in map")
+                        // https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961
+                    } else {
+                        // This code implements deduplication logic for spark.sql.mapKeyDedupPolicy=LAST_WIN (this is NOT default spark-config)
+                        keys_mask_one[cur_entry_idx] = true;
+                        values_mask_one[cur_entry_idx] = true;
+                        seen_keys.insert(key);
+                        new_last_offset += 1;
+                    }
+                }
             }
-            // else the result entry is NULL
+        } else {
+            // the result entry is NULL
             // both current row offsets are skipped
             // keys or values in the current row are marked false in the masks
-        } else if num_keys_entries != 0 {
-            let mut seen_keys = HashSet::new();
-
-            for cur_entry_idx in (0..num_keys_entries).rev() {
-                let key = ScalarValue::try_from_array(
-                    &flat_keys,
-                    cur_keys_offset + cur_entry_idx,
-                )?
-                .compacted();
-                if seen_keys.contains(&key) {
-                    // TODO: implement configuration and logic for spark.sql.mapKeyDedupPolicy=EXCEPTION (this is default spark-config)
-                    // exec_err!("invalid argument: duplicate keys in map")
-                    // https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961
-                } else {
-                    // This code implements deduplication logic for spark.sql.mapKeyDedupPolicy=LAST_WIN (this is NOT default spark-config)
-                    keys_mask_one[cur_entry_idx] = true;
-                    values_mask_one[cur_entry_idx] = true;
-                    seen_keys.insert(key);
-                    new_last_offset += 1;
-                }
-            }
         }
         keys_mask_builder.append_array(&keys_mask_one.into());
         values_mask_builder.append_array(&values_mask_one.into());