blaze: Revert "Improve StringView support for SUBSTR (apache#12044)", This reverts commit 7ead2ad

Author: richox

datafusion/functions/src/unicode/substr.rs

@@ -18,21 +18,21 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use crate::strings::make_and_append_view;
-use crate::utils::make_scalar_function;
 use arrow::array::{
-    Array, ArrayIter, ArrayRef, AsArray, Int64Array, NullBufferBuilder, StringArrayType,
-    StringViewArray, StringViewBuilder,
+    Array, ArrayIter, ArrayRef, AsArray, Int64Array, OffsetSizeTrait,
+    StringArrayType, StringViewBuilder,
 };
-use arrow::buffer::ScalarBuffer;
 use arrow::datatypes::DataType;
+
 use datafusion_common::cast::as_int64_array;
 use datafusion_common::{exec_err, plan_err, Result};
 use datafusion_expr::{
     ColumnarValue, Documentation, ScalarUDFImpl, Signature, Volatility,
 };
 use datafusion_macros::user_doc;
 
+use crate::utils::{make_scalar_function, utf8_to_str_type};
+
 #[user_doc(
     doc_section(label = "String Functions"),
     description = "Extracts a substring of a specified number of characters from a specific starting position in a string.",
@@ -44,7 +44,7 @@ use datafusion_macros::user_doc;
 | substr(Utf8("datafusion"),Int64(5),Int64(3)) |
 +----------------------------------------------+
 | fus                                          |
-+----------------------------------------------+ 
++----------------------------------------------+
 ```"#,
     standard_argument(name = "str", prefix = "String"),
     argument(
@@ -90,9 +90,8 @@ impl ScalarUDFImpl for SubstrFunc {
         &self.signature
     }
 
-    // `SubstrFunc` always generates `Utf8View` output for its efficiency.
-    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
-        Ok(DataType::Utf8View)
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        utf8_to_str_type(&arg_types[0], "substr")
     }
 
     fn invoke_with_args(
@@ -177,28 +176,21 @@ impl ScalarUDFImpl for SubstrFunc {
     }
 }
 
-/// Extracts the substring of string starting at the start'th character, and extending for count characters if that is specified. (Same as substring(string from start for count).)
-/// substr('alphabet', 3) = 'phabet'
-/// substr('alphabet', 3, 2) = 'ph'
-/// The implementation uses UTF-8 code points as characters
 pub fn substr(args: &[ArrayRef]) -> Result<ArrayRef> {
     match args[0].data_type() {
         DataType::Utf8 => {
             let string_array = args[0].as_string::<i32>();
-            string_substr::<_>(string_array, &args[1..])
+            calculate_substr::<_, i32>(string_array, &args[1..])
         }
         DataType::LargeUtf8 => {
             let string_array = args[0].as_string::<i64>();
-            string_substr::<_>(string_array, &args[1..])
+            calculate_substr::<_, i64>(string_array, &args[1..])
         }
         DataType::Utf8View => {
             let string_array = args[0].as_string_view();
-            string_view_substr(string_array, &args[1..])
+            calculate_substr::<_, i32>(string_array, &args[1..])
         }
-        other => exec_err!(
-            "Unsupported data type {other:?} for function substr,\
-            expected Utf8View, Utf8 or LargeUtf8."
-        ),
+        other => exec_err!("Unsupported data type {other:?} for function substr"),
     }
 }
 
@@ -312,120 +304,11 @@ fn enable_ascii_fast_path<'a, V: StringArrayType<'a>>(
     }
 }
 
-// The decoding process refs the trait at: arrow/arrow-data/src/byte_view.rs:44
-// From<u128> for ByteView
-fn string_view_substr(
-    string_view_array: &StringViewArray,
-    args: &[ArrayRef],
-) -> Result<ArrayRef> {
-    let mut views_buf = Vec::with_capacity(string_view_array.len());
-    let mut null_builder = NullBufferBuilder::new(string_view_array.len());
-
-    let start_array = as_int64_array(&args[0])?;
-    let count_array_opt = if args.len() == 2 {
-        Some(as_int64_array(&args[1])?)
-    } else {
-        None
-    };
-
-    let enable_ascii_fast_path =
-        enable_ascii_fast_path(&string_view_array, start_array, count_array_opt);
-
-    // In either case of `substr(s, i)` or `substr(s, i, cnt)`
-    // If any of input argument is `NULL`, the result is `NULL`
-    match args.len() {
-        1 => {
-            for ((str_opt, raw_view), start_opt) in string_view_array
-                .iter()
-                .zip(string_view_array.views().iter())
-                .zip(start_array.iter())
-            {
-                if let (Some(str), Some(start)) = (str_opt, start_opt) {
-                    let (start, end) =
-                        get_true_start_end(str, start, None, enable_ascii_fast_path);
-                    let substr = &str[start..end];
-
-                    make_and_append_view(
-                        &mut views_buf,
-                        &mut null_builder,
-                        raw_view,
-                        substr,
-                        start as u32,
-                    );
-                } else {
-                    null_builder.append_null();
-                    views_buf.push(0);
-                }
-            }
-        }
-        2 => {
-            let count_array = count_array_opt.unwrap();
-            for (((str_opt, raw_view), start_opt), count_opt) in string_view_array
-                .iter()
-                .zip(string_view_array.views().iter())
-                .zip(start_array.iter())
-                .zip(count_array.iter())
-            {
-                if let (Some(str), Some(start), Some(count)) =
-                    (str_opt, start_opt, count_opt)
-                {
-                    if count < 0 {
-                        return exec_err!(
-                            "negative substring length not allowed: substr(<str>, {start}, {count})"
-                        );
-                    } else {
-                        if start == i64::MIN {
-                            return exec_err!(
-                                "negative overflow when calculating skip value"
-                            );
-                        }
-                        let (start, end) = get_true_start_end(
-                            str,
-                            start,
-                            Some(count as u64),
-                            enable_ascii_fast_path,
-                        );
-                        let substr = &str[start..end];
-
-                        make_and_append_view(
-                            &mut views_buf,
-                            &mut null_builder,
-                            raw_view,
-                            substr,
-                            start as u32,
-                        );
-                    }
-                } else {
-                    null_builder.append_null();
-                    views_buf.push(0);
-                }
-            }
-        }
-        other => {
-            return exec_err!(
-                "substr was called with {other} arguments. It requires 2 or 3."
-            )
-        }
-    }
-
-    let views_buf = ScalarBuffer::from(views_buf);
-    let nulls_buf = null_builder.finish();
-
-    // Safety:
-    // (1) The blocks of the given views are all provided
-    // (2) Each of the range `view.offset+start..end` of view in views_buf is within
-    // the bounds of each of the blocks
-    unsafe {
-        let array = StringViewArray::new_unchecked(
-            views_buf,
-            string_view_array.data_buffers().to_vec(),
-            nulls_buf,
-        );
-        Ok(Arc::new(array) as ArrayRef)
-    }
-}
-
-fn string_substr<'a, V>(string_array: V, args: &[ArrayRef]) -> Result<ArrayRef>
+/// Extracts the substring of string starting at the start'th character, and extending for count characters if that is specified. (Same as substring(string from start for count).)
+/// substr('alphabet', 3) = 'phabet'
+/// substr('alphabet', 3, 2) = 'ph'
+/// The implementation uses UTF-8 code points as characters
+fn calculate_substr<'a, V, T>(string_array: V, args: &[ArrayRef]) -> Result<ArrayRef>
 where
     V: StringArrayType<'a>,
 {
@@ -507,8 +390,8 @@ where
 
 #[cfg(test)]
 mod tests {
-    use arrow::array::{Array, StringViewArray};
-    use arrow::datatypes::DataType::Utf8View;
+    use arrow::array::{Array, StringArray};
+    use arrow::datatypes::DataType::Utf8;
 
     use datafusion_common::{exec_err, Result, ScalarValue};
     use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
@@ -526,8 +409,8 @@ mod tests {
             ],
             Ok(None),
             &str,
-            Utf8View,
-            StringViewArray
+            Utf8,
+            StringArray
         );
         test_function!(
             SubstrFunc::new(),
@@ -539,35 +422,8 @@ mod tests {
             ],
             Ok(Some("alphabet")),
             &str,
-            Utf8View,
-            StringViewArray
-        );
-        test_function!(
-            SubstrFunc::new(),
-            vec![
-                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
-                    "this és longer than 12B"
-                )))),
-                ColumnarValue::Scalar(ScalarValue::from(5i64)),
-                ColumnarValue::Scalar(ScalarValue::from(2i64)),
-            ],
-            Ok(Some(" é")),
-            &str,
-            Utf8View,
-            StringViewArray
-        );
-        test_function!(
-            SubstrFunc::new(),
-            vec![
-                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
-                    "this is longer than 12B"
-                )))),
-                ColumnarValue::Scalar(ScalarValue::from(5i64)),
-            ],
-            Ok(Some(" is longer than 12B")),
-            &str,
-            Utf8View,
-            StringViewArray
+            Utf8,
+            StringArray
         );
         test_function!(
             SubstrFunc::new(),
@@ -579,8 +435,8 @@ mod tests {
             ],
             Ok(Some("ésoj")),
             &str,
-            Utf8View,
-            StringViewArray
+            Utf8,
+            StringArray
         );
         test_function!(
             SubstrFunc::new(),
@@ -593,8 +449,8 @@ mod tests {
             ],
             Ok(Some("ph")),
             &str,
-            Utf8View,
-            StringViewArray
+            Utf8,
+            StringArray
         );
         test_function!(
             SubstrFunc::new(),
@@ -607,8 +463,8 @@ mod tests {
             ],
             Ok(Some("phabet")),
             &str,
-            Utf8View,
-            StringViewArray
+            Utf8,
+            StringArray
         );
         test_function!(
             SubstrFunc::new(),