Add native stringview support for LEFT (apache#11976)

* add string view support for left * add tests for stringview support of LEFT
samuelcolvin · Aug 14, 2024 · b9961c3 · b9961c3
1 parent a00922b
commit b9961c3
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 8 deletions.
diff --git a/datafusion/functions/src/unicode/left.rs b/datafusion/functions/src/unicode/left.rs
@@ -19,10 +19,15 @@ use std::any::Any;
 use std::cmp::Ordering;
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait};
+use arrow::array::{
+    Array, ArrayAccessor, ArrayIter, ArrayRef, GenericStringArray, Int64Array,
+    OffsetSizeTrait,
+};
 use arrow::datatypes::DataType;
 
-use datafusion_common::cast::{as_generic_string_array, as_int64_array};
+use datafusion_common::cast::{
+    as_generic_string_array, as_int64_array, as_string_view_array,
+};
 use datafusion_common::exec_err;
 use datafusion_common::Result;
 use datafusion_expr::TypeSignature::Exact;
@@ -46,7 +51,11 @@ impl LeftFunc {
         use DataType::*;
         Self {
             signature: Signature::one_of(
-                vec![Exact(vec![Utf8, Int64]), Exact(vec![LargeUtf8, Int64])],
+                vec![
+                    Exact(vec![Utf8View, Int64]),
+                    Exact(vec![Utf8, Int64]),
+                    Exact(vec![LargeUtf8, Int64]),
+                ],
                 Volatility::Immutable,
             ),
         }
@@ -72,9 +81,14 @@ impl ScalarUDFImpl for LeftFunc {
 
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
         match args[0].data_type() {
-            DataType::Utf8 => make_scalar_function(left::<i32>, vec![])(args),
+            DataType::Utf8 | DataType::Utf8View => {
+                make_scalar_function(left::<i32>, vec![])(args)
+            }
             DataType::LargeUtf8 => make_scalar_function(left::<i64>, vec![])(args),
-            other => exec_err!("Unsupported data type {other:?} for function left"),
+            other => exec_err!(
+                "Unsupported data type {other:?} for function left,\
+                expected Utf8View, Utf8 or LargeUtf8."
+            ),
         }
     }
 }
@@ -83,10 +97,23 @@ impl ScalarUDFImpl for LeftFunc {
 /// left('abcde', 2) = 'ab'
 /// The implementation uses UTF-8 code points as characters
 pub fn left<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
-    let string_array = as_generic_string_array::<T>(&args[0])?;
     let n_array = as_int64_array(&args[1])?;
-    let result = string_array
-        .iter()
+
+    if args[0].data_type() == &DataType::Utf8View {
+        let string_array = as_string_view_array(&args[0])?;
+        left_impl::<T, _>(string_array, n_array)
+    } else {
+        let string_array = as_generic_string_array::<T>(&args[0])?;
+        left_impl::<T, _>(string_array, n_array)
+    }
+}
+
+fn left_impl<'a, T: OffsetSizeTrait, V: ArrayAccessor<Item = &'a str>>(
+    string_array: V,
+    n_array: &Int64Array,
+) -> Result<ArrayRef> {
+    let iter = ArrayIter::new(string_array);
+    let result = iter
         .zip(n_array.iter())
         .map(|(string, n)| match (string, n) {
             (Some(string), Some(n)) => match n.cmp(&0) {

diff --git a/datafusion/sqllogictest/test_files/string_view.slt b/datafusion/sqllogictest/test_files/string_view.slt
@@ -917,6 +917,29 @@ eng  (empty) ngpeng
 ael  (empty) hael
 NULL NULL    NULL
 
+## Ensure no casts for LEFT
+query TT
+EXPLAIN SELECT
+  LEFT(column1_utf8view, 3) as c2
+FROM test;
+----
+logical_plan
+01)Projection: left(test.column1_utf8view, Int64(3)) AS c2
+02)--TableScan: test projection=[column1_utf8view]
+
+# Test outputs of LEFT
+query TTT
+SELECT
+  LEFT(column1_utf8view, 3) as c1,
+  LEFT(column1_utf8view, 0) as c2,
+  LEFT(column1_utf8view, -3) as c3
+FROM test;
+----
+And  (empty) And
+Xia  (empty) Xiangp
+Rap  (empty) Raph
+NULL NULL    NULL
+
 ## Ensure no casts for RPAD
 ## TODO file ticket
 query TT