Skip to content

perf(substr_index): speed up scalar and Utf8View #21754

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

Merged
alamb merged 2 commits into from
Apr 25, 2026
+749 −136
Merged

perf(substr_index): speed up scalar and Utf8View #21754

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d5daf5a
perf(substr_index): speed up scalar and Utf8View
kumarUjjawal Apr 21, 2026
2e49afa
fixed sliced-array buffer over-allocation
kumarUjjawal Apr 22, 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
193 changes: 143 additions & 50 deletions datafusion/functions/benches/substr_index.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,20 @@
use std::hint::black_box;
use std::sync::Arc;

use arrow::array::{ArrayRef, Int64Array, StringArray};
use arrow::array::{ArrayRef, Int64Array, StringArray, StringViewArray};
use arrow::datatypes::{DataType, Field};
use criterion::{Criterion, criterion_group, criterion_main};
use datafusion_common::config::ConfigOptions;
use datafusion_common::{ScalarValue, config::ConfigOptions};
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
use datafusion_functions::unicode::substr_index;
use rand::Rng;
use rand::SeedableRng;
use rand::distr::{Alphanumeric, Uniform};
use rand::prelude::Distribution;
use rand::rngs::StdRng;

const ARRAY_DATA_SEED: u64 = 0x5EED_AAAA;
const SCALAR_DATA_SEED: u64 = 0x5EED_BBBB;

struct Filter<Dist, Test> {
dist: Dist,
Expand All @@ -48,28 +53,72 @@ where
}
}

fn data(
#[derive(Clone, Copy)]
enum StringRep {
Utf8,
Utf8View,
}

impl StringRep {
fn name(self) -> &'static str {
match self {
Self::Utf8 => "utf8",
Self::Utf8View => "utf8view",
}
}

fn data_type(self) -> DataType {
match self {
Self::Utf8 => DataType::Utf8,
Self::Utf8View => DataType::Utf8View,
}
}

fn array(self, values: &[String]) -> ArrayRef {
match self {
Self::Utf8 => Arc::new(StringArray::from(values.to_vec())) as ArrayRef,
Self::Utf8View => Arc::new(
values
.iter()
.map(|value| Some(value.as_str()))
.collect::<StringViewArray>(),
) as ArrayRef,
}
}

fn scalar(self, value: &str) -> ScalarValue {
match self {
Self::Utf8 => ScalarValue::Utf8(Some(value.to_string())),
Self::Utf8View => ScalarValue::Utf8View(Some(value.to_string())),
}
}
}

fn random_token<R: Rng + ?Sized>(rng: &mut R, len: usize) -> String {
rng.sample_iter(&Alphanumeric)
.take(len)
.map(char::from)
.collect()
}

fn array_data(
batch_size: usize,
single_char_delimiter: bool,
) -> (StringArray, StringArray, Int64Array) {
let dist = Filter {
dist: Uniform::new(-4, 5),
) -> (Vec<String>, Vec<String>, Vec<i64>) {
let count_dist = Filter {
dist: Uniform::new(-4, 5).expect("valid count distribution"),
test: |x: &i64| x != &0,
};
let mut rng = rand::rng();
let mut strings: Vec<String> = vec![];
let mut delimiters: Vec<String> = vec![];
let mut counts: Vec<i64> = vec![];
let mut rng = StdRng::seed_from_u64(ARRAY_DATA_SEED);
let mut strings = Vec::with_capacity(batch_size);
let mut delimiters = Vec::with_capacity(batch_size);
let mut counts = Vec::with_capacity(batch_size);

for _ in 0..batch_size {
let length = rng.random_range(20..50);
let base: String = (&mut rng)
.sample_iter(&Alphanumeric)
.take(length)
.map(char::from)
.collect();
let base = random_token(&mut rng, length);

let (string_value, delimiter): (String, String) = if single_char_delimiter {
let (string_value, delimiter) = if single_char_delimiter {
let char_idx = rng.random_range(0..base.chars().count());
let delimiter = base.chars().nth(char_idx).unwrap().to_string();
(base, delimiter)
Expand All @@ -80,7 +129,6 @@ fn data(

let delimiter_count = rng.random_range(1..4);
let mut result = String::new();

for i in 0..delimiter_count {
result.push_str(&base);
if i < delimiter_count - 1 {
Expand All @@ -90,32 +138,37 @@ fn data(
(result, delimiter)
};

let count = rng.sample(dist.dist.unwrap());

strings.push(string_value);
delimiters.push(delimiter);
counts.push(count);
counts.push(count_dist.sample(&mut rng));
}

(strings, delimiters, counts)
}

fn scalar_data(batch_size: usize, delimiter: &str) -> Vec<String> {
let mut rng = StdRng::seed_from_u64(SCALAR_DATA_SEED);
let mut strings = Vec::with_capacity(batch_size);

for _ in 0..batch_size {
let left_len = rng.random_range(12..24);
let middle_len = rng.random_range(12..24);
let right_len = rng.random_range(12..24);
let left = random_token(&mut rng, left_len);
let middle = random_token(&mut rng, middle_len);
let right = random_token(&mut rng, right_len);
strings.push(format!("{left}{delimiter}{middle}{delimiter}{right}"));
}

(
StringArray::from(strings),
StringArray::from(delimiters),
Int64Array::from(counts),
)
strings
}

fn run_benchmark(
b: &mut criterion::Bencher,
strings: StringArray,
delimiters: StringArray,
counts: Int64Array,
batch_size: usize,
args: &[ColumnarValue],
return_type: &DataType,
number_rows: usize,
) {
let strings = ColumnarValue::Array(Arc::new(strings) as ArrayRef);
let delimiters = ColumnarValue::Array(Arc::new(delimiters) as ArrayRef);
let counts = ColumnarValue::Array(Arc::new(counts) as ArrayRef);

let args = vec![strings, delimiters, counts];
let arg_fields = args
.iter()
.enumerate()
Expand All @@ -129,10 +182,10 @@ fn run_benchmark(
black_box(
substr_index()
.invoke_with_args(ScalarFunctionArgs {
args: args.clone(),
args: args.to_vec(),
arg_fields: arg_fields.clone(),
number_rows: batch_size,
return_field: Field::new("f", DataType::Utf8, true).into(),
number_rows,
return_field: Field::new("f", return_type.clone(), true).into(),
config_options: Arc::clone(&config_options),
})
.expect("substr_index should work on valid values"),
Expand All @@ -142,22 +195,62 @@ fn run_benchmark(

fn criterion_benchmark(c: &mut Criterion) {
let mut group = c.benchmark_group("substr_index");

let batch_sizes = [100, 1000, 10_000];

for batch_size in batch_sizes {
group.bench_function(
format!("substr_index_{batch_size}_single_delimiter"),
|b| {
let (strings, delimiters, counts) = data(batch_size, true);
run_benchmark(b, strings, delimiters, counts, batch_size);
},
);

group.bench_function(format!("substr_index_{batch_size}_long_delimiter"), |b| {
let (strings, delimiters, counts) = data(batch_size, false);
run_benchmark(b, strings, delimiters, counts, batch_size);
});
for rep in [StringRep::Utf8, StringRep::Utf8View] {
let rep_name = rep.name();

group.bench_function(
format!("substr_index_{rep_name}_{batch_size}_array_single_delimiter"),
|b| {
let (strings, delimiters, counts) = array_data(batch_size, true);
let args = vec![
ColumnarValue::Array(rep.array(&strings)),
ColumnarValue::Array(rep.array(&delimiters)),
ColumnarValue::Array(
Arc::new(Int64Array::from(counts)) as ArrayRef
),
];
run_benchmark(b, &args, &rep.data_type(), batch_size);
},
);

group.bench_function(
format!("substr_index_{rep_name}_{batch_size}_array_long_delimiter"),
|b| {
let (strings, delimiters, counts) = array_data(batch_size, false);
let args = vec![
ColumnarValue::Array(rep.array(&strings)),
ColumnarValue::Array(rep.array(&delimiters)),
ColumnarValue::Array(
Arc::new(Int64Array::from(counts)) as ArrayRef
),
];
run_benchmark(b, &args, &rep.data_type(), batch_size);
},
);

for (name, delimiter, count) in [
("single_delimiter_pos", ".", 1_i64),
("single_delimiter_neg", ".", -1_i64),
("long_delimiter_pos", "|||", 1_i64),
("long_delimiter_neg", "|||", -1_i64),
] {
group.bench_function(
format!("substr_index_{rep_name}_{batch_size}_scalar_{name}"),
|b| {
let strings = scalar_data(batch_size, delimiter);
let args = vec![
ColumnarValue::Array(rep.array(&strings)),
ColumnarValue::Scalar(rep.scalar(delimiter)),
ColumnarValue::Scalar(ScalarValue::Int64(Some(count))),
];
run_benchmark(b, &args, &rep.data_type(), batch_size);
},
);
}
}
}

group.finish();
Expand Down
Loading
Loading