datafusion_spark/function/string/
concat.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 arrow::datatypes::{DataType, Field};
19use datafusion_common::arrow::datatypes::FieldRef;
20use datafusion_common::{Result, ScalarValue};
21use datafusion_expr::ReturnFieldArgs;
22use datafusion_expr::{
23    ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
24};
25use datafusion_functions::string::concat::ConcatFunc;
26use std::sync::Arc;
27
28use crate::function::null_utils::{
29    NullMaskResolution, apply_null_mask, compute_null_mask,
30};
31
32/// Spark-compatible `concat` expression
33/// <https://spark.apache.org/docs/latest/api/sql/index.html#concat>
34///
35/// Concatenates multiple input strings into a single string.
36/// Returns NULL if any input is NULL.
37///
38/// Differences with DataFusion concat:
39/// - Support 0 arguments
40/// - Return NULL if any input is NULL
41#[derive(Debug, PartialEq, Eq, Hash)]
42pub struct SparkConcat {
43    signature: Signature,
44}
45
46impl Default for SparkConcat {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52impl SparkConcat {
53    pub fn new() -> Self {
54        Self {
55            signature: Signature::user_defined(Volatility::Immutable),
56        }
57    }
58}
59
60impl ScalarUDFImpl for SparkConcat {
61    fn name(&self) -> &str {
62        "concat"
63    }
64
65    fn signature(&self) -> &Signature {
66        &self.signature
67    }
68
69    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
70        spark_concat(args)
71    }
72
73    fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
74        if arg_types.is_empty() {
75            // Spark semantics: allow concat with zero arguments
76            Ok(vec![])
77        } else {
78            // Use concat coercion rules
79            ConcatFunc::new().coerce_types(arg_types)
80        }
81    }
82    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
83        datafusion_common::internal_err!(
84            "return_type should not be called for Spark concat"
85        )
86    }
87    fn return_field_from_args(&self, args: ReturnFieldArgs<'_>) -> Result<FieldRef> {
88        // Spark semantics: concat returns NULL if ANY input is NULL
89        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
90
91        let arg_types: Vec<DataType> = args
92            .arg_fields
93            .iter()
94            .map(|f| f.data_type().clone())
95            .collect();
96        let dt = ConcatFunc::new().return_type(&arg_types)?;
97
98        Ok(Arc::new(Field::new("concat", dt.clone(), nullable)))
99    }
100}
101
102/// Concatenates strings, returning NULL if any input is NULL
103/// This is a Spark-specific wrapper around DataFusion's concat that returns NULL
104/// if any argument is NULL (Spark behavior), whereas DataFusion's concat ignores NULLs.
105fn spark_concat(args: ScalarFunctionArgs) -> Result<ColumnarValue> {
106    let ScalarFunctionArgs {
107        args: arg_values,
108        arg_fields,
109        number_rows,
110        return_field,
111        config_options,
112    } = args;
113
114    // Handle zero-argument case: return empty string
115    if arg_values.is_empty() {
116        let return_type = return_field.data_type();
117        return Ok(ColumnarValue::Scalar(ScalarValue::new_default(
118            return_type,
119        )?));
120    }
121
122    // Step 1: Check for NULL mask in incoming args
123    let null_mask = compute_null_mask(&arg_values);
124
125    // If all scalars and any is NULL, return NULL immediately
126    if matches!(null_mask, NullMaskResolution::ReturnNull) {
127        let return_type = return_field.data_type();
128        return Ok(ColumnarValue::Scalar(ScalarValue::try_new_null(
129            return_type,
130        )?));
131    }
132
133    // Step 2: Delegate to DataFusion's concat
134    let concat_func = ConcatFunc::new();
135    let return_type = return_field.data_type().clone();
136    let func_args = ScalarFunctionArgs {
137        args: arg_values,
138        arg_fields,
139        number_rows,
140        return_field,
141        config_options,
142    };
143    let result = concat_func.invoke_with_args(func_args)?;
144
145    // Step 3: Apply NULL mask to result
146    apply_null_mask(result, null_mask, &return_type)
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use crate::function::utils::test::test_scalar_function;
153    use arrow::array::{Array, StringArray};
154
155    #[test]
156    fn test_concat_basic() -> Result<()> {
157        test_scalar_function!(
158            SparkConcat::new(),
159            vec![
160                ColumnarValue::Scalar(ScalarValue::Utf8(Some("Spark".to_string()))),
161                ColumnarValue::Scalar(ScalarValue::Utf8(Some("SQL".to_string()))),
162            ],
163            Ok(Some("SparkSQL")),
164            &str,
165            DataType::Utf8,
166            StringArray
167        );
168        Ok(())
169    }
170
171    #[test]
172    fn test_concat_with_null() -> Result<()> {
173        test_scalar_function!(
174            SparkConcat::new(),
175            vec![
176                ColumnarValue::Scalar(ScalarValue::Utf8(Some("Spark".to_string()))),
177                ColumnarValue::Scalar(ScalarValue::Utf8(Some("SQL".to_string()))),
178                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
179            ],
180            Ok(None),
181            &str,
182            DataType::Utf8,
183            StringArray
184        );
185        Ok(())
186    }
187
188    #[test]
189    fn test_spark_concat_return_field_non_nullable() -> Result<()> {
190        let func = SparkConcat::new();
191
192        let fields = vec![
193            Arc::new(Field::new("a", DataType::Utf8, false)),
194            Arc::new(Field::new("b", DataType::Utf8, false)),
195        ];
196
197        let args = ReturnFieldArgs {
198            arg_fields: &fields,
199            scalar_arguments: &[],
200        };
201
202        let field = func.return_field_from_args(args)?;
203
204        assert!(
205            !field.is_nullable(),
206            "Expected concat result to be non-nullable when all inputs are non-nullable"
207        );
208
209        Ok(())
210    }
211    #[test]
212    fn test_spark_concat_return_field_nullable() -> Result<()> {
213        let func = SparkConcat::new();
214
215        let fields = vec![
216            Arc::new(Field::new("a", DataType::Utf8, false)),
217            Arc::new(Field::new("b", DataType::Utf8, true)),
218        ];
219
220        let args = ReturnFieldArgs {
221            arg_fields: &fields,
222            scalar_arguments: &[],
223        };
224
225        let field = func.return_field_from_args(args)?;
226
227        assert!(
228            field.is_nullable(),
229            "Expected concat result to be nullable when any input is nullable"
230        );
231
232        Ok(())
233    }
234}
datafusion_spark/function/string/concat.rs

datafusion_spark/function/string/
concat.rs
