Skip to content

Commit d8d521a

Browse files
Omega359alamb
andauthored
Add benchmark for substr_index (#9878)
* Fix to_timestamp benchmark * Remove reference to simd and nightly build as simd is no longer an available feature in DataFusion and building with nightly may not be a good recommendation when getting started. * Fixed missing trim() function. * Add benchmark for substr_index * Add missing required-features * Update datafusion/functions/benches/substr_index.rs Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org> --------- Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent b698e5f commit d8d521a

2 files changed

Lines changed: 108 additions & 0 deletions

File tree

datafusion/functions/Cargo.toml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,3 +107,8 @@ required-features = ["datetime_expressions"]
107107
harness = false
108108
name = "to_char"
109109
required-features = ["datetime_expressions"]
110+
111+
[[bench]]
112+
harness = false
113+
name = "substr_index"
114+
required-features = ["unicode_expressions"]
Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
extern crate criterion;
19+
20+
use std::sync::Arc;
21+
22+
use arrow::array::{ArrayRef, Int64Array, StringArray};
23+
use criterion::{black_box, criterion_group, criterion_main, Criterion};
24+
use rand::distributions::{Alphanumeric, Uniform};
25+
use rand::prelude::Distribution;
26+
use rand::Rng;
27+
28+
use datafusion_expr::ColumnarValue;
29+
use datafusion_functions::unicode::substr_index;
30+
31+
struct Filter<Dist, Test> {
32+
dist: Dist,
33+
test: Test,
34+
}
35+
36+
impl<T, Dist, Test> Distribution<T> for Filter<Dist, Test>
37+
where
38+
Dist: Distribution<T>,
39+
Test: Fn(&T) -> bool,
40+
{
41+
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
42+
loop {
43+
let x = self.dist.sample(rng);
44+
if (self.test)(&x) {
45+
return x;
46+
}
47+
}
48+
}
49+
}
50+
51+
fn data() -> (StringArray, StringArray, Int64Array) {
52+
let dist = Filter {
53+
dist: Uniform::new(-4, 5),
54+
test: |x: &i64| x != &0,
55+
};
56+
let mut rng = rand::thread_rng();
57+
let mut strings: Vec<String> = vec![];
58+
let mut delimiters: Vec<String> = vec![];
59+
let mut counts: Vec<i64> = vec![];
60+
61+
for _ in 0..1000 {
62+
let length = rng.gen_range(20..50);
63+
let text: String = (&mut rng)
64+
.sample_iter(&Alphanumeric)
65+
.take(length)
66+
.map(char::from)
67+
.collect();
68+
let char = rng.gen_range(0..text.len());
69+
let delimiter = &text.chars().nth(char).unwrap();
70+
let count = rng.sample(&dist);
71+
72+
strings.push(text);
73+
delimiters.push(delimiter.to_string());
74+
counts.push(count);
75+
}
76+
77+
(
78+
StringArray::from(strings),
79+
StringArray::from(delimiters),
80+
Int64Array::from(counts),
81+
)
82+
}
83+
84+
fn criterion_benchmark(c: &mut Criterion) {
85+
c.bench_function("substr_index_array_array_1000", |b| {
86+
let (strings, delimiters, counts) = data();
87+
let strings = ColumnarValue::Array(Arc::new(strings) as ArrayRef);
88+
let delimiters = ColumnarValue::Array(Arc::new(delimiters) as ArrayRef);
89+
let counts = ColumnarValue::Array(Arc::new(counts) as ArrayRef);
90+
91+
let args = [strings, delimiters, counts];
92+
b.iter(|| {
93+
black_box(
94+
substr_index()
95+
.invoke(&args)
96+
.expect("substr_index should work on valid values"),
97+
)
98+
})
99+
});
100+
}
101+
102+
criterion_group!(benches, criterion_benchmark);
103+
criterion_main!(benches);

0 commit comments

Comments
 (0)