perf: Add append_with to string builders, use in replace (#22029)

## Which issue does this PR close?

- Closes #21997 (potentially).

## Rationale for this change

This PR adds two new APIs to `GenericStringArrayBuilder` and
`StringViewArrayBuilder`:

1. `append_with` appends a row whose bytes are produced by invoking a
closure that is passed a `StringWriter`
2. `append_byte_map` appends a row whose bytes are produced by mapping
each byte of the input with a byte-to-byte map closure.

For `StringViewArrayBuilder`, `StringWriter` is an append-only string
writer that switches between writing to a new inline view (for short
strings) or to the in-progress data block automatically. For
`GenericStringArrayBuilder`, `StringWriter` just appends to the value
buffer directly.

(We need two new APIs because `append_byte_map` vectorizes a lot better
than `append_with`, so callers that fit the byte-to-byte map pattern
should prefer it.)

Both of these new APIs allow string UDFs to avoid creating an
intermediate data copy in many cases. To illustrate this, this PR adopts
the new APIs in `replace`.

Benchmarks (Arm64):

  Group 1: ASCII single-byte fast path (StringArray)

  - size=1024 str_len=32 nulls=0.0 : 16.27 µs -> 12.83 µs (−21.1%)
  - size=1024 str_len=32 nulls=0.2 : 14.23 µs -> 12.10 µs (−15.0%)
  - size=1024 str_len=128 nulls=0.0 : 11.28 µs -> 8.21 µs (−27.3%)
  - size=1024 str_len=128 nulls=0.2 : 10.37 µs -> 7.79 µs (−24.9%)
  - size=4096 str_len=32 nulls=0.0 : 62.48 µs -> 49.50 µs (−20.8%)
  - size=4096 str_len=32 nulls=0.2 : 55.74 µs -> 46.66 µs (−16.3%)
  - size=4096 str_len=128 nulls=0.0 : 42.26 µs -> 29.06 µs (−31.2%)
  - size=4096 str_len=128 nulls=0.2 : 39.17 µs -> 28.52 µs (−27.2%)

  Group 2: Multi-byte StringArray — general writer path

  - size=1024 str_len=32 nulls=0.0 : 23.58 µs -> 21.75 µs (−7.8%)
  - size=1024 str_len=32 nulls=0.2 : 18.92 µs -> 17.41 µs (−8.0%)
  - size=1024 str_len=128 nulls=0.0 : 37.56 µs -> 35.33 µs (−5.9%)
  - size=1024 str_len=128 nulls=0.2 : 29.62 µs -> 28.71 µs (−3.1%)
  - size=4096 str_len=32 nulls=0.0 : 97.15 µs -> 88.92 µs (−8.5%)
  - size=4096 str_len=32 nulls=0.2 : 77.03 µs -> 71.43 µs (−7.3%)
  - size=4096 str_len=128 nulls=0.0 : 173.66 µs -> 163.68 µs (−5.7%)
  - size=4096 str_len=128 nulls=0.2 : 134.98 µs -> 128.56 µs (−4.8%)

  Group 3: Multi-byte StringViewArray — general writer path

  - size=1024 str_len=32 nulls=0.0 : 24.46 µs -> 22.18 µs (−9.3%)
  - size=1024 str_len=32 nulls=0.2 : 20.04 µs -> 17.71 µs (−11.7%)
  - size=1024 str_len=128 nulls=0.0 : 36.43 µs -> 35.79 µs (−1.8%)
  - size=1024 str_len=128 nulls=0.2 : 29.73 µs -> 28.70 µs (−3.5%)
  - size=4096 str_len=32 nulls=0.0 : 99.07 µs -> 89.68 µs (−9.5%)
  - size=4096 str_len=32 nulls=0.2 : 84.38 µs -> 72.46 µs (−14.1%)
- size=4096 str_len=128 nulls=0.0 : 169.27 µs -> 164.80 µs (−2.6%, n.s.)
- size=4096 str_len=128 nulls=0.2 : 133.79 µs -> 130.20 µs (−2.7%, n.s.)

  Group 4: Empty-from StringArray

  - size=1024 str_len=32 : 87.75 µs -> 50.64 µs (−42.3%)
  - size=1024 str_len=128 : 313.00 µs -> 187.77 µs (−40.0%)

  Group 5: Empty-from StringViewArray

  - size=1024 str_len=32 : 87.01 µs -> 50.10 µs (−42.4%)
  - size=1024 str_len=128 : 313.99 µs -> 190.17 µs (−39.4%)

## What changes are included in this PR?

* Add `append_byte_map` and `append_with` to both of the bulk-NULL
string builders
* Add unit tests
* Adopt the new APIs in `replace`

## Are these changes tested?

Yes; new tests added.

## Are there any user-facing changes?

No.

Author: neilconway

datafusion/functions/benches/replace.rs

@@ -162,6 +162,36 @@ fn criterion_benchmark(c: &mut Criterion) {
             }
         }
 
+        // Empty-`from` path: insert `to` between every char of the input and at
+        // both ends.
+        if size == 1024 {
+            for &str_len in &[32_usize, 128] {
+                let args = create_args::<i32>(size, str_len, false, 0, 3, 0.0);
+                group.bench_function(
+                    format!("replace_string_empty_from [size={size}, str_len={str_len}]"),
+                    |b| {
+                        b.iter(|| {
+                            let args_cloned = args.clone();
+                            black_box(invoke_replace_with_args(args_cloned, size))
+                        })
+                    },
+                );
+
+                let args = create_args::<i32>(size, str_len, true, 0, 3, 0.0);
+                group.bench_function(
+                    format!(
+                        "replace_string_view_empty_from [size={size}, str_len={str_len}]"
+                    ),
+                    |b| {
+                        b.iter(|| {
+                            let args_cloned = args.clone();
+                            black_box(invoke_replace_with_args(args_cloned, size))
+                        })
+                    },
+                );
+            }
+        }
+
         group.finish();
     }
 }

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_many([
         string_array.nulls(),
         from_array.nulls(),
@@ -184,19 +185,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);
         }
     }
 
@@ -212,7 +209,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_many([
         string_array.nulls(),
         from_array.nulls(),
@@ -232,71 +228,75 @@ 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) {
-    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);
-        }
-        return;
-    }
-
-    // Fast path for replacing a single ASCII character with another single ASCII character.
-    // Extends the buffer's underlying Vec<u8> directly, for performance.
-    if let ([from_byte], [to_byte]) = (from.as_bytes(), to.as_bytes())
+fn apply_replace<B: BulkNullStringArrayBuilder>(
+    builder: &mut B,
+    string: &str,
+    from: &str,
+    to: &str,
+) {
+    // Hot path: single ASCII byte → single ASCII byte. An ASCII byte (< 0x80)
+    // cannot appear inside a multi-byte UTF-8 sequence, so any multi-byte
+    // sequences in `string` pass through unchanged and output stays valid
+    // UTF-8.
+    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.
+        // SAFETY: see the contract above.
         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;
     }
 
+    if from.is_empty() {
+        // 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;
+    }
+
+    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};