Implement native support StringView for Ends With (#11924)

Signed-off-by: Chojan Shang <[email protected]>
apache · Aug 12, 2024 · 8deba02 · 8deba02
1 parent 032b9c9
commit 8deba02
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 22 deletions.
diff --git a/datafusion/functions/src/string/ends_with.rs b/datafusion/functions/src/string/ends_with.rs
@@ -18,12 +18,10 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, OffsetSizeTrait};
+use arrow::array::ArrayRef;
 use arrow::datatypes::DataType;
-use arrow::datatypes::DataType::Boolean;
 
-use datafusion_common::cast::as_generic_string_array;
-use datafusion_common::{exec_err, Result};
+use datafusion_common::{internal_err, Result};
 use datafusion_expr::TypeSignature::*;
 use datafusion_expr::{ColumnarValue, Volatility};
 use datafusion_expr::{ScalarUDFImpl, Signature};
@@ -43,14 +41,15 @@ impl Default for EndsWithFunc {
 
 impl EndsWithFunc {
     pub fn new() -> Self {
-        use DataType::*;
         Self {
             signature: Signature::one_of(
                 vec![
-                    Exact(vec![Utf8, Utf8]),
-                    Exact(vec![Utf8, LargeUtf8]),
-                    Exact(vec![LargeUtf8, Utf8]),
-                    Exact(vec![LargeUtf8, LargeUtf8]),
+                    // Planner attempts coercion to the target type starting with the most preferred candidate.
+                    // For example, given input `(Utf8View, Utf8)`, it first tries coercing to `(Utf8View, Utf8View)`.
+                    // If that fails, it proceeds to `(Utf8, Utf8)`.
+                    Exact(vec![DataType::Utf8View, DataType::Utf8View]),
+                    Exact(vec![DataType::Utf8, DataType::Utf8]),
+                    Exact(vec![DataType::LargeUtf8, DataType::LargeUtf8]),
                 ],
                 Volatility::Immutable,
             ),
@@ -72,27 +71,25 @@ impl ScalarUDFImpl for EndsWithFunc {
     }
 
     fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-        Ok(Boolean)
+        Ok(DataType::Boolean)
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
         match args[0].data_type() {
-            DataType::Utf8 => make_scalar_function(ends_with::<i32>, vec![])(args),
-            DataType::LargeUtf8 => make_scalar_function(ends_with::<i64>, vec![])(args),
+            DataType::Utf8View | DataType::Utf8 | DataType::LargeUtf8 => {
+                make_scalar_function(ends_with, vec![])(args)
+            }
             other => {
-                exec_err!("Unsupported data type {other:?} for function ends_with")
+                internal_err!("Unsupported data type {other:?} for function ends_with. Expected Utf8, LargeUtf8 or Utf8View")?
             }
         }
     }
 }
 
 /// Returns true if string ends with suffix.
 /// ends_with('alphabet', 'abet') = 't'
-pub fn ends_with<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
-    let left = as_generic_string_array::<T>(&args[0])?;
-    let right = as_generic_string_array::<T>(&args[1])?;
-
-    let result = arrow::compute::kernels::comparison::ends_with(left, right)?;
+pub fn ends_with(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let result = arrow::compute::kernels::comparison::ends_with(&args[0], &args[1])?;
 
     Ok(Arc::new(result) as ArrayRef)
 }

diff --git a/datafusion/sqllogictest/test_files/string_view.slt b/datafusion/sqllogictest/test_files/string_view.slt
@@ -618,17 +618,15 @@ logical_plan
 03)----TableScan: test projection=[column1_utf8view, column2_utf8view]
 
 ## Ensure no casts for ENDS_WITH
-## TODO https://github.com/apache/datafusion/issues/11852
 query TT
 EXPLAIN SELECT
   ENDS_WITH(column1_utf8view, 'foo') as c1,
   ENDS_WITH(column2_utf8view, column2_utf8view) as c2
 FROM test;
 ----
 logical_plan
-01)Projection: ends_with(CAST(test.column1_utf8view AS Utf8), Utf8("foo")) AS c1, ends_with(__common_expr_1, __common_expr_1) AS c2
-02)--Projection: CAST(test.column2_utf8view AS Utf8) AS __common_expr_1, test.column1_utf8view
-03)----TableScan: test projection=[column1_utf8view, column2_utf8view]
+01)Projection: ends_with(test.column1_utf8view, Utf8View("foo")) AS c1, ends_with(test.column2_utf8view, test.column2_utf8view) AS c2
+02)--TableScan: test projection=[column1_utf8view, column2_utf8view]
 
 ## Ensure no casts for LEVENSHTEIN
 ## TODO https://github.com/apache/datafusion/issues/11854