perf: Add append_with to string builders, use in replace

datafusion/functions/src/string/replace.rs

@@ -21,7 +21,9 @@ use arrow::array::{Array, ArrayRef, OffsetSizeTrait};
 use arrow::buffer::NullBuffer;
 use arrow::datatypes::DataType;
 
-use crate::strings::GenericStringArrayBuilder;
+use crate::strings::{
+    BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringWriter,
+};
 use crate::utils::{make_scalar_function, utf8_to_str_type};
 use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
 use datafusion_common::types::logical_string;
@@ -164,7 +166,6 @@ fn replace_view(args: &[ArrayRef]) -> Result<ArrayRef> {
 
     let len = string_array.len();
     let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(len, 0);
-    let mut buffer = String::new();
     let nulls = NullBuffer::union(
         NullBuffer::union(string_array.nulls(), from_array.nulls()).as_ref(),
         to_array.nulls(),
@@ -183,19 +184,15 @@ fn replace_view(args: &[ArrayRef]) -> Result<ArrayRef> {
             let string = unsafe { string_array.value_unchecked(i) };
             let from = unsafe { from_array.value_unchecked(i) };
             let to = unsafe { to_array.value_unchecked(i) };
-            buffer.clear();
-            replace_into_string(&mut buffer, string, from, to);
-            builder.append_value(&buffer);
+            apply_replace(&mut builder, string, from, to);
         }
     } else {
         for i in 0..len {
             // SAFETY: i < len, and no input has a null buffer.
             let string = unsafe { string_array.value_unchecked(i) };
             let from = unsafe { from_array.value_unchecked(i) };
             let to = unsafe { to_array.value_unchecked(i) };
-            buffer.clear();
-            replace_into_string(&mut buffer, string, from, to);
-            builder.append_value(&buffer);
+            apply_replace(&mut builder, string, from, to);
         }
     }
 
@@ -211,7 +208,6 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
 
     let len = string_array.len();
     let mut builder = GenericStringArrayBuilder::<T>::with_capacity(len, 0);
-    let mut buffer = String::new();
     let nulls = NullBuffer::union(
         NullBuffer::union(string_array.nulls(), from_array.nulls()).as_ref(),
         to_array.nulls(),
@@ -230,71 +226,77 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
             let string = unsafe { string_array.value_unchecked(i) };
             let from = unsafe { from_array.value_unchecked(i) };
             let to = unsafe { to_array.value_unchecked(i) };
-            buffer.clear();
-            replace_into_string(&mut buffer, string, from, to);
-            builder.append_value(&buffer);
+            apply_replace(&mut builder, string, from, to);
         }
     } else {
         for i in 0..len {
             // SAFETY: i < len, and no input has a null buffer.
             let string = unsafe { string_array.value_unchecked(i) };
             let from = unsafe { from_array.value_unchecked(i) };
             let to = unsafe { to_array.value_unchecked(i) };
-            buffer.clear();
-            replace_into_string(&mut buffer, string, from, to);
-            builder.append_value(&buffer);
+            apply_replace(&mut builder, string, from, to);
         }
     }
 
     Ok(Arc::new(builder.finish(nulls)?) as ArrayRef)
 }
 
-/// Helper function to perform string replacement into a reusable String buffer
 #[inline]
-fn replace_into_string(buffer: &mut String, string: &str, from: &str, to: &str) {
+fn apply_replace<B: BulkNullStringArrayBuilder>(
+    builder: &mut B,
+    string: &str,
+    from: &str,
+    to: &str,
+) {
     if from.is_empty() {
-        // When from is empty, insert 'to' at the beginning, between each character, and at the end
-        // This matches the behavior of str::replace()
-        buffer.push_str(to);
-        for ch in string.chars() {
-            buffer.push(ch);
-            buffer.push_str(to);
-        }
+        // Empty `from`: insert `to` before each character and at both ends
+        builder.append_with(|w| {
+            w.write_str(to);
+            for ch in string.chars() {
+                w.write_char(ch);
+                w.write_str(to);
+            }
+        });
         return;
     }
 
-    // Fast path for replacing a single ASCII character with another single ASCII character.
-    // Extends the buffer's underlying Vec<u8> directly, for performance.
+    // Fast-path for replacing one ASCII byte with another
     if let ([from_byte], [to_byte]) = (from.as_bytes(), to.as_bytes())
         && from_byte.is_ascii()
         && to_byte.is_ascii()
     {
-        // SAFETY: Replacing an ASCII byte with another ASCII byte preserves UTF-8 validity.
+        let from_byte = *from_byte;
+        let to_byte = *to_byte;
+        // SAFETY: an ASCII byte (< 0x80) cannot appear inside a multi-byte
+        // UTF-8 sequence, so the closure below only replaces ASCII bytes.
+        // `string` might be UTF-8, but any multi-byte sequences pass through
+        // unchanged. Output is valid UTF-8.
         unsafe {
-            buffer.as_mut_vec().extend(
-                string
-                    .as_bytes()
-                    .iter()
-                    .map(|&b| if b == *from_byte { *to_byte } else { b }),
-            );
+            builder.append_byte_map(string.as_bytes(), |b| {
+                if b == from_byte { to_byte } else { b }
+            });
         }
         return;
     }
 
+    builder.append_with(|w| replace_into_writer(w, string, from, to));
+}
+
+#[inline]
+fn replace_into_writer<W: StringWriter>(w: &mut W, string: &str, from: &str, to: &str) {
     let mut last_end = 0;
     for (start, _part) in string.match_indices(from) {
-        buffer.push_str(&string[last_end..start]);
-        buffer.push_str(to);
+        w.write_str(&string[last_end..start]);
+        w.write_str(to);
         last_end = start + from.len();
     }
-    buffer.push_str(&string[last_end..]);
+    w.write_str(&string[last_end..]);
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::utils::test::test_function;
-    use arrow::array::Array;
     use arrow::array::LargeStringArray;
     use arrow::array::StringArray;
     use arrow::datatypes::DataType::{LargeUtf8, Utf8};