Skip to content

Commit 55b2f09

Browse files
neilconwayneilconway
authored
perf: Use sub_view in trim, substr, split_part, and substring_index (#26121)
## Which issue does this PR close? - N/A ## Rationale for this change `sub_view` was introduced in #26039; it efficiently builds a view that is a substring of another view. Several string functions had private helpers that did the same thing or used `append_view`. This PR refactors those places to use `sub_view` instead. Both `append_view` and the private helpers were built on Arrow's `make_view`. `make_view` has some shortcomings: * For strings of up to 12 bytes, it does an indirect jump on the string's length. For real-world data with variable length strings, this jump is often poorly predicted. * Perhaps because it has a 13-arm match statement, `make_view` was marked as `inline(never)`, to avoid bloating the code at each call-site In contrast, `sub_view` replaces the indirect jump with shifts and masks, and is small enough to inline. Benchmarks show that adopting `sub_view` improves performance by 5-40%, depending on the workload. Along the way, simplify the `Trimmer` interface because `append_view` is no longer used. Benchmarks: (x86, AMD EPYC Milan): - ltrim/spaces: 52.4µs → 37.5µs, −28% - ltrim/heavy_padding: 285.2µs → 215.6µs, −24% - ltrim/nothing_to_trim: 51.8µs → 29.3µs, −43% - ltrim/char_set: 2.80ms → 2.80ms, 0% - rtrim/spaces: 53.6µs → 37.6µs, −30% - rtrim/heavy_padding: 322.0µs → 253.0µs, −21% - rtrim/nothing_to_trim: 51.4µs → 29.0µs, −44% - rtrim/char_set: 2.90ms → 2.80ms, −3% - btrim/spaces: 87.3µs → 56.6µs, −35% - btrim/heavy_padding: 458.6µs → 398.3µs, −13% - btrim/nothing_to_trim: 58.1µs → 42.0µs, −28% - btrim/char_set: 5.70ms → 5.50ms, −4% - substr/short_result: 114.6µs → 100.5µs, −12% - substr/long_result: 198.4µs → 181.8µs, −8% - substr/short_result_long_input: 93.5µs → 95.5µs, +2% - substr/no_count: 84.0µs → 62.3µs, −26% - substr/per_row_start: 149.0µs → 119.5µs, −20% - split_part/short_fields: 140.9µs → 113.7µs, −19% - split_part/long_fields: 137.6µs → 113.3µs, −18% - split_part/many_fields: 497.9µs → 427.6µs, −14% - split_part/negative_position: 98.6µs → 89.0µs, −10% - split_part/multi_char_delimiter: 163.8µs → 152.6µs, −7% - split_part/per_row_position: 458.5µs → 460.2µs, 0% - substr_index/array_single_delimiter: 295.1µs → 292.7µs, −1% - substr_index/array_long_delimiter: 471.7µs → 466.5µs, −1% - substr_index/scalar_single_delimiter_pos: 75.0µs → 72.0µs, −4% - substr_index/scalar_single_delimiter_neg: 79.3µs → 75.7µs, −5% - substr_index/scalar_long_delimiter_pos: 103.1µs → 104.5µs, +1% - substr_index/scalar_long_delimiter_neg: 219.2µs → 207.3µs, −5% - left/short_result: 37.2µs → 37.7µs, +1% - left/long_result: 38.2µs → 38.2µs, 0% - left/short_result_long_input: 50.4µs → 50.3µs, 0% - left/n_exceeds_len: 36.4µs → 34.6µs, −5% - left/negative_n: 34.7µs → 34.6µs, 0% - left/per_row_n: 33.5µs → 33.8µs, +1% - right/short_result: 48.1µs → 48.3µs, 0% - right/long_result: 52.2µs → 45.0µs, −14% - right/short_result_long_input: 63.1µs → 61.1µs, −3% - right/n_exceeds_len: 46.4µs → 41.9µs, −10% - right/negative_n: 45.7µs → 41.0µs, −10% - right/per_row_n: 45.3µs → 43.7µs, −4% ## What changes are included in this PR? * Refactor `trim`, `substr`, `split_part`, and `substring_index` to use `sub_view` * Remove `append_view` (no longer used) * Simplify `Trimmer` interface (byte offset no longer needed) * Optimize `sub_view` for long result strings ## What is the testing strategy for this PR? Covered by existing tests. ## Are there any user-facing changes? No.
1 parent 5c0e73f commit 55b2f09

9 files changed

Lines changed: 530 additions & 1087 deletions

File tree

‎datafusion/functions/benches/string_expressions/split_part.rs‎

Lines changed: 145 additions & 229 deletions
Original file line numberDiff line numberDiff line change
@@ -15,254 +15,170 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18-
use arrow::array::{ArrayRef, Int64Array, StringArray, StringViewArray};
19-
use arrow::datatypes::{DataType, Field};
20-
use criterion::{BenchmarkId, Criterion, criterion_group};
18+
use std::hint::black_box;
19+
use std::ops::{Range, RangeInclusive};
20+
use std::sync::Arc;
21+
22+
use arrow::array::{StringArray, StringViewArray};
23+
use arrow::datatypes::{Field, Int64Type};
24+
use arrow::util::bench_util::create_primitive_array_range;
25+
use criterion::{Criterion, criterion_group};
2126
use datafusion_common::ScalarValue;
2227
use datafusion_common::config::ConfigOptions;
23-
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF};
28+
use datafusion_expr::{ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs};
2429
use datafusion_functions::string::split_part;
2530
use rand::distr::Alphanumeric;
26-
use rand::prelude::StdRng;
31+
use rand::rngs::StdRng;
2732
use rand::{Rng, SeedableRng};
28-
use std::hint::black_box;
29-
use std::sync::Arc;
3033

31-
const N_ROWS: usize = 8192;
34+
const BATCH_SIZE: usize = 8192;
3235

33-
/// Creates an array of strings with `num_parts` random alphanumeric segments
34-
/// of `part_len` bytes each, joined by `delimiter`.
35-
fn gen_string_array(
36-
n_rows: usize,
37-
num_parts: usize,
38-
part_len: usize,
36+
/// How the position argument is passed.
37+
enum Position {
38+
/// The same position for every row, as in `split_part(s, '.', 2)`.
39+
Scalar(i64),
40+
/// A different position for each row, drawn from the range.
41+
PerRow(Range<i64>),
42+
}
43+
44+
/// Returns strings of alphanumeric fields joined by `delimiter`. Each row's
45+
/// number of fields and each field's length are drawn from the given ranges.
46+
fn create_strings(
47+
field_len: &RangeInclusive<usize>,
48+
num_fields: &RangeInclusive<usize>,
3949
delimiter: &str,
40-
use_string_view: bool,
41-
) -> ColumnarValue {
50+
) -> Vec<Option<String>> {
4251
let mut rng = StdRng::seed_from_u64(42);
43-
44-
let mut strings: Vec<String> = Vec::with_capacity(n_rows);
45-
for _ in 0..n_rows {
46-
let mut parts: Vec<String> = Vec::with_capacity(num_parts);
47-
for _ in 0..num_parts {
48-
let part: String = (&mut rng)
49-
.sample_iter(&Alphanumeric)
50-
.take(part_len)
51-
.map(char::from)
52+
(0..BATCH_SIZE)
53+
.map(|_| {
54+
if rng.random::<f32>() < 0.1 {
55+
return None;
56+
}
57+
let fields: Vec<String> = (0..rng.random_range(num_fields.clone()))
58+
.map(|_| {
59+
let len = rng.random_range(field_len.clone());
60+
(&mut rng)
61+
.sample_iter(&Alphanumeric)
62+
.take(len)
63+
.map(char::from)
64+
.collect()
65+
})
5266
.collect();
53-
parts.push(part);
54-
}
55-
strings.push(parts.join(delimiter));
56-
}
67+
Some(fields.join(delimiter))
68+
})
69+
.collect()
70+
}
5771

58-
if use_string_view {
59-
let string_array: StringViewArray = strings.into_iter().map(Some).collect();
60-
ColumnarValue::Array(Arc::new(string_array) as ArrayRef)
72+
fn create_args(
73+
field_len: &RangeInclusive<usize>,
74+
num_fields: &RangeInclusive<usize>,
75+
delimiter: &str,
76+
position: &Position,
77+
is_string_view: bool,
78+
) -> Vec<ColumnarValue> {
79+
let strings = create_strings(field_len, num_fields, delimiter);
80+
let string_arg = if is_string_view {
81+
ColumnarValue::Array(Arc::new(strings.into_iter().collect::<StringViewArray>()))
6182
} else {
62-
let string_array: StringArray = strings.into_iter().map(Some).collect();
63-
ColumnarValue::Array(Arc::new(string_array) as ArrayRef)
64-
}
65-
}
83+
ColumnarValue::Array(Arc::new(strings.into_iter().collect::<StringArray>()))
84+
};
6685

67-
#[expect(clippy::too_many_arguments)]
68-
fn bench_split_part(
69-
group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>,
70-
func: &ScalarUDF,
71-
config_options: &Arc<ConfigOptions>,
72-
name: &str,
73-
tag: &str,
74-
strings: ColumnarValue,
75-
delimiter: ColumnarValue,
76-
position: ColumnarValue,
77-
) {
78-
let args = vec![strings, delimiter, position];
79-
let arg_fields: Vec<_> = args
80-
.iter()
81-
.enumerate()
82-
.map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into())
83-
.collect();
84-
let return_type = match args[0].data_type() {
85-
DataType::Utf8View => DataType::Utf8View,
86-
_ => DataType::Utf8,
86+
let position_arg = match position {
87+
Position::Scalar(n) => ColumnarValue::Scalar(ScalarValue::Int64(Some(*n))),
88+
Position::PerRow(range) => ColumnarValue::Array(Arc::new(
89+
create_primitive_array_range::<Int64Type>(BATCH_SIZE, 0.0, range.clone()),
90+
)),
8791
};
88-
let return_field = Field::new("f", return_type, true).into();
8992

90-
group.bench_function(BenchmarkId::new(name, tag), |b| {
91-
b.iter(|| {
92-
black_box(
93-
func.invoke_with_args(ScalarFunctionArgs {
94-
args: args.clone(),
95-
arg_fields: arg_fields.clone(),
96-
number_rows: N_ROWS,
97-
return_field: Arc::clone(&return_field),
98-
config_options: Arc::clone(config_options),
99-
})
100-
.expect("split_part should work"),
101-
)
102-
})
103-
});
93+
vec![
94+
string_arg,
95+
ColumnarValue::Scalar(ScalarValue::from(delimiter)),
96+
position_arg,
97+
]
10498
}
10599

106100
fn criterion_benchmark(c: &mut Criterion) {
107-
let split_part_func = split_part();
101+
// Field lengths and counts vary within each case, as in real data: with
102+
// fixed lengths, per-row work that depends on them is unrealistically
103+
// predictable.
104+
let cases = [
105+
// Fields of up to 8 bytes, stored inline in StringView arrays (≤12 bytes).
106+
("short_fields", 1..=8, 3..=6, ".", Position::Scalar(2)),
107+
// Fields of more than 12 bytes, stored out of line.
108+
("long_fields", 16..=48, 3..=6, ".", Position::Scalar(2)),
109+
// A field near the middle of many fields.
110+
("many_fields", 1..=16, 20..=50, ".", Position::Scalar(10)),
111+
// A negative position counts fields from the end.
112+
(
113+
"negative_position",
114+
1..=32,
115+
3..=6,
116+
".",
117+
Position::Scalar(-1),
118+
),
119+
// A delimiter of more than one character.
120+
(
121+
"multi_char_delimiter",
122+
1..=32,
123+
3..=6,
124+
"~@~",
125+
Position::Scalar(2),
126+
),
127+
// The position computed per row.
128+
(
129+
"per_row_position",
130+
1..=32,
131+
3..=6,
132+
".",
133+
Position::PerRow(1..4),
134+
),
135+
];
136+
let function = split_part();
108137
let config_options = Arc::new(ConfigOptions::default());
109-
let mut group = c.benchmark_group("split_part");
110-
111-
// ── Scalar delimiter and position ────────────────
112-
113-
// Utf8, single-char delimiter, scalar args
114-
{
115-
let strings = gen_string_array(N_ROWS, 10, 8, ".", false);
116-
let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some(".".into())));
117-
let position = ColumnarValue::Scalar(ScalarValue::Int64(Some(1)));
118-
bench_split_part(
119-
&mut group,
120-
&split_part_func,
121-
&config_options,
122-
"scalar_utf8_single_char",
123-
"pos_first",
124-
strings,
125-
delimiter,
126-
position,
127-
);
128-
}
129-
130-
{
131-
let strings = gen_string_array(N_ROWS, 10, 8, ".", false);
132-
let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some(".".into())));
133-
let position = ColumnarValue::Scalar(ScalarValue::Int64(Some(5)));
134-
bench_split_part(
135-
&mut group,
136-
&split_part_func,
137-
&config_options,
138-
"scalar_utf8_single_char",
139-
"pos_middle",
140-
strings,
141-
delimiter,
142-
position,
143-
);
144-
}
145-
146-
{
147-
let strings = gen_string_array(N_ROWS, 10, 8, ".", false);
148-
let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some(".".into())));
149-
let position = ColumnarValue::Scalar(ScalarValue::Int64(Some(-1)));
150-
bench_split_part(
151-
&mut group,
152-
&split_part_func,
153-
&config_options,
154-
"scalar_utf8_single_char",
155-
"pos_negative",
156-
strings,
157-
delimiter,
158-
position,
159-
);
160-
}
161-
162-
// Utf8, multi-char delimiter, scalar args
163-
{
164-
let strings = gen_string_array(N_ROWS, 10, 8, "~@~", false);
165-
let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some("~@~".into())));
166-
let position = ColumnarValue::Scalar(ScalarValue::Int64(Some(5)));
167-
bench_split_part(
168-
&mut group,
169-
&split_part_func,
170-
&config_options,
171-
"scalar_utf8_multi_char",
172-
"pos_middle",
173-
strings,
174-
delimiter,
175-
position,
176-
);
177-
}
178-
179-
// Utf8, long strings, scalar args
180-
{
181-
let strings = gen_string_array(N_ROWS, 50, 16, ".", false);
182-
let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8(Some(".".into())));
183-
let position = ColumnarValue::Scalar(ScalarValue::Int64(Some(25)));
184-
bench_split_part(
185-
&mut group,
186-
&split_part_func,
187-
&config_options,
188-
"scalar_utf8_long_strings",
189-
"pos_middle",
190-
strings,
191-
delimiter,
192-
position,
193-
);
194-
}
195-
196-
// Utf8View, long parts, scalar args
197-
{
198-
let strings = gen_string_array(N_ROWS, 10, 32, ".", true);
199-
let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8View(Some(".".into())));
200-
let position = ColumnarValue::Scalar(ScalarValue::Int64(Some(5)));
201-
bench_split_part(
202-
&mut group,
203-
&split_part_func,
204-
&config_options,
205-
"scalar_utf8view_long_parts",
206-
"pos_middle",
207-
strings,
208-
delimiter,
209-
position,
210-
);
211-
}
212-
213-
// Utf8View, very long parts (256 bytes), position 1
214-
{
215-
let strings = gen_string_array(N_ROWS, 5, 256, ".", true);
216-
let delimiter = ColumnarValue::Scalar(ScalarValue::Utf8View(Some(".".into())));
217-
let position = ColumnarValue::Scalar(ScalarValue::Int64(Some(1)));
218-
bench_split_part(
219-
&mut group,
220-
&split_part_func,
221-
&config_options,
222-
"scalar_utf8view_very_long_parts",
223-
"pos_first",
224-
strings,
225-
delimiter,
226-
position,
227-
);
228-
}
229-
230-
// ── Array delimiter and position ─────────────────
231-
232-
// Utf8, single-char delimiter, array args
233-
{
234-
let strings = gen_string_array(N_ROWS, 10, 8, ".", false);
235-
let delimiters: StringArray = vec![Some("."); N_ROWS].into_iter().collect();
236-
let delimiter = ColumnarValue::Array(Arc::new(delimiters) as ArrayRef);
237-
let positions = ColumnarValue::Array(Arc::new(Int64Array::from(vec![5; N_ROWS])));
238-
bench_split_part(
239-
&mut group,
240-
&split_part_func,
241-
&config_options,
242-
"array_utf8_single_char",
243-
"pos_middle",
244-
strings,
245-
delimiter,
246-
positions,
247-
);
248-
}
249-
250-
// Utf8, multi-char delimiter, array args
251-
{
252-
let strings = gen_string_array(N_ROWS, 10, 8, "~@~", false);
253-
let delimiters: StringArray = vec![Some("~@~"); N_ROWS].into_iter().collect();
254-
let delimiter = ColumnarValue::Array(Arc::new(delimiters) as ArrayRef);
255-
let positions = ColumnarValue::Array(Arc::new(Int64Array::from(vec![5; N_ROWS])));
256-
bench_split_part(
257-
&mut group,
258-
&split_part_func,
259-
&config_options,
260-
"array_utf8_multi_char",
261-
"pos_middle",
262-
strings,
263-
delimiter,
264-
positions,
265-
);
138+
let mut group = c.benchmark_group(function.name().to_string());
139+
140+
for is_string_view in [false, true] {
141+
let array_type = if is_string_view {
142+
"string_view"
143+
} else {
144+
"string"
145+
};
146+
147+
for (case_name, field_len, num_fields, delimiter, position) in &cases {
148+
let args =
149+
create_args(field_len, num_fields, delimiter, position, is_string_view);
150+
151+
let arg_fields: Vec<_> = args
152+
.iter()
153+
.enumerate()
154+
.map(|(idx, arg)| {
155+
Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
156+
})
157+
.collect();
158+
let scalar_arguments = vec![None; arg_fields.len()];
159+
let return_field = function
160+
.return_field_from_args(ReturnFieldArgs {
161+
arg_fields: &arg_fields,
162+
scalar_arguments: &scalar_arguments,
163+
})
164+
.expect("should resolve return field");
165+
166+
group.bench_function(format!("{array_type} {case_name}"), |b| {
167+
b.iter(|| {
168+
black_box(
169+
function
170+
.invoke_with_args(ScalarFunctionArgs {
171+
args: args.clone(),
172+
arg_fields: arg_fields.clone(),
173+
number_rows: BATCH_SIZE,
174+
return_field: Arc::clone(&return_field),
175+
config_options: Arc::clone(&config_options),
176+
})
177+
.expect("should work"),
178+
)
179+
})
180+
});
181+
}
266182
}
267183

268184
group.finish();

0 commit comments

Comments
 (0)