Skip to content

Commit 4f36125

Browse files
committed
Add ASCII fast path optimization to replace function
Add a fast path for replacing single ASCII characters with another single ASCII character, matching Rust's str::replace() optimization. This enables vectorization and avoids UTF-8 boundary checking overhead. Changes: - Added ASCII character detection in replace_into_string() - When both 'from' and 'to' are single ASCII bytes, use direct byte mapping - Updated benchmark to include single ASCII character replacement tests Optimization: - Fast path operates directly on bytes using simple map operation - Compiler can vectorize the byte-wise replacement - Avoids overhead of match_indices() pattern matching for this common case Benchmark Results (Single ASCII Character Replacement) against previous commit: - size=1024, str_len=32: 29.5 µs → 21.4 µs (27% faster) - size=1024, str_len=128: 73.9 µs → 23.4 µs (68% faster) - size=4096, str_len=32: 121.8 µs → 85.6 µs (30% faster) - size=4096, str_len=128: 316.9 µs → 83.8 µs (74% faster) The optimization shows exceptional 27-74% improvements, with the benefit scaling dramatically with string length. For 128-character strings, we achieve over 3x speedup by enabling vectorization and eliminating pattern matching overhead. This addresses reviewer feedback about capturing Rust's str::replace() optimization tricks for single ASCII character replacements.
1 parent c10d410 commit 4f36125

File tree

2 files changed

+55
-11
lines changed

2 files changed

+55
-11
lines changed

datafusion/functions/benches/replace.rs

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,17 @@ fn create_args<O: OffsetSizeTrait>(
3535
size: usize,
3636
str_len: usize,
3737
force_view_types: bool,
38+
from_len: usize,
39+
to_len: usize,
3840
) -> Vec<ColumnarValue> {
3941
if force_view_types {
4042
let string_array =
4143
Arc::new(create_string_view_array_with_len(size, 0.1, str_len, false));
42-
let from_array = Arc::new(create_string_view_array_with_len(size, 0.1, 3, false));
43-
let to_array = Arc::new(create_string_view_array_with_len(size, 0.1, 5, false));
44+
let from_array = Arc::new(create_string_view_array_with_len(
45+
size, 0.1, from_len, false,
46+
));
47+
let to_array =
48+
Arc::new(create_string_view_array_with_len(size, 0.1, to_len, false));
4449
vec![
4550
ColumnarValue::Array(string_array),
4651
ColumnarValue::Array(from_array),
@@ -49,8 +54,8 @@ fn create_args<O: OffsetSizeTrait>(
4954
} else {
5055
let string_array =
5156
Arc::new(create_string_array_with_len::<O>(size, 0.1, str_len));
52-
let from_array = Arc::new(create_string_array_with_len::<O>(size, 0.1, 3));
53-
let to_array = Arc::new(create_string_array_with_len::<O>(size, 0.1, 5));
57+
let from_array = Arc::new(create_string_array_with_len::<O>(size, 0.1, from_len));
58+
let to_array = Arc::new(create_string_array_with_len::<O>(size, 0.1, to_len));
5459

5560
vec![
5661
ColumnarValue::Array(string_array),
@@ -87,9 +92,21 @@ fn criterion_benchmark(c: &mut Criterion) {
8792
group.sample_size(10);
8893
group.measurement_time(Duration::from_secs(10));
8994

90-
// Small strings
95+
// ASCII single character replacement (fast path)
9196
let str_len = 32;
92-
let args = create_args::<i32>(size, str_len, true);
97+
let args = create_args::<i32>(size, str_len, false, 1, 1);
98+
group.bench_function(
99+
format!("replace_string_ascii_single [size={size}, str_len={str_len}]"),
100+
|b| {
101+
b.iter(|| {
102+
let args_cloned = args.clone();
103+
black_box(invoke_replace_with_args(args_cloned, size))
104+
})
105+
},
106+
);
107+
108+
// Multi-character strings (general path)
109+
let args = create_args::<i32>(size, str_len, true, 3, 5);
93110
group.bench_function(
94111
format!("replace_string_view [size={size}, str_len={str_len}]"),
95112
|b| {
@@ -100,7 +117,7 @@ fn criterion_benchmark(c: &mut Criterion) {
100117
},
101118
);
102119

103-
let args = create_args::<i32>(size, str_len, false);
120+
let args = create_args::<i32>(size, str_len, false, 3, 5);
104121
group.bench_function(
105122
format!("replace_string [size={size}, str_len={str_len}]"),
106123
|b| {
@@ -111,7 +128,7 @@ fn criterion_benchmark(c: &mut Criterion) {
111128
},
112129
);
113130

114-
let args = create_args::<i64>(size, str_len, false);
131+
let args = create_args::<i64>(size, str_len, false, 3, 5);
115132
group.bench_function(
116133
format!("replace_large_string [size={size}, str_len={str_len}]"),
117134
|b| {
@@ -124,7 +141,18 @@ fn criterion_benchmark(c: &mut Criterion) {
124141

125142
// Larger strings
126143
let str_len = 128;
127-
let args = create_args::<i32>(size, str_len, true);
144+
let args = create_args::<i32>(size, str_len, false, 1, 1);
145+
group.bench_function(
146+
format!("replace_string_ascii_single [size={size}, str_len={str_len}]"),
147+
|b| {
148+
b.iter(|| {
149+
let args_cloned = args.clone();
150+
black_box(invoke_replace_with_args(args_cloned, size))
151+
})
152+
},
153+
);
154+
155+
let args = create_args::<i32>(size, str_len, true, 3, 5);
128156
group.bench_function(
129157
format!("replace_string_view [size={size}, str_len={str_len}]"),
130158
|b| {
@@ -135,7 +163,7 @@ fn criterion_benchmark(c: &mut Criterion) {
135163
},
136164
);
137165

138-
let args = create_args::<i32>(size, str_len, false);
166+
let args = create_args::<i32>(size, str_len, false, 3, 5);
139167
group.bench_function(
140168
format!("replace_string [size={size}, str_len={str_len}]"),
141169
|b| {
@@ -146,7 +174,7 @@ fn criterion_benchmark(c: &mut Criterion) {
146174
},
147175
);
148176

149-
let args = create_args::<i64>(size, str_len, false);
177+
let args = create_args::<i64>(size, str_len, false, 3, 5);
150178
group.bench_function(
151179
format!("replace_large_string [size={size}, str_len={str_len}]"),
152180
|b| {

datafusion/functions/src/string/replace.rs

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,22 @@ fn replace_into_string(buffer: &mut String, string: &str, from: &str, to: &str)
228228
return;
229229
}
230230

231+
// Fast path for replacing a single ASCII character with another single ASCII character
232+
// This matches Rust's str::replace() optimization and enables vectorization
233+
if let ([from_byte], [to_byte]) = (from.as_bytes(), to.as_bytes())
234+
&& from_byte.is_ascii()
235+
&& to_byte.is_ascii()
236+
{
237+
// SAFETY: We're replacing ASCII with ASCII, which preserves UTF-8 validity
238+
let replaced: Vec<u8> = string
239+
.as_bytes()
240+
.iter()
241+
.map(|b| if *b == *from_byte { *to_byte } else { *b })
242+
.collect();
243+
buffer.push_str(unsafe { std::str::from_utf8_unchecked(&replaced) });
244+
return;
245+
}
246+
231247
let mut last_end = 0;
232248
for (start, _part) in string.match_indices(from) {
233249
buffer.push_str(&string[last_end..start]);

0 commit comments

Comments
 (0)