datafusion_spark/function/map/
str_to_map.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::collections::{HashMap, HashSet};
19use std::sync::Arc;
20
21use arrow::array::{
22    Array, ArrayRef, MapBuilder, MapFieldNames, StringArrayType, StringBuilder,
23};
24use arrow::buffer::NullBuffer;
25use arrow::datatypes::{DataType, Field, FieldRef};
26use datafusion_common::cast::{
27    as_large_string_array, as_string_array, as_string_view_array,
28};
29use datafusion_common::{Result, exec_err, internal_err};
30use datafusion_expr::{
31    ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature,
32    TypeSignature, Volatility,
33};
34
35use crate::function::map::utils::map_type_from_key_value_types;
36use datafusion_common::config::MapKeyDedupPolicy;
37
38const DEFAULT_PAIR_DELIM: &str = ",";
39const DEFAULT_KV_DELIM: &str = ":";
40
41/// Spark-compatible `str_to_map` expression
42/// <https://spark.apache.org/docs/latest/api/sql/index.html#str_to_map>
43///
44/// Creates a map from a string by splitting on delimiters.
45/// str_to_map(text[, pairDelim[, keyValueDelim]]) -> Map<String, String>
46///
47/// - text: The input string
48/// - pairDelim: Delimiter between key-value pairs (default: ',')
49/// - keyValueDelim: Delimiter between key and value (default: ':')
50///
51/// # Duplicate Key Handling
52/// Mirrors Spark's [`spark.sql.mapKeyDedupPolicy`](https://github.com/apache/spark/blob/v4.0.0/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4502-L4511),
53/// wired through DataFusion's `datafusion.spark.map_key_dedup_policy`:
54/// - `EXCEPTION` (default): error on duplicate keys.
55/// - `LAST_WIN`: keep the last occurrence of each duplicate key.
56#[derive(Debug, PartialEq, Eq, Hash)]
57pub struct SparkStrToMap {
58    signature: Signature,
59}
60
61impl Default for SparkStrToMap {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl SparkStrToMap {
68    pub fn new() -> Self {
69        Self {
70            signature: Signature::one_of(
71                vec![
72                    // str_to_map(text)
73                    TypeSignature::String(1),
74                    // str_to_map(text, pairDelim)
75                    TypeSignature::String(2),
76                    // str_to_map(text, pairDelim, keyValueDelim)
77                    TypeSignature::String(3),
78                ],
79                Volatility::Immutable,
80            ),
81        }
82    }
83}
84
85impl ScalarUDFImpl for SparkStrToMap {
86    fn name(&self) -> &str {
87        "str_to_map"
88    }
89
90    fn signature(&self) -> &Signature {
91        &self.signature
92    }
93
94    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
95        internal_err!("return_field_from_args should be used instead")
96    }
97
98    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
99        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
100        let map_type = map_type_from_key_value_types(&DataType::Utf8, &DataType::Utf8);
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        let arrays: Vec<ArrayRef> = ColumnarValue::values_to_arrays(&args.args)?;
108        let result = str_to_map_inner(&arrays, last_value_wins)?;
109        Ok(ColumnarValue::Array(result))
110    }
111}
112
113fn str_to_map_inner(args: &[ArrayRef], last_value_wins: bool) -> Result<ArrayRef> {
114    match args.len() {
115        1 => match args[0].data_type() {
116            DataType::Utf8 => {
117                str_to_map_impl(as_string_array(&args[0])?, None, None, last_value_wins)
118            }
119            DataType::LargeUtf8 => str_to_map_impl(
120                as_large_string_array(&args[0])?,
121                None,
122                None,
123                last_value_wins,
124            ),
125            DataType::Utf8View => str_to_map_impl(
126                as_string_view_array(&args[0])?,
127                None,
128                None,
129                last_value_wins,
130            ),
131            other => exec_err!(
132                "Unsupported data type {other:?} for str_to_map, \
133                expected Utf8, LargeUtf8, or Utf8View"
134            ),
135        },
136        2 => match (args[0].data_type(), args[1].data_type()) {
137            (DataType::Utf8, DataType::Utf8) => str_to_map_impl(
138                as_string_array(&args[0])?,
139                Some(as_string_array(&args[1])?),
140                None,
141                last_value_wins,
142            ),
143            (DataType::LargeUtf8, DataType::LargeUtf8) => str_to_map_impl(
144                as_large_string_array(&args[0])?,
145                Some(as_large_string_array(&args[1])?),
146                None,
147                last_value_wins,
148            ),
149            (DataType::Utf8View, DataType::Utf8View) => str_to_map_impl(
150                as_string_view_array(&args[0])?,
151                Some(as_string_view_array(&args[1])?),
152                None,
153                last_value_wins,
154            ),
155            (t1, t2) => exec_err!(
156                "Unsupported data types ({t1:?}, {t2:?}) for str_to_map, \
157                expected matching Utf8, LargeUtf8, or Utf8View"
158            ),
159        },
160        3 => match (
161            args[0].data_type(),
162            args[1].data_type(),
163            args[2].data_type(),
164        ) {
165            (DataType::Utf8, DataType::Utf8, DataType::Utf8) => str_to_map_impl(
166                as_string_array(&args[0])?,
167                Some(as_string_array(&args[1])?),
168                Some(as_string_array(&args[2])?),
169                last_value_wins,
170            ),
171            (DataType::LargeUtf8, DataType::LargeUtf8, DataType::LargeUtf8) => {
172                str_to_map_impl(
173                    as_large_string_array(&args[0])?,
174                    Some(as_large_string_array(&args[1])?),
175                    Some(as_large_string_array(&args[2])?),
176                    last_value_wins,
177                )
178            }
179            (DataType::Utf8View, DataType::Utf8View, DataType::Utf8View) => {
180                str_to_map_impl(
181                    as_string_view_array(&args[0])?,
182                    Some(as_string_view_array(&args[1])?),
183                    Some(as_string_view_array(&args[2])?),
184                    last_value_wins,
185                )
186            }
187            (t1, t2, t3) => exec_err!(
188                "Unsupported data types ({t1:?}, {t2:?}, {t3:?}) for str_to_map, \
189                expected matching Utf8, LargeUtf8, or Utf8View"
190            ),
191        },
192        n => exec_err!("str_to_map expects 1-3 arguments, got {n}"),
193    }
194}
195
196fn str_to_map_impl<'a, V: StringArrayType<'a> + Copy>(
197    text_array: V,
198    pair_delim_array: Option<V>,
199    kv_delim_array: Option<V>,
200    last_value_wins: bool,
201) -> Result<ArrayRef> {
202    let num_rows = text_array.len();
203
204    // Precompute combined null buffer from all input arrays.
205    // NullBuffer::union_many performs a bitmap-level AND, which is more
206    // efficient than checking per-row nullability inline.
207    let combined_nulls = NullBuffer::union_many([
208        text_array.nulls(),
209        pair_delim_array.as_ref().and_then(|a| a.nulls()),
210        kv_delim_array.as_ref().and_then(|a| a.nulls()),
211    ]);
212
213    // Use field names matching map_type_from_key_value_types: "key" and "value"
214    let field_names = MapFieldNames {
215        entry: "entries".to_string(),
216        key: "key".to_string(),
217        value: "value".to_string(),
218    };
219    let mut map_builder = MapBuilder::new(
220        Some(field_names),
221        StringBuilder::new(),
222        StringBuilder::new(),
223    );
224
225    let mut seen_keys = HashSet::new();
226    // LAST_WIN buffers pairs to support in-place value overwrite at the key's
227    // first-seen position — matches Spark's `ArrayBasedMapBuilder`.
228    let mut pairs: Vec<(&str, Option<&str>)> = Vec::new();
229    let mut key_positions: HashMap<&str, usize> = HashMap::new();
230    for row_idx in 0..num_rows {
231        if combined_nulls.as_ref().is_some_and(|n| n.is_null(row_idx)) {
232            map_builder.append(false)?;
233            continue;
234        }
235
236        // Per-row delimiter extraction
237        let pair_delim =
238            pair_delim_array.map_or(DEFAULT_PAIR_DELIM, |a| a.value(row_idx));
239        let kv_delim = kv_delim_array.map_or(DEFAULT_KV_DELIM, |a| a.value(row_idx));
240
241        let text = text_array.value(row_idx);
242        if text.is_empty() {
243            // Empty string -> map with empty key and NULL value (Spark behavior)
244            map_builder.keys().append_value("");
245            map_builder.values().append_null();
246            map_builder.append(true)?;
247            continue;
248        }
249
250        if last_value_wins {
251            pairs.clear();
252            key_positions.clear();
253            for pair in text.split(pair_delim) {
254                if pair.is_empty() {
255                    continue;
256                }
257                let mut kv_iter = pair.splitn(2, kv_delim);
258                let key = kv_iter.next().unwrap_or("");
259                let value = kv_iter.next();
260                match key_positions.get(key) {
261                    Some(&idx) => pairs[idx].1 = value,
262                    None => {
263                        key_positions.insert(key, pairs.len());
264                        pairs.push((key, value));
265                    }
266                }
267            }
268            for (key, value) in &pairs {
269                map_builder.keys().append_value(key);
270                match value {
271                    Some(v) => map_builder.values().append_value(v),
272                    None => map_builder.values().append_null(),
273                }
274            }
275        } else {
276            seen_keys.clear();
277            for pair in text.split(pair_delim) {
278                if pair.is_empty() {
279                    continue;
280                }
281
282                let mut kv_iter = pair.splitn(2, kv_delim);
283                let key = kv_iter.next().unwrap_or("");
284                let value = kv_iter.next();
285
286                if !seen_keys.insert(key) {
287                    return exec_err!(
288                        "[DUPLICATED_MAP_KEY] Duplicate map key '{key}' was found, \
289                         please check the input data. To allow duplicate keys with \
290                         last-value-wins semantics, set \
291                         `datafusion.spark.map_key_dedup_policy` to `LAST_WIN`."
292                    );
293                }
294
295                map_builder.keys().append_value(key);
296                match value {
297                    Some(v) => map_builder.values().append_value(v),
298                    None => map_builder.values().append_null(),
299                }
300            }
301        }
302        map_builder.append(true)?;
303    }
304
305    Ok(Arc::new(map_builder.finish()))
306}
datafusion_spark/function/map/str_to_map.rs

datafusion_spark/function/map/
str_to_map.rs
