This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 6a13206 ARROW-11913: [Rust] Improve performance of StringBuilder by
delaying bitmap creation
6a13206 is described below
commit 6a1320604a2021711e637c2ca586b374a593d66e
Author: Ilya Biryukov <[email protected]>
AuthorDate: Sat Mar 13 05:55:45 2021 -0500
ARROW-11913: [Rust] Improve performance of StringBuilder by delaying bitmap
creation
`PrimitiveBuilder<u8>` was performing a lot of unnecessary bit operations.
One operation for each byte of each string one appends.
Delaying the creation of null bitmap builder for `PrimitiveBuilder`
significantly improves performance of `StringBuilder`. The optimization is
also generally useful as the null bitmap is not stored if resulting
arrays do not contain nulls.
Results from the added benchmark on my machine are:
```
bench_primitive/bench_string
time: [6.8971 ms 6.9123 ms 6.9335 ms]
thrpt: [937.48 MiB/s 940.35 MiB/s 942.43 MiB/s]
change:
time: [-49.585% -49.395% -49.204%] (p = 0.00 <
0.05)
thrpt: [+96.867% +97.610% +98.355%]
Performance has improved.
```
Closes #9662 from ilya-biryukov/optstr
Authored-by: Ilya Biryukov <[email protected]>
Signed-off-by: Andrew Lamb <[email protected]>
---
rust/arrow/benches/builder.rs | 40 +++++++++++++++++++++++++++++-
rust/arrow/src/array/builder.rs | 54 ++++++++++++++++++++++++++++++++---------
2 files changed, 81 insertions(+), 13 deletions(-)
diff --git a/rust/arrow/benches/builder.rs b/rust/arrow/benches/builder.rs
index f1db508..fd9f319 100644
--- a/rust/arrow/benches/builder.rs
+++ b/rust/arrow/benches/builder.rs
@@ -51,6 +51,20 @@ fn bench_primitive(c: &mut Criterion) {
group.finish();
}
+fn bench_primitive_nulls(c: &mut Criterion) {
+ let mut group = c.benchmark_group("bench_primitive_nulls");
+ group.bench_function("bench_primitive_nulls", |b| {
+ b.iter(|| {
+ let mut builder = UInt8Builder::new(64);
+ for _ in 0..NUM_BATCHES * BATCH_SIZE {
+ let _ = black_box(builder.append_null());
+ }
+ black_box(builder.finish());
+ })
+ });
+ group.finish();
+}
+
fn bench_bool(c: &mut Criterion) {
let data: Vec<bool> = seedable_rng()
.sample_iter(&Standard)
@@ -74,5 +88,29 @@ fn bench_bool(c: &mut Criterion) {
group.finish();
}
-criterion_group!(benches, bench_primitive, bench_bool);
+fn bench_string(c: &mut Criterion) {
+ const SAMPLE_STRING: &str = "sample string";
+ let mut group = c.benchmark_group("bench_primitive");
+ group.throughput(Throughput::Bytes(
+ ((BATCH_SIZE * NUM_BATCHES * SAMPLE_STRING.len()) as u32).into(),
+ ));
+ group.bench_function("bench_string", |b| {
+ b.iter(|| {
+ let mut builder = StringBuilder::new(64);
+ for _ in 0..NUM_BATCHES * BATCH_SIZE {
+ let _ = black_box(builder.append_value(SAMPLE_STRING));
+ }
+ black_box(builder.finish());
+ })
+ });
+ group.finish();
+}
+
+criterion_group!(
+ benches,
+ bench_primitive,
+ bench_primitive_nulls,
+ bench_bool,
+ bench_string
+);
criterion_main!(benches);
diff --git a/rust/arrow/src/array/builder.rs b/rust/arrow/src/array/builder.rs
index 32eea9c..1ac679d 100644
--- a/rust/arrow/src/array/builder.rs
+++ b/rust/arrow/src/array/builder.rs
@@ -526,7 +526,9 @@ impl ArrayBuilder for BooleanBuilder {
#[derive(Debug)]
pub struct PrimitiveBuilder<T: ArrowPrimitiveType> {
values_builder: BufferBuilder<T::Native>,
- bitmap_builder: BooleanBufferBuilder,
+ /// We only materialize the builder when we add `false`.
+ /// This optimization is **very** important for performance of
`StringBuilder`.
+ bitmap_builder: Option<BooleanBufferBuilder>,
}
impl<T: ArrowPrimitiveType> ArrayBuilder for PrimitiveBuilder<T> {
@@ -566,7 +568,7 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
pub fn new(capacity: usize) -> Self {
Self {
values_builder: BufferBuilder::<T::Native>::new(capacity),
- bitmap_builder: BooleanBufferBuilder::new(capacity),
+ bitmap_builder: None,
}
}
@@ -577,14 +579,17 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
/// Appends a value of type `T` into the builder
pub fn append_value(&mut self, v: T::Native) -> Result<()> {
- self.bitmap_builder.append(true);
+ if let Some(b) = self.bitmap_builder.as_mut() {
+ b.append(true);
+ }
self.values_builder.append(v);
Ok(())
}
/// Appends a null slot into the builder
pub fn append_null(&mut self) -> Result<()> {
- self.bitmap_builder.append(false);
+ self.materialize_bitmap_builder();
+ self.bitmap_builder.as_mut().unwrap().append(false);
self.values_builder.advance(1);
Ok(())
}
@@ -600,7 +605,9 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
/// Appends a slice of type `T` into the builder
pub fn append_slice(&mut self, v: &[T::Native]) -> Result<()> {
- self.bitmap_builder.append_n(v.len(), true);
+ if let Some(b) = self.bitmap_builder.as_mut() {
+ b.append_n(v.len(), true);
+ }
self.values_builder.append_slice(v);
Ok(())
}
@@ -616,7 +623,12 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
"Value and validity lengths must be equal".to_string(),
));
}
- self.bitmap_builder.append_slice(is_valid);
+ if is_valid.iter().any(|v| !*v) {
+ self.materialize_bitmap_builder();
+ }
+ if let Some(b) = self.bitmap_builder.as_mut() {
+ b.append_slice(is_valid);
+ }
self.values_builder.append_slice(values);
Ok(())
}
@@ -624,13 +636,17 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
/// Builds the `PrimitiveArray` and reset this builder.
pub fn finish(&mut self) -> PrimitiveArray<T> {
let len = self.len();
- let null_bit_buffer = self.bitmap_builder.finish();
- let null_count = len - null_bit_buffer.count_set_bits();
+ let null_bit_buffer = self.bitmap_builder.as_mut().map(|b| b.finish());
+ let null_count = len
+ - null_bit_buffer
+ .as_ref()
+ .map(|b| b.count_set_bits())
+ .unwrap_or(len);
let mut builder = ArrayData::builder(T::DATA_TYPE)
.len(len)
.add_buffer(self.values_builder.finish());
if null_count > 0 {
- builder = builder.null_bit_buffer(null_bit_buffer);
+ builder = builder.null_bit_buffer(null_bit_buffer.unwrap());
}
let data = builder.build();
PrimitiveArray::<T>::from(data)
@@ -639,8 +655,12 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
/// Builds the `DictionaryArray` and reset this builder.
pub fn finish_dict(&mut self, values: ArrayRef) -> DictionaryArray<T> {
let len = self.len();
- let null_bit_buffer = self.bitmap_builder.finish();
- let null_count = len - null_bit_buffer.count_set_bits();
+ let null_bit_buffer = self.bitmap_builder.as_mut().map(|b| b.finish());
+ let null_count = len
+ - null_bit_buffer
+ .as_ref()
+ .map(|b| b.count_set_bits())
+ .unwrap_or(len);
let data_type = DataType::Dictionary(
Box::new(T::DATA_TYPE),
Box::new(values.data_type().clone()),
@@ -649,11 +669,21 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
.len(len)
.add_buffer(self.values_builder.finish());
if null_count > 0 {
- builder = builder.null_bit_buffer(null_bit_buffer);
+ builder = builder.null_bit_buffer(null_bit_buffer.unwrap());
}
builder = builder.add_child_data(values.data());
DictionaryArray::<T>::from(builder.build())
}
+
+ fn materialize_bitmap_builder(&mut self) {
+ if self.bitmap_builder.is_some() {
+ return;
+ }
+ let mut b = BooleanBufferBuilder::new(0);
+ b.reserve(self.values_builder.capacity());
+ b.append_n(self.values_builder.len, true);
+ self.bitmap_builder = Some(b);
+ }
}
/// Array builder for `ListArray`
