Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
neilconway wants to merge 4 commits into
base: main
Choose a base branch
from
+697 −51
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
77525bb
.
neilconway May 4, 2026
7b84ffe
.
neilconway May 7, 2026
389e5c2
Merge remote-tracking branch 'origin/main' into neilc/perf-builder-ap…
neilconway May 8, 2026
4dc4c57
Comment cleanup
neilconway May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions datafusion/functions/benches/replace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Expand Down
76 changes: 39 additions & 37 deletions datafusion/functions/src/string/replace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(),
Expand All @@ -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);
}
}

Expand All @@ -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(),
Expand All @@ -232,71 +228,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);
Copy link
Copy Markdown
Contributor

@lyne7-sc lyne7-sc May 7, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wondering about the from.is_empty() branch — the per-write trade-off looks unusually unfavorable here. Is it worth adding a from_empty benchmark case?

Copy link
Copy Markdown
Contributor

@lyne7-sc lyne7-sc May 7, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI, I ran the benchmarks locally for the empty_from cases and it does look like this branch causes a regression.

group                                                                                   main                                   pr-22029-bench
-----                                                                                   ----                                   --------------
replace size=1024/replace_string_empty_from [size=1024, str_len=128, nulls=0.2]         1.00    280.0±1.26µs        ? ?/sec    1.41    394.3±2.06µs        ? ?/sec
replace size=1024/replace_string_empty_from [size=1024, str_len=128, nulls=0]           1.00    339.0±2.10µs        ? ?/sec    1.48   502.4±11.19µs        ? ?/sec
replace size=1024/replace_string_empty_from [size=1024, str_len=32, nulls=0.2]          1.00     78.9±0.71µs        ? ?/sec    1.27    100.1±0.64µs        ? ?/sec
replace size=1024/replace_string_empty_from [size=1024, str_len=32, nulls=0]            1.00     97.4±0.59µs        ? ?/sec    1.33    129.8±2.01µs        ? ?/sec
replace size=1024/replace_string_view_empty_from [size=1024, str_len=128, nulls=0.2]    1.00    281.4±3.73µs        ? ?/sec    1.41    396.2±6.00µs        ? ?/sec
replace size=1024/replace_string_view_empty_from [size=1024, str_len=128, nulls=0]      1.00    338.2±2.78µs        ? ?/sec    1.47    498.1±4.54µs        ? ?/sec
replace size=1024/replace_string_view_empty_from [size=1024, str_len=32, nulls=0.2]     1.00     78.3±0.18µs        ? ?/sec    1.28    100.5±0.28µs        ? ?/sec
replace size=1024/replace_string_view_empty_from [size=1024, str_len=32, nulls=0]       1.00     97.8±0.70µs        ? ?/sec    1.31    128.0±1.11µs        ? ?/sec
replace size=4096/replace_string_empty_from [size=4096, str_len=128, nulls=0.2]         1.00   1140.6±6.91µs        ? ?/sec    1.43  1625.6±11.65µs        ? ?/sec
replace size=4096/replace_string_empty_from [size=4096, str_len=128, nulls=0]           1.00  1411.9±17.41µs        ? ?/sec    1.50      2.1±0.01ms        ? ?/sec
replace size=4096/replace_string_empty_from [size=4096, str_len=32, nulls=0.2]          1.00    317.6±2.31µs        ? ?/sec    1.28    405.6±1.73µs        ? ?/sec
replace size=4096/replace_string_empty_from [size=4096, str_len=32, nulls=0]            1.00    396.4±3.03µs        ? ?/sec    1.29    511.2±5.86µs        ? ?/sec
replace size=4096/replace_string_view_empty_from [size=4096, str_len=128, nulls=0.2]    1.00   1147.6±9.87µs        ? ?/sec    1.42  1624.8±13.37µs        ? ?/sec
replace size=4096/replace_string_view_empty_from [size=4096, str_len=128, nulls=0]      1.00  1433.2±23.33µs        ? ?/sec    1.46      2.1±0.01ms        ? ?/sec
replace size=4096/replace_string_view_empty_from [size=4096, str_len=32, nulls=0.2]     1.00    318.2±1.08µs        ? ?/sec    1.29    409.3±2.62µs        ? ?/sec
replace size=4096/replace_string_view_empty_from [size=4096, str_len=32, nulls=0]       1.00    397.6±5.38µs        ? ?/sec    1.30    517.0±6.47µs        ? ?/sec

Copy link
Copy Markdown
Contributor Author

@neilconway neilconway May 8, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Really interesting! Thanks for raising this.

I didn't quite follow your comment about the per-write tradeoff: I don't think there's anything fundamental about append_with that should be slower for the many-small-writes case (and we still are able to avoid a final memcpy). However, I can reproduce the significant slowdown you suggested for the "empty from" case. I dug into why, and it seems that repeatedly extending an Arrow MutableBuffer is relatively slow: StringWriter::write_str() does MutableBuffer::extend_from_slice(&[byte]), which results in a libc memcpy for every call, which is obviously slower than per-char String::push(c), as the original code does. I don't think MutableBuffer::extend_from_slice(&[byte]) is inherently slow, but there was enough helper functions / abstractions here that LLVM didn't inline, which lead to the per-call memcpy.

A few options:

  1. Ignore it, because replace with empty-from is a corner-case. But it's unfortunate to leave append_with with a performance footgun like this.
  2. Fix it by reverting to a mut String buffer for the empty-from case. Fixes this specific workload but doesn't fix the underlyling issue in StringWriter.
  3. Optimize StringWriter for small-string writes.

I've implemented (3). A combination of special-casing the string length and marking functions as #[inline(always)] appears to convince LLVM to vectorize this code path, which is a nice win (30-40% faster than main for the empty-from). All the inlining might in theory cause code block for other callers but it doesn't appear to regress any other replace benchmarks, at least.

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};
Expand Down
Loading
Loading