Optimize left, right using Utf8View

Author: neilconway

datafusion/functions/src/unicode/common.rs

@@ -18,11 +18,12 @@
 //! Common utilities for implementing unicode functions
 
 use arrow::array::{
-    Array, ArrayAccessor, ArrayIter, ArrayRef, ByteView, GenericStringArray, Int64Array,
-    OffsetSizeTrait, StringViewArray, make_view,
+    Array, ArrayRef, ByteView, GenericStringArray, Int64Array, OffsetSizeTrait,
+    StringViewArray, make_view,
 };
 use arrow::datatypes::DataType;
 use arrow_buffer::{NullBuffer, ScalarBuffer};
+use datafusion_common::Result;
 use datafusion_common::ScalarValue;
 use datafusion_common::cast::{
     as_generic_string_array, as_int64_array, as_string_view_array,
@@ -99,17 +100,17 @@ fn left_right_byte_length(string: &str, n: i64) -> usize {
 /// General implementation for `left` and `right` functions
 pub(crate) fn general_left_right<F: LeftRightSlicer>(
     args: &[ArrayRef],
-) -> datafusion_common::Result<ArrayRef> {
+) -> Result<ArrayRef> {
     let n_array = as_int64_array(&args[1])?;
 
     match args[0].data_type() {
         DataType::Utf8 => {
             let string_array = as_generic_string_array::<i32>(&args[0])?;
-            general_left_right_array::<i32, _, F>(string_array, n_array)
+            general_left_right_array::<i32, F>(string_array, n_array)
         }
         DataType::LargeUtf8 => {
             let string_array = as_generic_string_array::<i64>(&args[0])?;
-            general_left_right_array::<i64, _, F>(string_array, n_array)
+            general_left_right_array::<i64, F>(string_array, n_array)
         }
         DataType::Utf8View => {
             let string_view_array = as_string_view_array(&args[0])?;
@@ -119,83 +120,125 @@ pub(crate) fn general_left_right<F: LeftRightSlicer>(
     }
 }
 
-/// `general_left_right` implementation for strings
-fn general_left_right_array<
-    'a,
-    T: OffsetSizeTrait,
-    V: ArrayAccessor<Item = &'a str>,
-    F: LeftRightSlicer,
->(
-    string_array: V,
+/// Returns true if all offsets in the array fit in u32, meaning the values
+/// buffer can be referenced by StringView's u32 offset field.
+fn values_fit_in_u32<T: OffsetSizeTrait>(string_array: &GenericStringArray<T>) -> bool {
+    string_array
+        .offsets()
+        .last()
+        .map(|offset| offset.as_usize() <= u32::MAX as usize)
+        .unwrap_or(true)
+}
+
+/// `left`/`right` for Utf8/LargeUtf8 input.
+///
+/// When offsets fit in u32, produces a zero-copy `StringViewArray` with views
+/// pointing into the input values buffer. Otherwise falls back to building a
+/// `StringViewArray` by copying.
+fn general_left_right_array<T: OffsetSizeTrait, F: LeftRightSlicer>(
+    string_array: &GenericStringArray<T>,
     n_array: &Int64Array,
-) -> datafusion_common::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)) => {
-                let range = F::slice(string, n);
-                // Extract a given range from a byte-indexed slice
-                Some(&string[range])
-            }
-            _ => None,
-        })
-        .collect::<GenericStringArray<T>>();
+) -> Result<ArrayRef> {
+    if !values_fit_in_u32(string_array) {
+        let result = string_array
+            .iter()
+            .zip(n_array.iter())
+            .map(|(string, n)| match (string, n) {
+                (Some(string), Some(n)) => Some(&string[F::slice(string, n)]),
+                _ => None,
+            })
+            .collect::<StringViewArray>();
+        return Ok(Arc::new(result) as ArrayRef);
+    }
+
+    let len = string_array.len();
+    let offsets = string_array.value_offsets();
+    let nulls = NullBuffer::union(string_array.nulls(), n_array.nulls());
 
-    Ok(Arc::new(result) as ArrayRef)
+    let mut views_buf = Vec::with_capacity(len);
+    let mut has_out_of_line = false;
+
+    for (i, offset) in offsets.iter().enumerate().take(len) {
+        if nulls.as_ref().is_some_and(|n| !n.is_valid(i)) {
+            views_buf.push(0);
+            continue;
+        }
+
+        // SAFETY: we just checked validity above
+        let string = unsafe { string_array.value_unchecked(i) };
+        let n = n_array.value(i);
+        let range = F::slice(string, n);
+        let result_bytes = &string.as_bytes()[range.clone()];
+        if result_bytes.len() > 12 {
+            has_out_of_line = true;
+        }
+
+        let buf_offset = offset.as_usize() as u32 + range.start as u32;
+        views_buf.push(make_view(result_bytes, 0, buf_offset));
+    }
+
+    let views = ScalarBuffer::from(views_buf);
+    let data_buffers = if has_out_of_line {
+        vec![string_array.values().clone()]
+    } else {
+        vec![]
+    };
+
+    // SAFETY:
+    // - Each view is produced by `make_view` with correct bytes and offset
+    // - Out-of-line views reference buffer index 0, which is the original
+    //   values buffer included in data_buffers when has_out_of_line is true
+    // - values_fit_in_u32 guarantees all offsets fit in u32
+    unsafe {
+        let array = StringViewArray::new_unchecked(views, data_buffers, nulls);
+        Ok(Arc::new(array) as ArrayRef)
+    }
 }
 
-/// `general_left_right` implementation for StringViewArray
+/// `general_left_right` for StringViewArray input.
 fn general_left_right_view<F: LeftRightSlicer>(
     string_view_array: &StringViewArray,
     n_array: &Int64Array,
-) -> datafusion_common::Result<ArrayRef> {
-    let len = n_array.len();
-
+) -> Result<ArrayRef> {
     let views = string_view_array.views();
-    // Every string in StringViewArray has one corresponding view in `views`
-    debug_assert!(views.len() == string_view_array.len());
-
-    // Compose null buffer at once
-    let string_nulls = string_view_array.nulls();
-    let n_nulls = n_array.nulls();
-    let new_nulls = NullBuffer::union(string_nulls, n_nulls);
+    let new_nulls = NullBuffer::union(string_view_array.nulls(), n_array.nulls());
+    let len = n_array.len();
+    let mut has_out_of_line = false;
 
     let new_views = (0..len)
         .map(|idx| {
-            let view = views[idx];
-
-            let is_valid = match &new_nulls {
-                Some(nulls_buf) => nulls_buf.is_valid(idx),
-                None => true,
-            };
-
-            if is_valid {
-                let string: &str = string_view_array.value(idx);
-                let n = n_array.value(idx);
-
-                // Input string comes from StringViewArray, so it should fit in 32-bit length
-                let range = F::slice(string, n);
-                let result_bytes = &string.as_bytes()[range.clone()];
-
-                let byte_view = ByteView::from(view);
-                // New offset starts at 0 for left, and at `range.start` for right,
-                // which is encoded in the given range
-                let new_offset = byte_view.offset + (range.start as u32);
-                // Reuse buffer
-                make_view(result_bytes, byte_view.buffer_index, new_offset)
-            } else {
-                // For nulls, keep the original view
-                view
+            if new_nulls.as_ref().is_some_and(|n| !n.is_valid(idx)) {
+                return 0;
             }
+
+            // SAFETY: we just checked validity above
+            let string: &str = unsafe { string_view_array.value_unchecked(idx) };
+            let n = n_array.value(idx);
+
+            let range = F::slice(string, n);
+            let result_bytes = &string.as_bytes()[range.clone()];
+            if result_bytes.len() > 12 {
+                has_out_of_line = true;
+            }
+
+            let byte_view = ByteView::from(views[idx]);
+            let new_offset = byte_view.offset + (range.start as u32);
+            make_view(result_bytes, byte_view.buffer_index, new_offset)
         })
         .collect::<Vec<u128>>();
 
-    // Buffers are unchanged
-    let result = StringViewArray::try_new(
-        ScalarBuffer::from(new_views),
-        Vec::from(string_view_array.data_buffers()),
-        new_nulls,
-    )?;
-    Ok(Arc::new(result) as ArrayRef)
+    let views = ScalarBuffer::from(new_views);
+    let data_buffers = if has_out_of_line {
+        string_view_array.data_buffers().to_vec()
+    } else {
+        vec![]
+    };
+
+    // SAFETY:
+    // - Each view is produced by `make_view` with correct bytes and offset
+    // - Out-of-line views reuse the original buffer index and adjusted offset
+    unsafe {
+        let array = StringViewArray::new_unchecked(views, data_buffers, new_nulls);
+        Ok(Arc::new(array) as ArrayRef)
+    }
 }

datafusion/functions/src/unicode/left.rs

@@ -79,8 +79,8 @@ impl ScalarUDFImpl for LeftFunc {
         &self.signature
     }
 
-    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
-        Ok(arg_types[0].clone())
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Utf8View)
     }
 
     /// Returns first n characters in the string, or when n is negative, returns all but last |n| characters.
@@ -108,8 +108,8 @@ impl ScalarUDFImpl for LeftFunc {
 
 #[cfg(test)]
 mod tests {
-    use arrow::array::{Array, StringArray, StringViewArray};
-    use arrow::datatypes::DataType::{Utf8, Utf8View};
+    use arrow::array::{Array, StringViewArray};
+    use arrow::datatypes::DataType::Utf8View;
 
     use datafusion_common::{Result, ScalarValue};
     use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
@@ -127,8 +127,8 @@ mod tests {
             ],
             Ok(Some("ab")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -138,8 +138,8 @@ mod tests {
             ],
             Ok(Some("abcde")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -149,8 +149,8 @@ mod tests {
             ],
             Ok(Some("abc")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -160,8 +160,8 @@ mod tests {
             ],
             Ok(Some("")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -171,8 +171,8 @@ mod tests {
             ],
             Ok(Some("")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -182,8 +182,8 @@ mod tests {
             ],
             Ok(Some("")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -193,8 +193,8 @@ mod tests {
             ],
             Ok(None),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -204,8 +204,8 @@ mod tests {
             ],
             Ok(None),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -215,8 +215,8 @@ mod tests {
             ],
             Ok(Some("joséé")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         test_function!(
             LeftFunc::new(),
@@ -226,8 +226,8 @@ mod tests {
             ],
             Ok(Some("joséé")),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
         #[cfg(not(feature = "unicode_expressions"))]
         test_function!(
@@ -240,8 +240,8 @@ mod tests {
                 "function left requires compilation with feature flag: unicode_expressions."
             ),
             &str,
-            Utf8,
-            StringArray
+            Utf8View,
+            StringViewArray
         );
 
         // StringView cases
@@ -307,8 +307,8 @@ mod tests {
                 ],
                 Ok(Some(expected.as_str())),
                 &str,
-                Utf8,
-                StringArray
+                Utf8View,
+                StringViewArray
             );
         }