datafusion_spark/function/map/
map_from_entries.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use crate::function::map::utils::{
21    get_list_offsets, get_list_values, map_from_keys_values_offsets_nulls,
22    map_type_from_key_value_types,
23};
24use arrow::array::{Array, ArrayRef, NullBufferBuilder, StructArray};
25use arrow::buffer::NullBuffer;
26use arrow::datatypes::{DataType, Field, FieldRef};
27use datafusion_common::config::MapKeyDedupPolicy;
28use datafusion_common::utils::take_function_args;
29use datafusion_common::{Result, exec_err, internal_err};
30use datafusion_expr::{
31    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32    Volatility,
33};
34use datafusion_functions::utils::make_scalar_function;
35
36/// Spark-compatible `map_from_entries` expression
37/// <https://spark.apache.org/docs/latest/api/sql/index.html#map_from_entries>
38#[derive(Debug, PartialEq, Eq, Hash)]
39pub struct MapFromEntries {
40    signature: Signature,
41}
42
43impl Default for MapFromEntries {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49impl MapFromEntries {
50    pub fn new() -> Self {
51        Self {
52            signature: Signature::array(Volatility::Immutable),
53        }
54    }
55}
56
57impl ScalarUDFImpl for MapFromEntries {
58    fn name(&self) -> &str {
59        "map_from_entries"
60    }
61
62    fn signature(&self) -> &Signature {
63        &self.signature
64    }
65
66    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
67        internal_err!("return_field_from_args should be used instead")
68    }
69
70    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
71        let [entries_field] = args.arg_fields else {
72            return exec_err!("map_from_entries: expected one argument");
73        };
74
75        let (entries_element_field, entries_element_type) =
76            match entries_field.data_type() {
77                DataType::List(field)
78                | DataType::LargeList(field)
79                | DataType::FixedSizeList(field, _) => {
80                    Ok((field.as_ref(), field.data_type()))
81                }
82                wrong_type => exec_err!(
83                    "map_from_entries: expected array<struct<key, value>>, got {:?}",
84                    wrong_type
85                ),
86            }?;
87
88        let (keys_type, values_type) = match entries_element_type {
89            DataType::Struct(fields) if fields.len() == 2 => {
90                Ok((fields[0].data_type(), fields[1].data_type()))
91            }
92            wrong_type => exec_err!(
93                "map_from_entries: expected array<struct<key, value>>, got {:?}",
94                wrong_type
95            ),
96        }?;
97
98        let map_type = map_type_from_key_value_types(keys_type, values_type);
99        let nullable = entries_field.is_nullable() || entries_element_field.is_nullable();
100
101        Ok(Arc::new(Field::new(self.name(), map_type, nullable)))
102    }
103
104    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
105        let last_value_wins =
106            args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin;
107        make_scalar_function(
108            move |args: &[ArrayRef]| map_from_entries_inner(args, last_value_wins),
109            vec![],
110        )(&args.args)
111    }
112}
113
114fn map_from_entries_inner(args: &[ArrayRef], last_value_wins: bool) -> Result<ArrayRef> {
115    let [entries] = take_function_args("map_from_entries", args)?;
116    let entries_offsets = get_list_offsets(entries)?;
117    let entries_values = get_list_values(entries)?;
118
119    let (flat_keys, flat_values) =
120        match entries_values.as_any().downcast_ref::<StructArray>() {
121            Some(a) => Ok((a.column(0), a.column(1))),
122            None => exec_err!(
123                "map_from_entries: expected array<struct<key, value>>, got {:?}",
124                entries_values.data_type()
125            ),
126        }?;
127
128    let entries_with_nulls = entries_values.nulls().and_then(|entries_inner_nulls| {
129        let mut builder = NullBufferBuilder::new_with_len(0);
130        let mut cur_offset = entries_offsets
131            .first()
132            .map(|offset| *offset as usize)
133            .unwrap_or(0);
134
135        for next_offset in entries_offsets.iter().skip(1) {
136            let num_entries = *next_offset as usize - cur_offset;
137            builder.append(
138                entries_inner_nulls
139                    .slice(cur_offset, num_entries)
140                    .null_count()
141                    == 0,
142            );
143            cur_offset = *next_offset as usize;
144        }
145        builder.finish()
146    });
147
148    let res_nulls = NullBuffer::union(entries.nulls(), entries_with_nulls.as_ref());
149
150    map_from_keys_values_offsets_nulls(
151        flat_keys,
152        flat_values,
153        &entries_offsets,
154        &entries_offsets,
155        None,
156        res_nulls.as_ref(),
157        last_value_wins,
158    )
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use arrow::datatypes::Fields;
165
166    fn make_entries_field(array_nullable: bool, element_nullable: bool) -> FieldRef {
167        let struct_type = DataType::Struct(Fields::from(vec![
168            Field::new("key", DataType::Int32, false),
169            Field::new("value", DataType::Utf8, true),
170        ]));
171        Arc::new(Field::new(
172            "entries",
173            DataType::List(Arc::new(Field::new("item", struct_type, element_nullable))),
174            array_nullable,
175        ))
176    }
177
178    #[test]
179    fn test_map_from_entries_nullability_matches_input() {
180        let func = MapFromEntries::new();
181        let expected_type =
182            map_type_from_key_value_types(&DataType::Int32, &DataType::Utf8);
183
184        // Non-nullable array and elements => non-nullable result
185        let non_nullable_field = make_entries_field(false, false);
186        let result = func
187            .return_field_from_args(ReturnFieldArgs {
188                arg_fields: &[Arc::clone(&non_nullable_field)],
189                scalar_arguments: &[None],
190            })
191            .expect("should infer field");
192        assert!(!result.is_nullable());
193        assert_eq!(result.data_type(), &expected_type);
194
195        // Nullable elements should make result nullable even if array is non-nullable
196        let element_nullable_field = make_entries_field(false, true);
197        let result = func
198            .return_field_from_args(ReturnFieldArgs {
199                arg_fields: &[Arc::clone(&element_nullable_field)],
200                scalar_arguments: &[None],
201            })
202            .expect("should infer field");
203        assert!(result.is_nullable());
204        assert_eq!(result.data_type(), &expected_type);
205
206        // Nullable array should also yield nullable result
207        let array_nullable_field = make_entries_field(true, false);
208        let result = func
209            .return_field_from_args(ReturnFieldArgs {
210                arg_fields: &[Arc::clone(&array_nullable_field)],
211                scalar_arguments: &[None],
212            })
213            .expect("should infer field");
214        assert!(result.is_nullable());
215        assert_eq!(result.data_type(), &expected_type);
216    }
217}
datafusion_spark/function/map/map_from_entries.rs

datafusion_spark/function/map/
map_from_entries.rs
