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 std::any::Any;
19
20use crate::function::map::utils::{
21    get_element_type, get_list_offsets, get_list_values,
22    map_from_keys_values_offsets_nulls, map_type_from_key_value_types,
23};
24use arrow::array::{Array, ArrayRef, NullArray};
25use arrow::compute::kernels::cast;
26use arrow::datatypes::DataType;
27use datafusion_common::utils::take_function_args;
28use datafusion_common::Result;
29use datafusion_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
30use datafusion_functions::utils::make_scalar_function;
31
32/// Spark-compatible `map_from_arrays` expression
33/// <https://spark.apache.org/docs/latest/api/sql/index.html#map_from_arrays>
34#[derive(Debug, PartialEq, Eq, Hash)]
35pub struct MapFromArrays {
36    signature: Signature,
37}
38
39impl Default for MapFromArrays {
40    fn default() -> Self {
41        Self::new()
42    }
43}
44
45impl MapFromArrays {
46    pub fn new() -> Self {
47        Self {
48            signature: Signature::any(2, Volatility::Immutable),
49        }
50    }
51}
52
53impl ScalarUDFImpl for MapFromArrays {
54    fn as_any(&self) -> &dyn Any {
55        self
56    }
57
58    fn name(&self) -> &str {
59        "map_from_arrays"
60    }
61
62    fn signature(&self) -> &Signature {
63        &self.signature
64    }
65
66    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
67        let [key_type, value_type] = take_function_args("map_from_arrays", arg_types)?;
68        Ok(map_type_from_key_value_types(
69            get_element_type(key_type)?,
70            get_element_type(value_type)?,
71        ))
72    }
73
74    fn invoke_with_args(
75        &self,
76        args: datafusion_expr::ScalarFunctionArgs,
77    ) -> Result<ColumnarValue> {
78        make_scalar_function(map_from_arrays_inner, vec![])(&args.args)
79    }
80}
81
82fn map_from_arrays_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
83    let [keys, values] = take_function_args("map_from_arrays", args)?;
84
85    if matches!(keys.data_type(), DataType::Null)
86        || matches!(values.data_type(), DataType::Null)
87    {
88        return Ok(cast(
89            &NullArray::new(keys.len()),
90            &map_type_from_key_value_types(
91                get_element_type(keys.data_type())?,
92                get_element_type(values.data_type())?,
93            ),
94        )?);
95    }
96
97    map_from_keys_values_offsets_nulls(
98        get_list_values(keys)?,
99        get_list_values(values)?,
100        &get_list_offsets(keys)?,
101        &get_list_offsets(values)?,
102        keys.nulls(),
103        values.nulls(),
104    )
105}
datafusion_spark/function/map/map_from_arrays.rs

datafusion_spark/function/map/
map_from_arrays.rs
