datafusion_spark/function/string/
concat_ws.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
18//! Spark-compatible `concat_ws`: joins strings (and array elements) with a separator.
19//!
20//! Null scalar args and null array elements are skipped; a null separator yields a
21//! null row. Non-string args are coerced to STRING; list args (`List`, `LargeList`,
22//! `ListView`, `LargeListView`, `FixedSizeList`) expand their elements.
23//!
24//! Differences with DataFusion core `concat_ws`:
25//! - Accepts list arguments and expands their elements
26//! - Always returns Utf8 (Spark's `STRING` type)
27//! - Coerces non-string scalars (numbers, booleans, dates, ...) to Utf8
28
29use std::fmt::Write as _;
30use std::sync::Arc;
31
32use arrow::array::{
33    Array, ArrayRef, AsArray, GenericListArray, LargeStringArray, OffsetSizeTrait,
34    StringArray, StringBuilder, StringViewArray,
35};
36use arrow::datatypes::{DataType, Field};
37use datafusion_common::{Result, ScalarValue};
38use datafusion_expr::{
39    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
40};
41
42use crate::function::error_utils::{
43    invalid_arg_count_exec_err, unsupported_data_type_exec_err,
44};
45
46#[derive(Debug, PartialEq, Eq, Hash)]
47pub struct SparkConcatWs {
48    signature: Signature,
49}
50
51impl Default for SparkConcatWs {
52    fn default() -> Self {
53        Self::new()
54    }
55}
56
57impl SparkConcatWs {
58    pub fn new() -> Self {
59        Self {
60            signature: Signature::user_defined(Volatility::Immutable),
61        }
62    }
63}
64
65impl ScalarUDFImpl for SparkConcatWs {
66    fn name(&self) -> &str {
67        "concat_ws"
68    }
69
70    fn signature(&self) -> &Signature {
71        &self.signature
72    }
73
74    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
75        Ok(DataType::Utf8)
76    }
77
78    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
79        if arg_types.is_empty() {
80            return Err(invalid_arg_count_exec_err("concat_ws", (1, i32::MAX), 0));
81        }
82        Ok(arg_types
83            .iter()
84            .enumerate()
85            .map(|(i, dt)| match dt {
86                DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => dt.clone(),
87                // Non-separator list args expand their elements at runtime.
88                // Normalize the list variant so the kernel only sees
89                // List/LargeList, AND force the element type to Utf8 so the
90                // planner inserts a cast for non-string children (Spark
91                // coerces them to STRING the same way it does for scalars).
92                DataType::List(f)
93                | DataType::ListView(f)
94                | DataType::FixedSizeList(f, _)
95                    if i > 0 =>
96                {
97                    DataType::List(Arc::new(Field::new(
98                        f.name(),
99                        DataType::Utf8,
100                        f.is_nullable(),
101                    )))
102                }
103                DataType::LargeList(f) | DataType::LargeListView(f) if i > 0 => {
104                    DataType::LargeList(Arc::new(Field::new(
105                        f.name(),
106                        DataType::Utf8,
107                        f.is_nullable(),
108                    )))
109                }
110                // Spark casts everything else (numbers, booleans, dates,
111                // binary, null...) to STRING.
112                _ => DataType::Utf8,
113            })
114            .collect())
115    }
116
117    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
118        // Only separator provided → empty string (or NULL if separator is null).
119        // Arg-count validation happens in coerce_types at planning time.
120        if args.args.len() == 1 {
121            return only_separator(&args.args[0]);
122        }
123
124        spark_concat_ws(&args.args, args.number_rows)
125    }
126}
127
128fn only_separator(sep: &ColumnarValue) -> Result<ColumnarValue> {
129    match sep {
130        ColumnarValue::Scalar(s) if s.is_null() => {
131            Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
132        }
133        ColumnarValue::Scalar(_) => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
134            String::new(),
135        )))),
136        ColumnarValue::Array(arr) => {
137            let mut builder = StringBuilder::with_capacity(arr.len(), 0);
138            for row_idx in 0..arr.len() {
139                if arr.is_null(row_idx) {
140                    builder.append_null();
141                } else {
142                    builder.append_value("");
143                }
144            }
145            Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef))
146        }
147    }
148}
149
150fn spark_concat_ws(args: &[ColumnarValue], num_rows: usize) -> Result<ColumnarValue> {
151    let arrays = ColumnarValue::values_to_arrays(args)?;
152    let sep_view = StringView::try_new(&arrays[0])?;
153    let arg_views: Vec<ArgView> = arrays[1..]
154        .iter()
155        .map(ArgView::try_new)
156        .collect::<Result<_>>()?;
157
158    let mut builder = StringBuilder::with_capacity(num_rows, num_rows * 16);
159
160    for row_idx in 0..num_rows {
161        if sep_view.is_null(row_idx) {
162            builder.append_null();
163            continue;
164        }
165
166        // Write parts directly into the builder via its `fmt::Write` impl;
167        // `append_value("")` then finalises the row (offset + validity) with
168        // no extra copy from an intermediate `String`.
169        let separator = sep_view.value(row_idx);
170        let mut first = true;
171        for view in &arg_views {
172            view.write_row(row_idx, separator, &mut builder, &mut first)?;
173        }
174        builder.append_value("");
175    }
176
177    Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef))
178}
179
180/// Typed view over a string array that downcasts once and exposes
181/// per-row access without further dispatch.
182enum StringView<'a> {
183    Utf8(&'a StringArray),
184    LargeUtf8(&'a LargeStringArray),
185    Utf8View(&'a StringViewArray),
186}
187
188impl<'a> StringView<'a> {
189    fn try_new(arr: &'a ArrayRef) -> Result<Self> {
190        match arr.data_type() {
191            DataType::Utf8 => Ok(Self::Utf8(arr.as_string::<i32>())),
192            DataType::LargeUtf8 => Ok(Self::LargeUtf8(arr.as_string::<i64>())),
193            DataType::Utf8View => Ok(Self::Utf8View(arr.as_string_view())),
194            other => Err(unsupported_data_type_exec_err("concat_ws", "STRING", other)),
195        }
196    }
197
198    fn value(&self, idx: usize) -> &str {
199        match self {
200            Self::Utf8(a) => a.value(idx),
201            Self::LargeUtf8(a) => a.value(idx),
202            Self::Utf8View(a) => a.value(idx),
203        }
204    }
205
206    fn is_null(&self, idx: usize) -> bool {
207        match self {
208            Self::Utf8(a) => a.is_null(idx),
209            Self::LargeUtf8(a) => a.is_null(idx),
210            Self::Utf8View(a) => a.is_null(idx),
211        }
212    }
213}
214
215/// Per-argument view: a string array or a list of strings. The downcast
216/// happens once at construction time. `DataType::Null` cannot appear here —
217/// `coerce_types` rewrites it to `Utf8` before invocation.
218enum ArgView<'a> {
219    Str(StringView<'a>),
220    List(&'a GenericListArray<i32>),
221    LargeList(&'a GenericListArray<i64>),
222}
223
224impl<'a> ArgView<'a> {
225    fn try_new(arr: &'a ArrayRef) -> Result<Self> {
226        match arr.data_type() {
227            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
228                Ok(Self::Str(StringView::try_new(arr)?))
229            }
230            DataType::List(_) => Ok(Self::List(arr.as_list::<i32>())),
231            DataType::LargeList(_) => Ok(Self::LargeList(arr.as_list::<i64>())),
232            other => Err(unsupported_data_type_exec_err(
233                "concat_ws",
234                "STRING or ARRAY<STRING>",
235                other,
236            )),
237        }
238    }
239
240    fn write_row(
241        &self,
242        row_idx: usize,
243        sep: &str,
244        builder: &mut StringBuilder,
245        first: &mut bool,
246    ) -> Result<()> {
247        match self {
248            Self::Str(view) => {
249                if !view.is_null(row_idx) {
250                    push_part(builder, view.value(row_idx), sep, first);
251                }
252            }
253            Self::List(list) => write_list_row(*list, row_idx, sep, builder, first)?,
254            Self::LargeList(list) => write_list_row(*list, row_idx, sep, builder, first)?,
255        }
256        Ok(())
257    }
258}
259
260fn write_list_row<O: OffsetSizeTrait>(
261    list: &GenericListArray<O>,
262    row_idx: usize,
263    sep: &str,
264    builder: &mut StringBuilder,
265    first: &mut bool,
266) -> Result<()> {
267    if list.is_null(row_idx) {
268        return Ok(());
269    }
270    let values = list.value(row_idx);
271    // An empty array (e.g. `array()`) contributes nothing — Spark renders it
272    // as the empty string, not an error.
273    if values.is_empty() {
274        return Ok(());
275    }
276    let view = StringView::try_new(&values)?;
277    for i in 0..values.len() {
278        if !view.is_null(i) {
279            push_part(builder, view.value(i), sep, first);
280        }
281    }
282    Ok(())
283}
284
285// `StringBuilder::write_str` only does `extend_from_slice` and never errors;
286// the `.expect(..)` is a documentation hint, not a real failure path.
287fn push_part(builder: &mut StringBuilder, part: &str, sep: &str, first: &mut bool) {
288    if !*first {
289        builder
290            .write_str(sep)
291            .expect("StringBuilder::write_str is infallible");
292    }
293    *first = false;
294    builder
295        .write_str(part)
296        .expect("StringBuilder::write_str is infallible");
297}
datafusion_spark/function/string/concat_ws.rs

datafusion_spark/function/string/
concat_ws.rs
