diff --git a/arrow-buffer/Cargo.toml b/arrow-buffer/Cargo.toml index 39bdf82120b4..1d71d5c288a3 100644 --- a/arrow-buffer/Cargo.toml +++ b/arrow-buffer/Cargo.toml @@ -68,5 +68,9 @@ harness = false name = "mutable_buffer_extend" harness = false +[[bench]] +name = "null_buffer_builder" +harness = false + [lints] workspace = true diff --git a/arrow-buffer/benches/null_buffer_builder.rs b/arrow-buffer/benches/null_buffer_builder.rs new file mode 100644 index 000000000000..00d1ded9d147 --- /dev/null +++ b/arrow-buffer/benches/null_buffer_builder.rs @@ -0,0 +1,59 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_buffer::NullBufferBuilder; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use std::hint; + +const WORDS: usize = 512; +const LEN: usize = WORDS * 64; + +fn criterion_benchmark(c: &mut Criterion) { + let very_sparse_nulls = (0..WORDS) + .map(|i| if i % 32 == 0 { u64::MAX - 1 } else { u64::MAX }) + .collect(); + let sparse_nulls = (0..WORDS) + .map(|i| if i % 8 == 0 { u64::MAX - 1 } else { u64::MAX }) + .collect(); + let inputs = [ + ("all_valid", vec![u64::MAX; WORDS]), + ("very_sparse_nulls", very_sparse_nulls), + ("sparse_nulls", sparse_nulls), + ("dense_nulls", vec![0x0000_0000_0000_FFFF; WORDS]), + ("alternating", vec![0xAAAA_AAAA_AAAA_AAAA; WORDS]), + ]; + + let mut group = c.benchmark_group("NullBufferBuilder append"); + group.throughput(Throughput::Elements(LEN as u64)); + + for (name, words) in inputs { + group.bench_with_input(BenchmarkId::new("word", name), &words, |b, words| { + b.iter(|| { + let mut builder = NullBufferBuilder::new(LEN); + for &word in words { + builder.append_word(hint::black_box(word), 64); + } + hint::black_box(builder.finish()) + }) + }); + } + + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/arrow-buffer/src/builder/null.rs b/arrow-buffer/src/builder/null.rs index ee58ab591ec2..5d9331146a83 100644 --- a/arrow-buffer/src/builder/null.rs +++ b/arrow-buffer/src/builder/null.rs @@ -149,6 +149,27 @@ impl NullBufferBuilder { } } + /// Appends the low `count` bits from `word` as validity values. + /// + /// `word` is treated as a packed LSB-first bitmap, where set bits indicate + /// non-null values. Only the lowest `count` bits are appended; higher bits + /// are ignored. + /// + /// This does not materialize the null buffer when all appended bits are set. + /// + /// # Panics + /// + /// Panics if `count` is greater than 64. + #[inline] + pub fn append_word(&mut self, word: u64, count: usize) { + assert!(count <= 64, "count must be less than or equal to 64"); + if word.trailing_ones() as usize >= count { + self.append_n_non_nulls(count); + } else { + self.materialize_if_needed().append_word(word, count) + } + } + /// Sets a bit in the builder at `index` /// /// # Panics @@ -339,6 +360,35 @@ mod tests { assert_eq!(&[0b1011_u8], buf.validity()); } + #[test] + fn test_append_word_all_valid() { + let mut builder = NullBufferBuilder::new(0); + builder.append_word(u64::MAX, 64); + builder.append_word(0b0111, 3); + + assert_eq!(builder.len(), 67); + assert_eq!(builder.allocated_size(), 0); + assert!(builder.finish().is_none()); + } + + #[test] + fn test_append_word_mixed() { + let mut builder = NullBufferBuilder::new(0); + builder.append_n_non_nulls(3); + builder.append_word(0b1100_1010, 6); + builder.append_word(0b1111, 4); + + let buffer = builder.finish().unwrap(); + assert_eq!(&[0b0101_0111, 0b0001_1110], buffer.validity()); + } + + #[test] + #[should_panic(expected = "count must be less than or equal to 64")] + fn test_append_word_count_too_large() { + let mut builder = NullBufferBuilder::new(0); + builder.append_word(u64::MAX, 65); + } + #[test] fn test_null_buffer_builder_is_valid() { let mut builder = NullBufferBuilder::new(0);