Skip to content

Commit 0626ca3

Browse files
authored
perf: Optimize Utf8View string concat (#21535)
## Which issue does this PR close? - Closes #21534. ## Rationale for this change Optimize `||` for `Utf8View` values in two ways: 1. We previously checked two `Option`s for every row, to check for NULL values. It is faster to precompute the NULL bitmap and then just check a single bitmap. Annoyingly, `StringViewBuilder` still requires constructing the result NULL bitmap itself incrementally, but fixing that would be a more invasive change. 2. We previously constructed the result value with `write!({l}, {r})`, which goes through the `fmt` machinery and is very slow. It is more efficient to just `write_str` twice. Benchmarks (Arm64): ``` - concat_utf8view/concat/nulls_0: 289.9µs → 140.2µs (-51.6%) - concat_utf8view/concat/nulls_10: 293.5µs → 154.7µs (-47.3%) - concat_utf8view/concat/nulls_50: 197.9µs → 95.0µs (-52.0%) ``` ## What changes are included in this PR? * Add benchmark for string concatenation * Implement optimizations described above ## Are these changes tested? Yes. ## Are there any user-facing changes? No.
1 parent eaf0a41 commit 0626ca3

3 files changed

Lines changed: 115 additions & 13 deletions

File tree

datafusion/physical-expr/Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,5 +87,9 @@ name = "binary_op"
8787
harness = false
8888
name = "simplify"
8989

90+
[[bench]]
91+
harness = false
92+
name = "string_concat"
93+
9094
[package.metadata.cargo-machete]
9195
ignored = ["half"]
Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
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+
use arrow::array::StringViewArray;
19+
use arrow::datatypes::{DataType, Field, Schema};
20+
use arrow::record_batch::RecordBatch;
21+
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
22+
use datafusion_expr::Operator;
23+
use datafusion_physical_expr::PhysicalExpr;
24+
use datafusion_physical_expr::expressions::{BinaryExpr, Column};
25+
use rand::rngs::StdRng;
26+
use rand::{Rng, SeedableRng};
27+
use std::hint::black_box;
28+
use std::sync::Arc;
29+
30+
const NUM_ROWS: usize = 8192;
31+
const SEED: u64 = 42;
32+
33+
fn create_string_view_array(
34+
num_rows: usize,
35+
str_len: usize,
36+
null_density: f64,
37+
seed: u64,
38+
) -> StringViewArray {
39+
let mut rng = StdRng::seed_from_u64(seed);
40+
let values: Vec<Option<String>> = (0..num_rows)
41+
.map(|_| {
42+
if rng.random::<f64>() < null_density {
43+
None
44+
} else {
45+
let s: String = (0..str_len)
46+
.map(|_| rng.random_range(b'a'..=b'z') as char)
47+
.collect();
48+
Some(s)
49+
}
50+
})
51+
.collect();
52+
StringViewArray::from_iter(values)
53+
}
54+
55+
fn bench_concat_utf8view(c: &mut Criterion) {
56+
let mut group = c.benchmark_group("concat_utf8view");
57+
58+
let schema = Arc::new(Schema::new(vec![
59+
Field::new("left", DataType::Utf8View, true),
60+
Field::new("right", DataType::Utf8View, true),
61+
]));
62+
63+
// left || right
64+
let expr = BinaryExpr::new(
65+
Arc::new(Column::new("left", 0)),
66+
Operator::StringConcat,
67+
Arc::new(Column::new("right", 1)),
68+
);
69+
70+
for null_density in [0.0, 0.1, 0.5] {
71+
let left = create_string_view_array(NUM_ROWS, 16, null_density, SEED);
72+
let right = create_string_view_array(NUM_ROWS, 16, null_density, SEED + 1);
73+
74+
let batch =
75+
RecordBatch::try_new(schema.clone(), vec![Arc::new(left), Arc::new(right)])
76+
.unwrap();
77+
78+
let label = format!("nulls_{}", (null_density * 100.0) as u32);
79+
group.bench_with_input(
80+
BenchmarkId::new("concat", &label),
81+
&null_density,
82+
|b, _| {
83+
b.iter(|| {
84+
black_box(expr.evaluate(black_box(&batch)).unwrap());
85+
})
86+
},
87+
);
88+
}
89+
90+
group.finish();
91+
}
92+
93+
criterion_group!(benches, bench_concat_utf8view);
94+
criterion_main!(benches);

datafusion/physical-expr/src/expressions/binary/kernels.rs

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
//! This module contains computation kernels that are specific to
1919
//! datafusion and not (yet) targeted to port upstream to arrow
2020
use arrow::array::*;
21+
use arrow::buffer::NullBuffer;
2122
use arrow::compute::kernels::bitwise::{
2223
bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_shift_left,
2324
bitwise_shift_left_scalar, bitwise_shift_right, bitwise_shift_right_scalar,
@@ -177,24 +178,27 @@ pub fn concat_elements_utf8view(
177178
right.len()
178179
)));
179180
}
180-
let capacity = left.len();
181-
let mut result = StringViewBuilder::with_capacity(capacity);
181+
let mut result = StringViewBuilder::with_capacity(left.len());
182182

183-
// Avoid reallocations by writing to a reused buffer (note we
184-
// could be even more efficient r by creating the view directly
185-
// here and avoid the buffer but that would be more complex)
183+
// Avoid reallocations by writing to a reused buffer (note we could be even
184+
// more efficient by creating the view directly here and avoid the buffer
185+
// but that would be more complex)
186186
let mut buffer = String::new();
187187

188-
for (left, right) in left.iter().zip(right.iter()) {
189-
if let (Some(left), Some(right)) = (left, right) {
190-
use std::fmt::Write;
188+
// Pre-compute combined null bitmap, so the per-row NULL check is more
189+
// efficient
190+
let nulls = NullBuffer::union(left.nulls(), right.nulls());
191+
192+
for i in 0..left.len() {
193+
if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
194+
result.append_null();
195+
} else {
196+
let l = left.value(i);
197+
let r = right.value(i);
191198
buffer.clear();
192-
write!(&mut buffer, "{left}{right}")
193-
.expect("writing into string buffer failed");
199+
buffer.push_str(l);
200+
buffer.push_str(r);
194201
result.try_append_value(&buffer)?;
195-
} else {
196-
// at least one of the values is null, so the output is also null
197-
result.append_null()
198202
}
199203
}
200204
Ok(result.finish())

0 commit comments

Comments
 (0)