feat(spark): add concat_ws with array support (#20928)

davidlghellin · web-flow · commit e5f7af15a6da · 2026-06-13T01:37:17.000Z
## Which issue does this PR close? - Part of #15914 ## Rationale for this change DataFusion core's `concat_ws` does not support array arguments. Spark's `concat_ws(sep, ...)` accepts both scalar strings and arrays, expanding array elements and skipping nulls. This is needed for Spark compatibility in the `datafusion-spark` crate. ## What changes are included in this PR? - New `SparkConcatWs` UDF in `datafusion/spark/src/function/string/concat_ws.rs` - Supports `concat_ws(sep, str1, str2, ...)` with scalar strings - Supports array arguments: `concat_ws(',', array('a', 'b'), 'c')` → `"a,b,c"` - Null scalars and null array elements are skipped (Spark behavior) - Null separator returns NULL - Zero value arguments (`concat_ws(',')`) returns empty string - Supports Utf8, LargeUtf8, Utf8View, List, and LargeList types - Registered the function in `mod.rs` (`make_udf_function!`, `export_functions!`, `functions()`) - Replaced commented-out SLT tests with 14 working test cases covering basic usage, arrays, mixed arguments, nulls, column expressions, and edge cases ## Are these changes tested? Yes. - 7 unit tests in `concat_ws.rs` (basic, null values skipped, null separator, list arrays, list with nulls, mixed scalar+list, multiple rows) - 14 SLT tests in `spark/string/concat_ws.slt` covering scalars, arrays, nulls, column expressions, and edge cases ## Are there any user-facing changes? No. This is a new function in the `datafusion-spark` crate only.
diff --git a/datafusion/spark/src/function/string/concat_ws.rs b/datafusion/spark/src/function/string/concat_ws.rs
@@ -0,0 +1,297 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Spark-compatible `concat_ws`: joins strings (and array elements) with a separator.
+//!
+//! Null scalar args and null array elements are skipped; a null separator yields a
+//! null row. Non-string args are coerced to STRING; list args (`List`, `LargeList`,
+//! `ListView`, `LargeListView`, `FixedSizeList`) expand their elements.
+//!
+//! Differences with DataFusion core `concat_ws`:
+//! - Accepts list arguments and expands their elements
+//! - Always returns Utf8 (Spark's `STRING` type)
+//! - Coerces non-string scalars (numbers, booleans, dates, ...) to Utf8
+
+use std::fmt::Write as _;
+use std::sync::Arc;
+
+use arrow::array::{
+    Array, ArrayRef, AsArray, GenericListArray, LargeStringArray, OffsetSizeTrait,
+    StringArray, StringBuilder, StringViewArray,
+};
+use arrow::datatypes::{DataType, Field};
+use datafusion_common::{Result, ScalarValue};
+use datafusion_expr::{
+    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
+};
+
+use crate::function::error_utils::{
+    invalid_arg_count_exec_err, unsupported_data_type_exec_err,
+};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkConcatWs {
+    signature: Signature,
+}
+
+impl Default for SparkConcatWs {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SparkConcatWs {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::user_defined(Volatility::Immutable),
+        }
+    }
+}
+
+impl ScalarUDFImpl for SparkConcatWs {
+    fn name(&self) -> &str {
+        "concat_ws"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Utf8)
+    }
+
+    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+        if arg_types.is_empty() {
+            return Err(invalid_arg_count_exec_err("concat_ws", (1, i32::MAX), 0));
+        }
+        Ok(arg_types
+            .iter()
+            .enumerate()
+            .map(|(i, dt)| match dt {
+                DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => dt.clone(),
+                // Non-separator list args expand their elements at runtime.
+                // Normalize the list variant so the kernel only sees
+                // List/LargeList, AND force the element type to Utf8 so the
+                // planner inserts a cast for non-string children (Spark
+                // coerces them to STRING the same way it does for scalars).
+                DataType::List(f)
+                | DataType::ListView(f)
+                | DataType::FixedSizeList(f, _)
+                    if i > 0 =>
+                {
+                    DataType::List(Arc::new(Field::new(
+                        f.name(),
+                        DataType::Utf8,
+                        f.is_nullable(),
+                    )))
+                }
+                DataType::LargeList(f) | DataType::LargeListView(f) if i > 0 => {
+                    DataType::LargeList(Arc::new(Field::new(
+                        f.name(),
+                        DataType::Utf8,
+                        f.is_nullable(),
+                    )))
+                }
+                // Spark casts everything else (numbers, booleans, dates,
+                // binary, null...) to STRING.
+                _ => DataType::Utf8,
+            })
+            .collect())
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        // Only separator provided → empty string (or NULL if separator is null).
+        // Arg-count validation happens in coerce_types at planning time.
+        if args.args.len() == 1 {
+            return only_separator(&args.args[0]);
+        }
+
+        spark_concat_ws(&args.args, args.number_rows)
+    }
+}
+
+fn only_separator(sep: &ColumnarValue) -> Result<ColumnarValue> {
+    match sep {
+        ColumnarValue::Scalar(s) if s.is_null() => {
+            Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
+        }
+        ColumnarValue::Scalar(_) => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+            String::new(),
+        )))),
+        ColumnarValue::Array(arr) => {
+            let mut builder = StringBuilder::with_capacity(arr.len(), 0);
+            for row_idx in 0..arr.len() {
+                if arr.is_null(row_idx) {
+                    builder.append_null();
+                } else {
+                    builder.append_value("");
+                }
+            }
+            Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef))
+        }
+    }
+}
+
+fn spark_concat_ws(args: &[ColumnarValue], num_rows: usize) -> Result<ColumnarValue> {
+    let arrays = ColumnarValue::values_to_arrays(args)?;
+    let sep_view = StringView::try_new(&arrays[0])?;
+    let arg_views: Vec<ArgView> = arrays[1..]
+        .iter()
+        .map(ArgView::try_new)
+        .collect::<Result<_>>()?;
+
+    let mut builder = StringBuilder::with_capacity(num_rows, num_rows * 16);
+
+    for row_idx in 0..num_rows {
+        if sep_view.is_null(row_idx) {
+            builder.append_null();
+            continue;
+        }
+
+        // Write parts directly into the builder via its `fmt::Write` impl;
+        // `append_value("")` then finalises the row (offset + validity) with
+        // no extra copy from an intermediate `String`.
+        let separator = sep_view.value(row_idx);
+        let mut first = true;
+        for view in &arg_views {
+            view.write_row(row_idx, separator, &mut builder, &mut first)?;
+        }
+        builder.append_value("");
+    }
+
+    Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef))
+}
+
+/// Typed view over a string array that downcasts once and exposes
+/// per-row access without further dispatch.
+enum StringView<'a> {
+    Utf8(&'a StringArray),
+    LargeUtf8(&'a LargeStringArray),
+    Utf8View(&'a StringViewArray),
+}
+
+impl<'a> StringView<'a> {
+    fn try_new(arr: &'a ArrayRef) -> Result<Self> {
+        match arr.data_type() {
+            DataType::Utf8 => Ok(Self::Utf8(arr.as_string::<i32>())),
+            DataType::LargeUtf8 => Ok(Self::LargeUtf8(arr.as_string::<i64>())),
+            DataType::Utf8View => Ok(Self::Utf8View(arr.as_string_view())),
+            other => Err(unsupported_data_type_exec_err("concat_ws", "STRING", other)),
+        }
+    }
+
+    fn value(&self, idx: usize) -> &str {
+        match self {
+            Self::Utf8(a) => a.value(idx),
+            Self::LargeUtf8(a) => a.value(idx),
+            Self::Utf8View(a) => a.value(idx),
+        }
+    }
+
+    fn is_null(&self, idx: usize) -> bool {
+        match self {
+            Self::Utf8(a) => a.is_null(idx),
+            Self::LargeUtf8(a) => a.is_null(idx),
+            Self::Utf8View(a) => a.is_null(idx),
+        }
+    }
+}
+
+/// Per-argument view: a string array or a list of strings. The downcast
+/// happens once at construction time. `DataType::Null` cannot appear here —
+/// `coerce_types` rewrites it to `Utf8` before invocation.
+enum ArgView<'a> {
+    Str(StringView<'a>),
+    List(&'a GenericListArray<i32>),
+    LargeList(&'a GenericListArray<i64>),
+}
+
+impl<'a> ArgView<'a> {
+    fn try_new(arr: &'a ArrayRef) -> Result<Self> {
+        match arr.data_type() {
+            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => {
+                Ok(Self::Str(StringView::try_new(arr)?))
+            }
+            DataType::List(_) => Ok(Self::List(arr.as_list::<i32>())),
+            DataType::LargeList(_) => Ok(Self::LargeList(arr.as_list::<i64>())),
+            other => Err(unsupported_data_type_exec_err(
+                "concat_ws",
+                "STRING or ARRAY<STRING>",
+                other,
+            )),
+        }
+    }
+
+    fn write_row(
+        &self,
+        row_idx: usize,
+        sep: &str,
+        builder: &mut StringBuilder,
+        first: &mut bool,
+    ) -> Result<()> {
+        match self {
+            Self::Str(view) => {
+                if !view.is_null(row_idx) {
+                    push_part(builder, view.value(row_idx), sep, first);
+                }
+            }
+            Self::List(list) => write_list_row(*list, row_idx, sep, builder, first)?,
+            Self::LargeList(list) => write_list_row(*list, row_idx, sep, builder, first)?,
+        }
+        Ok(())
+    }
+}
+
+fn write_list_row<O: OffsetSizeTrait>(
+    list: &GenericListArray<O>,
+    row_idx: usize,
+    sep: &str,
+    builder: &mut StringBuilder,
+    first: &mut bool,
+) -> Result<()> {
+    if list.is_null(row_idx) {
+        return Ok(());
+    }
+    let values = list.value(row_idx);
+    // An empty array (e.g. `array()`) contributes nothing — Spark renders it
+    // as the empty string, not an error.
+    if values.is_empty() {
+        return Ok(());
+    }
+    let view = StringView::try_new(&values)?;
+    for i in 0..values.len() {
+        if !view.is_null(i) {
+            push_part(builder, view.value(i), sep, first);
+        }
+    }
+    Ok(())
+}
+
+// `StringBuilder::write_str` only does `extend_from_slice` and never errors;
+// the `.expect(..)` is a documentation hint, not a real failure path.
+fn push_part(builder: &mut StringBuilder, part: &str, sep: &str, first: &mut bool) {
+    if !*first {
+        builder
+            .write_str(sep)
+            .expect("StringBuilder::write_str is infallible");
+    }
+    *first = false;
+    builder
+        .write_str(part)
+        .expect("StringBuilder::write_str is infallible");
+}
diff --git a/datafusion/spark/src/function/string/mod.rs b/datafusion/spark/src/function/string/mod.rs
@@ -19,6 +19,7 @@ pub mod ascii;
 pub mod base64;
 pub mod char;
 pub mod concat;
+pub mod concat_ws;
 pub mod elt;
 pub mod format_string;
 pub mod ilike;
@@ -40,6 +41,7 @@ make_udf_function!(ascii::SparkAscii, ascii);
 make_udf_function!(base64::SparkBase64, base64);
 make_udf_function!(char::CharFunc, char);
 make_udf_function!(concat::SparkConcat, concat);
+make_udf_function!(concat_ws::SparkConcatWs, concat_ws);
 make_udf_function!(ilike::SparkILike, ilike);
 make_udf_function!(length::SparkLengthFunc, length);
 make_udf_function!(elt::SparkElt, elt);
@@ -77,6 +79,11 @@ pub mod expr_fn {
         "Concatenates multiple input strings into a single string. Returns NULL if any input is NULL.",
         args
     ));
+    export_functions!((
+        concat_ws,
+        "Concatenates strings with separator. Supports arrays. Null values are skipped.",
+        sep args
+    ));
     export_functions!((
         elt,
         "Returns the n-th input (1-indexed), e.g. returns 2nd input when n is 2. The function returns NULL if the index is 0 or exceeds the length of the array.",
@@ -142,6 +149,7 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
         base64(),
         char(),
         concat(),
+        concat_ws(),
         elt(),
         ilike(),
         length(),
diff --git a/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt b/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt
