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 1 commit into
base: main
Choose a base branch
from
+579 −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
1 commit
Select commit
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
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(
NullBuffer::union(string_array.nulls(), from_array.nulls()).as_ref(),
to_array.nulls(),
Expand All @@ -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);
}
}

Expand All @@ -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(),
Expand All @@ -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);
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

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