perf: improve performance of string replace (17-32% faster)

viirya · viirya · commit c10d4100bdf3 · 2025-12-28T11:32:07.000-08:00
Use a reusable String buffer instead of allocating a new String for each
row. This optimization achieves 17-32% performance improvement across different
string types and sizes by avoiding per-row allocations.

Benchmark results:
| Benchmark     | Array Size | String Length | Baseline (µs) | Optimized (µs) | Improvement |
|---------------|------------|---------------|---------------|----------------|-------------|
| string_view   | 1024       | 32            | 32.53         | 22.32          | 31.4% faster|
| string        | 1024       | 32            | 31.89         | 21.49          | 32.6% faster|
| large_string  | 1024       | 32            | 31.75         | 22.01          | 30.7% faster|
| string_view   | 1024       | 128           | 49.51         | 36.11          | 27.1% faster|
| string        | 1024       | 128           | 48.91         | 34.90          | 28.6% faster|
| large_string  | 1024       | 128           | 49.78         | 35.42          | 28.8% faster|
| string_view   | 4096       | 32            | 133.67        | 95.93          | 28.2% faster|
| string        | 4096       | 32            | 131.48        | 91.73          | 30.2% faster|
| large_string  | 4096       | 32            | 129.61        | 92.82          | 28.4% faster|
| string_view   | 4096       | 128           | 191.50        | 153.74         | 19.7% faster|
| string        | 4096       | 128           | 185.27        | 149.37         | 19.4% faster|
| large_string  | 4096       | 128           | 187.82        | 154.32         | 17.8% faster|
diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml
@@ -210,6 +210,11 @@ harness = false
 name = "repeat"
 required-features = ["string_expressions"]
 
+[[bench]]
+harness = false
+name = "replace"
+required-features = ["string_expressions"]
+
 [[bench]]
 harness = false
 name = "random"
diff --git a/datafusion/functions/benches/replace.rs b/datafusion/functions/benches/replace.rs
@@ -0,0 +1,165 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+extern crate criterion;
+
+use arrow::array::OffsetSizeTrait;
+use arrow::datatypes::{DataType, Field};
+use arrow::util::bench_util::{
+    create_string_array_with_len, create_string_view_array_with_len,
+};
+use criterion::{Criterion, SamplingMode, criterion_group, criterion_main};
+use datafusion_common::DataFusionError;
+use datafusion_common::config::ConfigOptions;
+use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
+use datafusion_functions::string;
+use std::hint::black_box;
+use std::sync::Arc;
+use std::time::Duration;
+
+fn create_args<O: OffsetSizeTrait>(
+    size: usize,
+    str_len: usize,
+    force_view_types: bool,
+) -> Vec<ColumnarValue> {
+    if force_view_types {
+        let string_array =
+            Arc::new(create_string_view_array_with_len(size, 0.1, str_len, false));
+        let from_array = Arc::new(create_string_view_array_with_len(size, 0.1, 3, false));
+        let to_array = Arc::new(create_string_view_array_with_len(size, 0.1, 5, false));
+        vec![
+            ColumnarValue::Array(string_array),
+            ColumnarValue::Array(from_array),
+            ColumnarValue::Array(to_array),
+        ]
+    } else {
+        let string_array =
+            Arc::new(create_string_array_with_len::<O>(size, 0.1, str_len));
+        let from_array = Arc::new(create_string_array_with_len::<O>(size, 0.1, 3));
+        let to_array = Arc::new(create_string_array_with_len::<O>(size, 0.1, 5));
+
+        vec![
+            ColumnarValue::Array(string_array),
+            ColumnarValue::Array(from_array),
+            ColumnarValue::Array(to_array),
+        ]
+    }
+}
+
+fn invoke_replace_with_args(
+    args: Vec<ColumnarValue>,
+    number_rows: usize,
+) -> Result<ColumnarValue, DataFusionError> {
+    let arg_fields = args
+        .iter()
+        .enumerate()
+        .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into())
+        .collect::<Vec<_>>();
+    let config_options = Arc::new(ConfigOptions::default());
+
+    string::replace().invoke_with_args(ScalarFunctionArgs {
+        args,
+        arg_fields,
+        number_rows,
+        return_field: Field::new("f", DataType::Utf8, true).into(),
+        config_options: Arc::clone(&config_options),
+    })
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    for size in [1024, 4096] {
+        let mut group = c.benchmark_group(format!("replace size={size}"));
+        group.sampling_mode(SamplingMode::Flat);
+        group.sample_size(10);
+        group.measurement_time(Duration::from_secs(10));
+
+        // Small strings
+        let str_len = 32;
+        let args = create_args::<i32>(size, str_len, true);
+        group.bench_function(
+            format!("replace_string_view [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, false);
+        group.bench_function(
+            format!("replace_string [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::<i64>(size, str_len, false);
+        group.bench_function(
+            format!("replace_large_string [size={size}, str_len={str_len}]"),
+            |b| {
+                b.iter(|| {
+                    let args_cloned = args.clone();
+                    black_box(invoke_replace_with_args(args_cloned, size))
+                })
+            },
+        );
+
+        // Larger strings
+        let str_len = 128;
+        let args = create_args::<i32>(size, str_len, true);
+        group.bench_function(
+            format!("replace_string_view [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, false);
+        group.bench_function(
+            format!("replace_string [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::<i64>(size, str_len, false);
+        group.bench_function(
+            format!("replace_large_string [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();
+    }
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);
diff --git a/datafusion/functions/src/string/replace.rs b/datafusion/functions/src/string/replace.rs
@@ -18,7 +18,7 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait, StringArray};
+use arrow::array::{ArrayRef, GenericStringBuilder, OffsetSizeTrait};
 use arrow::datatypes::DataType;
 
 use crate::utils::{make_scalar_function, utf8_to_str_type};
@@ -165,17 +165,25 @@ fn replace_view(args: &[ArrayRef]) -> Result<ArrayRef> {
     let from_array = as_string_view_array(&args[1])?;
     let to_array = as_string_view_array(&args[2])?;
 
-    let result = string_array
+    let mut builder = GenericStringBuilder::<i32>::new();
+    let mut buffer = String::new();
+
+    for ((string, from), to) in string_array
         .iter()
         .zip(from_array.iter())
         .zip(to_array.iter())
-        .map(|((string, from), to)| match (string, from, to) {
-            (Some(string), Some(from), Some(to)) => Some(string.replace(from, to)),
-            _ => None,
-        })
-        .collect::<StringArray>();
+    {
+        match (string, from, to) {
+            (Some(string), Some(from), Some(to)) => {
+                buffer.clear();
+                replace_into_string(&mut buffer, string, from, to);
+                builder.append_value(&buffer);
+            }
+            _ => builder.append_null(),
+        }
+    }
 
-    Ok(Arc::new(result) as ArrayRef)
+    Ok(Arc::new(builder.finish()) as ArrayRef)
 }
 
 /// Replaces all occurrences in string of substring from with substring to.
@@ -185,17 +193,48 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
     let from_array = as_generic_string_array::<T>(&args[1])?;
     let to_array = as_generic_string_array::<T>(&args[2])?;
 
-    let result = string_array
+    let mut builder = GenericStringBuilder::<T>::new();
+    let mut buffer = String::new();
+
+    for ((string, from), to) in string_array
         .iter()
         .zip(from_array.iter())
         .zip(to_array.iter())
-        .map(|((string, from), to)| match (string, from, to) {
-            (Some(string), Some(from), Some(to)) => Some(string.replace(from, to)),
-            _ => None,
-        })
-        .collect::<GenericStringArray<T>>();
+    {
+        match (string, from, to) {
+            (Some(string), Some(from), Some(to)) => {
+                buffer.clear();
+                replace_into_string(&mut buffer, string, from, to);
+                builder.append_value(&buffer);
+            }
+            _ => builder.append_null(),
+        }
+    }
+
+    Ok(Arc::new(builder.finish()) as ArrayRef)
+}
 
-    Ok(Arc::new(result) 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;
+    }
+
+    let mut last_end = 0;
+    for (start, _part) in string.match_indices(from) {
+        buffer.push_str(&string[last_end..start]);
+        buffer.push_str(to);
+        last_end = start + from.len();
+    }
+    buffer.push_str(&string[last_end..]);
 }
 
 #[cfg(test)]
