Update the CONCAT scalar function to support Utf8View

datafusion/functions/src/string/common.rs

@@ -255,22 +255,29 @@ pub(crate) enum ColumnarValueRef<'a> {
     Scalar(&'a [u8]),
     NullableArray(&'a StringArray),
     NonNullableArray(&'a StringArray),
+    NullableStringViewArray(&'a StringViewArray),
+    NonNullableStringViewArray(&'a StringViewArray),
 }
 
 impl<'a> ColumnarValueRef<'a> {
     #[inline]
     pub fn is_valid(&self, i: usize) -> bool {
         match &self {
             Self::Scalar(_) | Self::NonNullableArray(_) => true,
+            Self::NonNullableStringViewArray(_) => true,
             Self::NullableArray(array) => array.is_valid(i),
+            Self::NullableStringViewArray(array) => array.is_valid(i),
         }
     }
 
     #[inline]
     pub fn nulls(&self) -> Option<NullBuffer> {
         match &self {
-            Self::Scalar(_) | Self::NonNullableArray(_) => None,
+            Self::Scalar(_)
+            | Self::NonNullableArray(_)
+            | Self::NonNullableStringViewArray(_) => None,
             Self::NullableArray(array) => array.nulls().cloned(),
+            Self::NullableStringViewArray(array) => array.nulls().cloned(),
         }
     }
 }
@@ -389,10 +396,20 @@ impl StringArrayBuilder {
                         .extend_from_slice(array.value(i).as_bytes());
                 }
             }
+            ColumnarValueRef::NullableStringViewArray(array) => {
+                if !CHECK_VALID || array.is_valid(i) {
+                    self.value_buffer
+                        .extend_from_slice(array.value(i).as_bytes());
+                }
+            }
             ColumnarValueRef::NonNullableArray(array) => {
                 self.value_buffer
                     .extend_from_slice(array.value(i).as_bytes());
             }
+            ColumnarValueRef::NonNullableStringViewArray(array) => {
+                self.value_buffer
+                    .extend_from_slice(array.value(i).as_bytes());
+            }
         }
     }
 

datafusion/functions/src/string/concat.rs

@@ -15,14 +15,13 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::array::{Array, StringViewArray};
+use arrow::datatypes::DataType;
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::datatypes::DataType;
-use arrow::datatypes::DataType::Utf8;
-
-use datafusion_common::cast::as_string_array;
-use datafusion_common::{internal_err, Result, ScalarValue};
+use datafusion_common::cast::{as_string_array, as_string_view_array};
+use datafusion_common::{internal_err, plan_err, Result, ScalarValue};
 use datafusion_expr::expr::ScalarFunction;
 use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
 use datafusion_expr::{lit, ColumnarValue, Expr, Volatility};
@@ -46,7 +45,7 @@ impl ConcatFunc {
     pub fn new() -> Self {
         use DataType::*;
         Self {
-            signature: Signature::variadic(vec![Utf8], Volatility::Immutable),
+            signature: Signature::variadic(vec![Utf8, Utf8View], Volatility::Immutable),
         }
     }
 }
@@ -64,13 +63,18 @@ impl ScalarUDFImpl for ConcatFunc {
         &self.signature
     }
 
-    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-        Ok(Utf8)
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        use DataType::*;
+        Ok(match &arg_types[0] {
+            Utf8View => Utf8View,
+            _ => Utf8,
+        })
     }
 
     /// Concatenates the text representations of all the arguments. NULL arguments are ignored.
     /// concat('abcde', 2, NULL, 22) = 'abcde222'
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+        let args_datatype = args[0].data_type();
         let array_len = args
             .iter()
             .filter_map(|x| match x {
@@ -87,7 +91,18 @@ impl ScalarUDFImpl for ConcatFunc {
                     result.push_str(v);
                 }
             }
-            return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))));
+
+            return match args_datatype {
+                DataType::Utf8View => {
+                    Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result))))
+                }
+                DataType::Utf8 => {
+                    Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))))
+                }
+                other => {
+                    plan_err!("Concat function does not support datatype of {other}")
+                }
+            };
         }
 
         // Array
@@ -103,15 +118,40 @@ impl ScalarUDFImpl for ConcatFunc {
                         columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
                     }
                 }
+                ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => {
+                    if let Some(s) = maybe_value {
+                        data_size += s.len() * len;
+                        columns.push(ColumnarValueRef::Scalar(s.as_bytes()));
+                    }
+                }
                 ColumnarValue::Array(array) => {
-                    let string_array = as_string_array(array)?;
-                    data_size += string_array.values().len();
-                    let column = if array.is_nullable() {
-                        ColumnarValueRef::NullableArray(string_array)
-                    } else {
-                        ColumnarValueRef::NonNullableArray(string_array)
+                    match array.data_type() {
+                        DataType::Utf8 | DataType::LargeUtf8 => {
+                            let string_array = as_string_array(array)?;
+
+                            data_size += string_array.values().len();
+                            let column = if array.is_nullable() {
+                                ColumnarValueRef::NullableArray(string_array)
+                            } else {
+                                ColumnarValueRef::NonNullableArray(string_array)
+                            };
+                            columns.push(column);
+                        },
+                        DataType::Utf8View => {
+                            let string_array = as_string_view_array(array)?;
+
+                            data_size += string_array.len();
+                            let column = if array.is_nullable() {
+                                ColumnarValueRef::NullableStringViewArray(string_array)
+                            } else {
+                                ColumnarValueRef::NonNullableStringViewArray(string_array)
+                            };
+                            columns.push(column);
+                        },
+                        other => {
+                            return plan_err!("Input was {other} which is not a supported datatype for concat function")
+                        }
                     };
-                    columns.push(column);
                 }
                 _ => unreachable!(),
             }
@@ -124,7 +164,20 @@ impl ScalarUDFImpl for ConcatFunc {
                 .for_each(|column| builder.write::<true>(column, i));
             builder.append_offset();
         }
-        Ok(ColumnarValue::Array(Arc::new(builder.finish(None))))
+        let string_array = builder.finish(None);
+
+        match args_datatype {
+            DataType::Utf8 | DataType::LargeUtf8 => {
+                Ok(ColumnarValue::Array(Arc::new(string_array)))
+            }
+            DataType::Utf8View => {
+                let string_array_iter = string_array.into_iter();
+                Ok(ColumnarValue::Array(Arc::new(StringViewArray::from_iter(
+                    string_array_iter,
+                ))))
+            }
+            _ => unreachable!(),
+        }
     }
 
     /// Simplify the `concat` function by
@@ -151,11 +204,11 @@ pub fn simplify_concat(args: Vec<Expr>) -> Result<ExprSimplifyResult> {
     for arg in args.clone() {
         match arg {
             // filter out `null` args
-            Expr::Literal(ScalarValue::Utf8(None) | ScalarValue::LargeUtf8(None)) => {}
+            Expr::Literal(ScalarValue::Utf8(None) | ScalarValue::LargeUtf8(None) | ScalarValue::Utf8View(None)) => {}
             // All literals have been converted to Utf8 or LargeUtf8 in type_coercion.
             // Concatenate it with the `contiguous_scalar`.
             Expr::Literal(
-                ScalarValue::Utf8(Some(v)) | ScalarValue::LargeUtf8(Some(v)),
+                ScalarValue::Utf8(Some(v)) | ScalarValue::LargeUtf8(Some(v)) | ScalarValue::Utf8View(Some(v)),
             ) => contiguous_scalar += &v,
             Expr::Literal(x) => {
                 return internal_err!(
@@ -197,6 +250,7 @@ mod tests {
     use crate::utils::test::test_function;
     use arrow::array::Array;
     use arrow::array::{ArrayRef, StringArray};
+    use DataType::*;
 
     #[test]
     fn test_functions() -> Result<()> {

datafusion/sqllogictest/test_files/string_view.slt

@@ -776,7 +776,7 @@ EXPLAIN SELECT
 FROM test;
 ----
 logical_plan
-01)Projection: concat(CAST(test.column1_utf8view AS Utf8), CAST(test.column2_utf8view AS Utf8)) AS c
+01)Projection: concat(test.column1_utf8view, test.column2_utf8view) AS c
 02)--TableScan: test projection=[column1_utf8view, column2_utf8view]
 
 ## Ensure no casts for CONCAT_WS
@@ -863,6 +863,39 @@ XIANGPENG
 RAPHAEL
 NULL
 
+## Should run CONCAT successfully
+query T
+SELECT
+  concat(column1_utf8view, column2_utf8view) as c
+FROM test;
+----
+AndrewX
+XiangpengXiangpeng
+RaphaelR
+R
+
+## Should run CONCAT successfully with utf8 and utf8view
+query T
+SELECT
+  concat(column1_utf8view, column2_utf8) as c
+FROM test;
+----
+AndrewX
+XiangpengXiangpeng
+RaphaelR
+R
+
+## Should run CONCAT successfully with utf8 utf8view and largeutf8
+query T
+SELECT
+  concat(column1_utf8view, column2_utf8, column2_large_utf8) as c
+FROM test;
+----
+AndrewXX
+XiangpengXiangpengXiangpeng
+RaphaelRR
+RR
+
 ## Ensure no casts for LPAD
 query TT
 EXPLAIN SELECT
@@ -1307,3 +1340,4 @@ select column2|| ' ' ||column3 from temp;
 ----
 rust fast
 datafusion cool
+