Improve performance of trim for string view (10%)

datafusion/functions/src/string/btrim.rs

@@ -82,7 +82,11 @@ impl ScalarUDFImpl for BTrimFunc {
     }
 
     fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        utf8_to_str_type(&arg_types[0], "btrim")
+        if arg_types[0] == DataType::Utf8View {
+            Ok(DataType::Utf8View)
+        } else {
+            utf8_to_str_type(&arg_types[0], "btrim")
+        }
     }
 
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
@@ -106,3 +110,133 @@ impl ScalarUDFImpl for BTrimFunc {
         &self.aliases
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use arrow::array::{Array, StringArray, StringViewArray};
+    use arrow::datatypes::DataType::{Utf8, Utf8View};
+
+    use datafusion_common::{Result, ScalarValue};
+    use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+    use crate::string::btrim::BTrimFunc;
+    use crate::utils::test::test_function;
+
+    #[test]
+    fn test_functions() {
+        test_function!(
+            BTrimFunc::new(),
+            &[ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
+                String::from("alphabet  ")
+            ))),],
+            Ok(Some("alphabet")),
+            &str,
+            Utf8View,
+            StringViewArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
+                String::from("  alphabet  ")
+            ))),],
+            Ok(Some("alphabet")),
+            &str,
+            Utf8View,
+            StringViewArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
+                    "alphabet"
+                )))),
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from("t")))),
+            ],
+            Ok(Some("alphabe")),
+            &str,
+            Utf8View,
+            StringViewArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
+                    "alphabet"
+                )))),
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
+                    "alphabe"
+                )))),
+            ],
+            Ok(Some("t")),
+            &str,
+            Utf8View,
+            StringViewArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8View(Some(String::from(
+                    "alphabet"
+                )))),
+                ColumnarValue::Scalar(ScalarValue::Utf8View(None)),
+            ],
+            Ok(None),
+            &str,
+            Utf8View,
+            StringViewArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+                String::from("alphabet  ")
+            ))),],
+            Ok(Some("alphabet")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[ColumnarValue::Scalar(ScalarValue::Utf8(Some(
+                String::from("alphabet  ")
+            ))),],
+            Ok(Some("alphabet")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("t")))),
+            ],
+            Ok(Some("alphabe")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabe")))),
+            ],
+            Ok(Some("t")),
+            &str,
+            Utf8,
+            StringArray
+        );
+        test_function!(
+            BTrimFunc::new(),
+            &[
+                ColumnarValue::Scalar(ScalarValue::Utf8(Some(String::from("alphabet")))),
+                ColumnarValue::Scalar(ScalarValue::Utf8(None)),
+            ],
+            Ok(None),
+            &str,
+            Utf8,
+            StringArray
+        );
+    }
+}

datafusion/functions/src/string/common.rs

@@ -21,17 +21,39 @@ use std::fmt::{Display, Formatter};
 use std::sync::Arc;
 
 use arrow::array::{
-    new_null_array, Array, ArrayAccessor, ArrayDataBuilder, ArrayIter, ArrayRef,
-    GenericStringArray, GenericStringBuilder, LargeStringArray, OffsetSizeTrait,
-    StringArray, StringBuilder, StringViewArray, StringViewBuilder,
+    make_view, new_null_array, Array, ArrayAccessor, ArrayDataBuilder, ArrayIter,
+    ArrayRef, ByteView, GenericStringArray, GenericStringBuilder, LargeStringArray,
+    OffsetSizeTrait, StringArray, StringBuilder, StringViewArray, StringViewBuilder,
 };
 use arrow::buffer::{Buffer, MutableBuffer, NullBuffer};
 use arrow::datatypes::DataType;
+use arrow_buffer::{NullBufferBuilder, ScalarBuffer};
 use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
 use datafusion_common::Result;
 use datafusion_common::{exec_err, ScalarValue};
 use datafusion_expr::ColumnarValue;
 
+/// Make a `u128` based on the given substr, start(offset to view.offset), and
+/// push into to the given buffers
+pub(crate) fn make_and_append_view(
+    views_buffer: &mut Vec<u128>,
+    null_builder: &mut NullBufferBuilder,
+    raw_view: &u128,
+    substr: &str,
+    start: u32,
+) {
+    let substr_len = substr.len();
+    let sub_view = if substr_len > 12 {
+        let view = ByteView::from(*raw_view);
+        make_view(substr.as_bytes(), view.buffer_index, view.offset + start)
+    } else {
+        // inline value does not need block id or offset
+        make_view(substr.as_bytes(), 0, 0)
+    };
+    views_buffer.push(sub_view);
+    null_builder.append_non_null();
+}
+
 pub(crate) enum TrimType {
     Left,
     Right,
@@ -72,65 +94,126 @@ pub(crate) fn general_trim<T: OffsetSizeTrait>(
     };
 
     if use_string_view {
-        string_view_trim::<T>(func, args)
+        string_view_trim(func, args)
     } else {
         string_trim::<T>(func, args)
     }
 }
 
 // removing 'a will cause compiler complaining lifetime of `func`
-fn string_view_trim<'a, T: OffsetSizeTrait>(
-    func: fn(&'a str, &'a str) -> &'a str,
+fn string_view_trim<'a>(
+    trim_func: fn(&'a str, &'a str) -> &'a str,
     args: &'a [ArrayRef],
 ) -> Result<ArrayRef> {
-    let string_array = as_string_view_array(&args[0])?;
+    let string_view_array = as_string_view_array(&args[0])?;
+    let mut views_buf = Vec::with_capacity(string_view_array.len());
+    let mut null_builder = NullBufferBuilder::new(string_view_array.len());
 
     match args.len() {
         1 => {
-            let result = string_array
-                .iter()
-                .map(|string| string.map(|string: &str| func(string, " ")))
-                .collect::<GenericStringArray<T>>();
-
-            Ok(Arc::new(result) as ArrayRef)
+            let array_iter = string_view_array.iter();
+            let views_iter = string_view_array.views().iter();
+            for (src_str_opt, raw_view) in array_iter.zip(views_iter) {
+                trim_and_append_str(
+                    src_str_opt,
+                    Some(" "),
+                    trim_func,
+                    &mut views_buf,
+                    &mut null_builder,
+                    raw_view,
+                );
+            }
         }
         2 => {
             let characters_array = as_string_view_array(&args[1])?;
 
             if characters_array.len() == 1 {
+                // Only one `trim characters` exist
                 if characters_array.is_null(0) {
                     return Ok(new_null_array(
                         // The schema is expecting utf8 as null
-                        &DataType::Utf8,
-                        string_array.len(),
+                        &DataType::Utf8View,
+                        string_view_array.len(),
                     ));
                 }
 
                 let characters = characters_array.value(0);
-                let result = string_array
-                    .iter()
-                    .map(|item| item.map(|string| func(string, characters)))
-                    .collect::<GenericStringArray<T>>();
-                return Ok(Arc::new(result) as ArrayRef);
+                let array_iter = string_view_array.iter();
+                let views_iter = string_view_array.views().iter();
+                for (src_str_opt, raw_view) in array_iter.zip(views_iter) {
+                    trim_and_append_str(
+                        src_str_opt,
+                        Some(characters),
+                        trim_func,
+                        &mut views_buf,
+                        &mut null_builder,
+                        raw_view,
+                    );
+                }
+            } else {
+                // A specific `trim characters` for a row in the string view array
+                let characters_iter = characters_array.iter();
+                let array_iter = string_view_array.iter();
+                let views_iter = string_view_array.views().iter();
+                for ((src_str_opt, raw_view), characters_opt) in
+                    array_iter.zip(views_iter).zip(characters_iter)
+                {
+                    trim_and_append_str(
+                        src_str_opt,
+                        characters_opt,
+                        trim_func,
+                        &mut views_buf,
+                        &mut null_builder,
+                        raw_view,
+                    );
+                }
             }
-
-            let result = string_array
-                .iter()
-                .zip(characters_array.iter())
-                .map(|(string, characters)| match (string, characters) {
-                    (Some(string), Some(characters)) => Some(func(string, characters)),
-                    _ => None,
-                })
-                .collect::<GenericStringArray<T>>();
-
-            Ok(Arc::new(result) as ArrayRef)
         }
         other => {
-            exec_err!(
+            return exec_err!(
             "Function TRIM was called with {other} arguments. It requires at least 1 and at most 2."
-            )
+            );
         }
     }
+
+    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 trim_and_append_str<'a>(
+    src_str_opt: Option<&'a str>,
+    trim_characters_opt: Option<&'a str>,
+    trim_func: fn(&'a str, &'a str) -> &'a str,
+    views_buf: &mut Vec<u128>,
+    null_builder: &mut NullBufferBuilder,
+    raw: &u128,
+) {
+    if let (Some(src_str), Some(characters)) = (src_str_opt, trim_characters_opt) {
+        let trim_str = trim_func(src_str, characters);
+
+        // Safety:
+        // `trim_str` is computed from `str::trim_xxx_matches`,
+        // and its addr is ensured to be >= `origin_str`'s
+        let start = unsafe { trim_str.as_ptr().offset_from(src_str.as_ptr()) as u32 };
+
+        make_and_append_view(views_buf, null_builder, raw, trim_str, start);
+    } else {
+        null_builder.append_null();
+        views_buf.push(0);
+    }
 }
 
 fn string_trim<'a, T: OffsetSizeTrait>(