[arrow-rs] branch master updated: Faster BooleanBufferBuilder::append_n for true values (#3596)

[tustvold]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 025ffd0b2 Faster BooleanBufferBuilder::append_n for true values (#3596)
025ffd0b2 is described below

commit 025ffd0b226fcaca24ff1e9c64c3189c0926f587
Author: Raphael Taylor-Davies <1781103+tustv...@users.noreply.github.com>
AuthorDate: Tue Jan 24 15:33:18 2023 +0000

    Faster BooleanBufferBuilder::append_n for true values (#3596)
---
 arrow-array/src/builder/boolean_buffer_builder.rs | 27 ++++++++++++++++++-----
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/arrow-array/src/builder/boolean_buffer_builder.rs 
b/arrow-array/src/builder/boolean_buffer_builder.rs
index a0fdea948..ac2a96fea 100644
--- a/arrow-array/src/builder/boolean_buffer_builder.rs
+++ b/arrow-array/src/builder/boolean_buffer_builder.rs
@@ -146,12 +146,27 @@ impl BooleanBufferBuilder {
     /// Appends n `additional` bits of value `v` into the buffer
     #[inline]
     pub fn append_n(&mut self, additional: usize, v: bool) {
-        self.advance(additional);
-        if additional > 0 && v {
-            let offset = self.len() - additional;
-            (0..additional).for_each(|i| unsafe {
-                bit_util::set_bit_raw(self.buffer.as_mut_ptr(), offset + i)
-            })
+        match v {
+            true => {
+                let new_len = self.len + additional;
+                let new_len_bytes = bit_util::ceil(new_len, 8);
+                let cur_remainder = self.len % 8;
+                let new_remainder = new_len % 8;
+
+                if cur_remainder != 0 {
+                    // Pad last byte with 1s
+                    *self.buffer.as_slice_mut().last_mut().unwrap() |=
+                        !((1 << cur_remainder) - 1)
+                }
+                self.buffer.resize(new_len_bytes, 0xFF);
+                if new_remainder != 0 {
+                    // Clear remaining bits
+                    *self.buffer.as_slice_mut().last_mut().unwrap() &=
+                        (1 << new_remainder) - 1
+                }
+                self.len = new_len;
+            }
+            false => self.advance(additional),
         }
     }