implement utf8_view for replace

datafusion/functions/src/string/replace.rs

@@ -18,16 +18,16 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait};
+use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait, StringViewArray};
 use arrow::datatypes::DataType;
 
-use datafusion_common::cast::as_generic_string_array;
+use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
 use datafusion_common::{exec_err, Result};
 use datafusion_expr::TypeSignature::*;
 use datafusion_expr::{ColumnarValue, Volatility};
 use datafusion_expr::{ScalarUDFImpl, Signature};
 
-use crate::utils::{make_scalar_function, utf8_to_str_type};
+use crate::utils::make_scalar_function;
 
 #[derive(Debug)]
 pub struct ReplaceFunc {
@@ -45,7 +45,11 @@ impl ReplaceFunc {
         use DataType::*;
         Self {
             signature: Signature::one_of(
-                vec![Exact(vec![Utf8, Utf8, Utf8])],
+                vec![
+                    Exact(vec![Utf8View, Utf8View, Utf8View]), 
+                    Exact(vec![Utf8, Utf8, Utf8]), 
+                    Exact(vec![LargeUtf8, LargeUtf8, LargeUtf8])
+                ],
                 Volatility::Immutable,
             ),
         }
@@ -66,20 +70,46 @@ impl ScalarUDFImpl for ReplaceFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        utf8_to_str_type(&arg_types[0], "replace")
+        match arg_types[0].clone() {
+            DataType::Utf8 => return Ok(DataType::Utf8),
+            DataType::LargeUtf8 => return Ok(DataType::LargeUtf8),
+            DataType::Utf8View => return Ok(DataType::Utf8View),
+            other => {
+                exec_err!("Unsupported data type {other:?} for function replace")
+            }
+        }
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
         match args[0].data_type() {
             DataType::Utf8 => make_scalar_function(replace::<i32>, vec![])(args),
             DataType::LargeUtf8 => make_scalar_function(replace::<i64>, vec![])(args),
+            DataType::Utf8View => make_scalar_function(replace_view, vec![])(args),
             other => {
                 exec_err!("Unsupported data type {other:?} for function replace")
             }
         }
     }
 }
 
+fn replace_view(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let string_array  = as_string_view_array(&args[0])?;
+    let from_array = as_string_view_array(&args[1])?;
+    let to_array = as_string_view_array(&args[2])?;
+
+    let result = string_array
+        .iter()
+        .zip(from_array.iter())
+        .zip(to_array.iter())
+        .map(|((string, from), to)| match (string, from, to) {
+            (Some(string), Some(from), Some(to)) => Some(string.replace(from, to)),
+            _ => None,
+        })
+        .collect::<StringViewArray>();
+
+    Ok(Arc::new(result) as ArrayRef)
+
+}
 /// Replaces all occurrences in string of substring from with substring to.
 /// replace('abcdefabcdef', 'cd', 'XX') = 'abXXefabXXef'
 fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
@@ -100,4 +130,59 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
     Ok(Arc::new(result) as ArrayRef)
 }
 
-mod test {}
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::utils::test::test_function;
+    use arrow::array::Array;
+    use arrow::array::LargeStringArray;
+    use arrow::array::StringArray;
+    use arrow::array::StringViewArray;
+    use arrow::datatypes::DataType::{Utf8, LargeUtf8, Utf8View};
+    use datafusion_common::ScalarValue;
+
+    #[test]
+    fn test_functions() -> Result<()> {
+        test_function!(
+            ReplaceFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("aabbdqcbb")))),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("bb")))),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("ccc")))),
+            ],
+            Ok(Some("aacccdqcccc")),
+            &str,
+            Utf8,
+            StringArray
+        );
+
+        test_function!(
+            ReplaceFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(String::from("aabbb")))),
+                ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(String::from("bbb")))),
+                ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(String::from("cc")))),
+            ],
+            Ok(Some("aacc")),
+            &str,
+            LargeUtf8,
+            LargeStringArray
+        );
+
+        test_function!(
+            ReplaceFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("aabbbcw")))),
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("bb")))),
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("cc")))),
+            ],
+            Ok(Some("aaccbcw")),
+            &str,
+            Utf8View,
+            StringViewArray
+        );
+
+        Ok(())
+    }
+
+}

datafusion/sqllogictest/test_files/string_view.slt

@@ -718,9 +718,8 @@ EXPLAIN SELECT
 FROM test;
 ----
 logical_plan
-01)Projection: replace(__common_expr_1, Utf8("foo"), Utf8("bar")) AS c1, replace(__common_expr_1, CAST(test.column2_utf8view AS Utf8), Utf8("bar")) AS c2
-02)--Projection: CAST(test.column1_utf8view AS Utf8) AS __common_expr_1, test.column2_utf8view
-03)----TableScan: test projection=[column1_utf8view, column2_utf8view]
+01)Projection: replace(test.column1_utf8view, Utf8View("foo"), Utf8View("bar")) AS c1, replace(test.column1_utf8view, test.column2_utf8view, Utf8View("bar")) AS c2
+02)--TableScan: test projection=[column1_utf8view, column2_utf8view]
 
 ## Ensure no casts for REVERSE
 ## TODO file ticket