datafusion_spark/function/map/
map_from_arrays.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 crate::function::map::utils::{
19    get_element_type, get_list_offsets, get_list_values,
20    map_from_keys_values_offsets_nulls, map_type_from_key_value_types,
21};
22use arrow::array::{Array, ArrayRef, NullArray};
23use arrow::compute::kernels::cast;
24use arrow::datatypes::{DataType, Field, FieldRef};
25use datafusion_common::config::MapKeyDedupPolicy;
26use datafusion_common::utils::take_function_args;
27use datafusion_common::{Result, internal_err};
28use datafusion_expr::{
29    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
30    Volatility,
31};
32use datafusion_functions::utils::make_scalar_function;
33use std::sync::Arc;
34
35/// Spark-compatible `map_from_arrays` expression
36/// <https://spark.apache.org/docs/latest/api/sql/index.html#map_from_arrays>
37#[derive(Debug, PartialEq, Eq, Hash)]
38pub struct MapFromArrays {
39    signature: Signature,
40}
41
42impl Default for MapFromArrays {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl MapFromArrays {
49    pub fn new() -> Self {
50        Self {
51            signature: Signature::any(2, Volatility::Immutable),
52        }
53    }
54}
55
56impl ScalarUDFImpl for MapFromArrays {
57    fn name(&self) -> &str {
58        "map_from_arrays"
59    }
60
61    fn signature(&self) -> &Signature {
62        &self.signature
63    }
64
65    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
66        internal_err!("return_field_from_args should be used instead")
67    }
68
69    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
70        let [keys_field, values_field] = args.arg_fields else {
71            return internal_err!("map_from_arrays expects exactly 2 arguments");
72        };
73
74        let map_type = map_type_from_key_value_types(
75            get_element_type(keys_field.data_type())?,
76            get_element_type(values_field.data_type())?,
77        );
78        // Spark marks map_from_arrays as null intolerant, so the output is
79        // nullable if either input is nullable.
80        let nullable = keys_field.is_nullable() || values_field.is_nullable();
81        Ok(Arc::new(Field::new(self.name(), map_type, nullable)))
82    }
83
84    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
85        let last_value_wins =
86            args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin;
87        make_scalar_function(
88            move |args: &[ArrayRef]| map_from_arrays_inner(args, last_value_wins),
89            vec![],
90        )(&args.args)
91    }
92}
93
94fn map_from_arrays_inner(args: &[ArrayRef], last_value_wins: bool) -> Result<ArrayRef> {
95    let [keys, values] = take_function_args("map_from_arrays", args)?;
96
97    if *keys.data_type() == DataType::Null || *values.data_type() == DataType::Null {
98        return Ok(cast(
99            &NullArray::new(keys.len()),
100            &map_type_from_key_value_types(
101                get_element_type(keys.data_type())?,
102                get_element_type(values.data_type())?,
103            ),
104        )?);
105    }
106
107    map_from_keys_values_offsets_nulls(
108        get_list_values(keys)?,
109        get_list_values(values)?,
110        &get_list_offsets(keys)?,
111        &get_list_offsets(values)?,
112        keys.nulls(),
113        values.nulls(),
114        last_value_wins,
115    )
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn test_map_from_arrays_nullability_and_type() {
124        let func = MapFromArrays::new();
125
126        let keys_field: FieldRef = Arc::new(Field::new(
127            "keys",
128            DataType::List(Arc::new(Field::new("item", DataType::Int32, false))),
129            false,
130        ));
131        let values_field: FieldRef = Arc::new(Field::new(
132            "values",
133            DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
134            false,
135        ));
136
137        let out = func
138            .return_field_from_args(ReturnFieldArgs {
139                arg_fields: &[Arc::clone(&keys_field), Arc::clone(&values_field)],
140                scalar_arguments: &[None, None],
141            })
142            .expect("return_field_from_args should succeed");
143
144        let expected_type =
145            map_type_from_key_value_types(&DataType::Int32, &DataType::Utf8);
146        assert_eq!(out.data_type(), &expected_type);
147        assert!(
148            !out.is_nullable(),
149            "map_from_arrays should be non-nullable when both inputs are non-nullable"
150        );
151
152        let nullable_keys: FieldRef = Arc::new(Field::new(
153            "keys",
154            DataType::List(Arc::new(Field::new("item", DataType::Int32, false))),
155            true,
156        ));
157
158        let out_nullable = func
159            .return_field_from_args(ReturnFieldArgs {
160                arg_fields: &[nullable_keys, values_field],
161                scalar_arguments: &[None, None],
162            })
163            .expect("return_field_from_args should succeed");
164
165        assert!(
166            out_nullable.is_nullable(),
167            "map_from_arrays should be nullable when any input is nullable"
168        );
169    }
170}
datafusion_spark/function/map/map_from_arrays.rs

datafusion_spark/function/map/
map_from_arrays.rs
