arrow_buffer/builder/
boolean.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::bit_util::apply_bitwise_binary_op;
6use crate::{BooleanBuffer, Buffer, MutableBuffer, NullBuffer, bit_util};
7use std::ops::Range;
8
9/// Builder for [`BooleanBuffer`]
10///
11/// Builds a packed buffer of bits representing boolean values. Each bit in the
12/// buffer corresponds to a boolean value,
13///
14/// # See Also
15///
16/// * [`NullBufferBuilder`] for building [`BooleanBuffer`]s for representing nulls
17/// * [`BufferBuilder`] for building [`Buffer`]s
18///
19/// # Example
20/// ```
21/// # use arrow_buffer::builder::BooleanBufferBuilder;
22/// let mut builder = BooleanBufferBuilder::new(10);
23/// builder.append(true);
24/// builder.append(false);
25/// builder.append_n(3, true); // append 3 trues
26/// let buffer = builder.build();
27/// assert_eq!(buffer.len(), 5); // 5 bits appended
28/// assert_eq!(buffer.values(), &[0b00011101_u8]); // packed bits
29///```
30///
31/// [`BufferBuilder`]: crate::builder::BufferBuilder
32/// [`NullBufferBuilder`]: crate::builder::NullBufferBuilder
33#[derive(Debug)]
34pub struct BooleanBufferBuilder {
35    buffer: MutableBuffer,
36    len: usize,
37}
38
39impl BooleanBufferBuilder {
40    /// Creates a new `BooleanBufferBuilder` with sufficient space for
41    /// `capacity` bits (not bytes).
42    ///
43    /// The capacity is rounded up to the nearest multiple of 8 for the
44    /// allocation.
45    #[inline]
46    pub fn new(capacity: usize) -> Self {
47        let byte_capacity = bit_util::ceil(capacity, 8);
48        let buffer = MutableBuffer::new(byte_capacity);
49        Self { buffer, len: 0 }
50    }
51
52    /// Creates a new `BooleanBufferBuilder` from [`MutableBuffer`] of `len`
53    ///
54    /// # Panics
55    ///
56    /// Panics if `len > buffer.len() * 8`
57    pub fn new_from_buffer(buffer: MutableBuffer, len: usize) -> Self {
58        assert!(len <= buffer.len() * 8);
59        let mut s = Self {
60            len: buffer.len() * 8,
61            buffer,
62        };
63        s.truncate(len);
64        s
65    }
66
67    /// Returns the length of the buffer
68    #[inline]
69    pub fn len(&self) -> usize {
70        self.len
71    }
72
73    /// Sets a bit in the buffer at `index`
74    ///
75    /// # Panics
76    ///
77    /// Panics if `index / 8 >= self.as_slice().len()`
78    #[inline]
79    pub fn set_bit(&mut self, index: usize, v: bool) {
80        if v {
81            bit_util::set_bit(self.buffer.as_mut(), index);
82        } else {
83            bit_util::unset_bit(self.buffer.as_mut(), index);
84        }
85    }
86
87    /// Gets a bit in the buffer at `index`
88    ///
89    /// # Panics
90    ///
91    /// Panics if `index / 8 >= self.as_slice().len()`
92    #[inline]
93    pub fn get_bit(&self, index: usize) -> bool {
94        bit_util::get_bit(self.buffer.as_slice(), index)
95    }
96
97    /// Returns true if empty
98    #[inline]
99    pub fn is_empty(&self) -> bool {
100        self.len == 0
101    }
102
103    /// Returns the capacity of the buffer, in bits (not bytes)
104    ///
105    /// Note this
106    ///
107    /// # Example
108    /// ```
109    /// # use arrow_buffer::builder::BooleanBufferBuilder;
110    /// // empty requires 0 bytes
111    /// let b = BooleanBufferBuilder::new(0);
112    /// assert_eq!(0, b.capacity());
113    /// // Creating space for 1 bit results in 64 bytes (space for 512 bits)
114    /// // (64 is the minimum allocation size for 64 bit architectures)
115    /// let mut b = BooleanBufferBuilder::new(1);
116    /// assert_eq!(512, b.capacity());
117    /// // 1000 bits requires 128 bytes (space for 1024 bits)
118    /// b.append_n(1000, true);
119    /// assert_eq!(1024, b.capacity());
120    /// ```
121    #[inline]
122    pub fn capacity(&self) -> usize {
123        self.buffer.capacity() * 8
124    }
125
126    /// Advances the buffer by `additional` bits
127    #[inline]
128    pub fn advance(&mut self, additional: usize) {
129        let new_len = self.len + additional;
130        let new_len_bytes = bit_util::ceil(new_len, 8);
131        if new_len_bytes > self.buffer.len() {
132            self.buffer.resize(new_len_bytes, 0);
133        }
134        self.len = new_len;
135    }
136
137    /// Truncates the builder to the given length
138    ///
139    /// If `len` is greater than the buffer's current length, this has no effect
140    #[inline]
141    pub fn truncate(&mut self, len: usize) {
142        if len > self.len {
143            return;
144        }
145
146        let new_len_bytes = bit_util::ceil(len, 8);
147        self.buffer.truncate(new_len_bytes);
148        self.len = len;
149
150        let remainder = self.len % 8;
151        if remainder != 0
152            && let Some(last) = self.buffer.as_mut().last_mut()
153        {
154            let mask = (1_u8 << remainder).wrapping_sub(1);
155            *last &= mask;
156        }
157    }
158
159    /// Reserve space to at least `additional` new bits.
160    /// Capacity will be `>= self.len() + additional`.
161    #[inline]
162    pub fn reserve(&mut self, additional: usize) {
163        let capacity = self.len + additional;
164        if capacity > self.capacity() {
165            // convert differential to bytes
166            let additional = bit_util::ceil(capacity, 8) - self.buffer.len();
167            self.buffer.reserve(additional);
168        }
169    }
170
171    /// Resizes the buffer, either truncating its contents (with no change in capacity), or
172    /// growing it (potentially reallocating it) and writing `false` in the newly available bits.
173    #[inline]
174    pub fn resize(&mut self, len: usize) {
175        match len.checked_sub(self.len) {
176            Some(delta) => self.advance(delta),
177            None => self.truncate(len),
178        }
179    }
180
181    /// Appends a boolean `v` into the buffer
182    #[inline]
183    pub fn append(&mut self, v: bool) {
184        self.advance(1);
185        if v {
186            unsafe { bit_util::set_bit_raw(self.buffer.as_mut_ptr(), self.len - 1) };
187        }
188    }
189
190    /// Appends the low `count` bits from `word` into the buffer.
191    ///
192    /// `word` is treated as a packed LSB-first bitmap. Only the lowest
193    /// `count` bits are appended; higher bits are ignored.
194    ///
195    /// This is significantly faster than calling [`Self::append`] in a
196    /// loop when the caller already has bits packed into a `u64`.
197    #[inline]
198    pub fn append_word(&mut self, word: u64, count: usize) {
199        debug_assert!(count <= 64);
200        let mask = (u64::MAX >> ((64 - count) & 63)) * ((count != 0) as u64);
201        let word = word & mask;
202
203        let new_len = self.len + count;
204        let new_len_bytes = bit_util::ceil(new_len, 8);
205        if new_len_bytes > self.buffer.len() {
206            self.buffer.resize(new_len_bytes, 0);
207        }
208
209        let bit_offset = self.len & 7;
210        let byte_start = self.len / 8;
211        let buf = self.buffer.as_slice_mut();
212
213        // Shift the word to align with the bit offset within the
214        // current byte, then OR it into the buffer. The shift merges
215        // correctly with any existing bits in the partial trailing byte.
216        let shifted = word << bit_offset;
217        let shifted_bytes = shifted.to_le_bytes();
218        let bytes_to_write = bit_util::ceil(count + bit_offset, 8).min(8);
219        for i in 0..bytes_to_write {
220            buf[byte_start + i] |= shifted_bytes[i];
221        }
222        // When bit_offset > 0 the shift can overflow into a 9th byte.
223        if bit_offset > 0 && count + bit_offset > 64 {
224            buf[byte_start + 8] |= (word >> (64 - bit_offset)) as u8;
225        }
226
227        self.len = new_len;
228    }
229
230    /// Appends n `additional` bits of value `v` into the buffer
231    #[inline]
232    pub fn append_n(&mut self, additional: usize, v: bool) {
233        match v {
234            true => {
235                let new_len = self.len + additional;
236                let new_len_bytes = bit_util::ceil(new_len, 8);
237                let cur_remainder = self.len % 8;
238                let new_remainder = new_len % 8;
239
240                if cur_remainder != 0
241                    && let Some(last) = self.buffer.as_slice_mut().last_mut()
242                {
243                    // Pad last byte with 1s
244                    *last |= !((1 << cur_remainder) - 1);
245                }
246                self.buffer.resize(new_len_bytes, 0xFF);
247                if new_remainder != 0
248                    && let Some(last) = self.buffer.as_slice_mut().last_mut()
249                {
250                    // Clear remaining bits
251                    *last &= (1 << new_remainder) - 1;
252                }
253                self.len = new_len;
254            }
255            false => self.advance(additional),
256        }
257    }
258
259    /// Appends a slice of booleans into the buffer
260    #[inline]
261    pub fn append_slice(&mut self, slice: &[bool]) {
262        let additional = slice.len();
263        self.advance(additional);
264
265        let offset = self.len() - additional;
266
267        // Up to 8 values touch at most two bytes, so pack them directly. This
268        // includes 8 values at an unaligned offset, which leave no whole byte
269        // for `pack_slice` to pack.
270        if additional <= 8 {
271            if additional == 0 {
272                return;
273            }
274            // At most 8 bits starting at bit 7 end at bit 14, so they span at most two bytes
275            let bits = (Self::pack_bits(slice) as u16) << (offset % 8);
276            let buffer = self.buffer.as_slice_mut();
277            buffer[offset / 8] |= bits as u8;
278            if let Some(next) = buffer.get_mut(offset / 8 + 1) {
279                *next |= (bits >> 8) as u8;
280            }
281        } else {
282            self.pack_slice(offset, slice)
283        }
284    }
285
286    /// Packs up to 8 bools into a byte, `slice[i]` at bit `i`,
287    /// without a branch per value
288    #[inline]
289    fn pack_bits(slice: &[bool]) -> u8 {
290        slice
291            .iter()
292            .enumerate()
293            .fold(0, |acc, (i, v)| acc | ((*v as u8) << i))
294    }
295
296    /// Writes `slice` into the already allocated bits starting at `offset`
297    ///
298    /// Kept out of line so that the short slice path of [`Self::append_slice`]
299    /// stays small when inlined
300    #[inline(never)]
301    fn pack_slice(&mut self, offset: usize, slice: &[bool]) {
302        let additional = slice.len();
303        let buffer = self.buffer.as_slice_mut();
304
305        // Fill the first byte up to the next byte boundary
306        let head = ((8 - offset % 8) % 8).min(additional);
307        let (head_slice, rest) = slice.split_at(head);
308        buffer[offset / 8] |= Self::pack_bits(head_slice) << (offset % 8);
309
310        // Pack 8 bools into each byte without branching: a `bool` is a byte
311        // holding 0 or 1, and the multiply moves byte `i` to bit `56 + i`
312        let byte_offset = (offset + head) / 8;
313        let (chunks, remainder) = rest.as_chunks::<8>();
314        for (dst, chunk) in buffer[byte_offset..].iter_mut().zip(chunks) {
315            let bytes: [u8; 8] = std::array::from_fn(|i| chunk[i] as u8);
316            let packed = u64::from_le_bytes(bytes).wrapping_mul(0x0102_0408_1020_4080) >> 56;
317            *dst = packed as u8;
318        }
319
320        // The remainder starts on a byte boundary
321        if !remainder.is_empty() {
322            buffer[(self.len - remainder.len()) / 8] |= Self::pack_bits(remainder);
323        }
324    }
325
326    /// Append `range` bits from `to_set`
327    ///
328    /// `to_set` is a slice of bits packed LSB-first into `[u8]`
329    ///
330    /// # Panics
331    ///
332    /// Panics if `to_set` does not contain `ceil(range.end / 8)` bytes
333    pub fn append_packed_range(&mut self, range: Range<usize>, to_set: &[u8]) {
334        let offset_write = self.len;
335        let len = range.end - range.start;
336        // allocate new bits as 0
337        self.advance(len);
338        // copy bits from to_set into self.buffer a word at a time
339        apply_bitwise_binary_op(
340            self.buffer.as_slice_mut(),
341            offset_write,
342            to_set,
343            range.start,
344            len,
345            |_a, b| b, // copy bits from to_set
346        );
347    }
348
349    /// Append [`BooleanBuffer`] to this [`BooleanBufferBuilder`]
350    pub fn append_buffer(&mut self, buffer: &BooleanBuffer) {
351        let range = buffer.offset()..buffer.offset() + buffer.len();
352        self.append_packed_range(range, buffer.values())
353    }
354
355    /// Returns the packed bits
356    pub fn as_slice(&self) -> &[u8] {
357        self.buffer.as_slice()
358    }
359
360    /// Returns the packed bits
361    pub fn as_slice_mut(&mut self) -> &mut [u8] {
362        self.buffer.as_slice_mut()
363    }
364
365    /// Resets this builder and returns a [`BooleanBuffer`].
366    ///
367    /// Use [`Self::build`] when you don't need to reuse this builder.
368    #[inline]
369    pub fn finish(&mut self) -> BooleanBuffer {
370        let buf = std::mem::replace(&mut self.buffer, MutableBuffer::new(0));
371        let len = std::mem::replace(&mut self.len, 0);
372        BooleanBuffer::new(buf.into(), 0, len)
373    }
374
375    /// Builds a [`BooleanBuffer`] without resetting the builder.
376    ///
377    /// This consumes the builder. Use [`Self::finish`] to reuse it.
378    #[inline]
379    pub fn build(self) -> BooleanBuffer {
380        BooleanBuffer::new(self.buffer.into(), 0, self.len)
381    }
382
383    /// Builds the [BooleanBuffer] without resetting the builder.
384    pub fn finish_cloned(&self) -> BooleanBuffer {
385        BooleanBuffer::new(Buffer::from_slice_ref(self.as_slice()), 0, self.len)
386    }
387
388    /// Extends the builder from a trusted length iterator of booleans.
389    /// # Safety
390    /// Callers must ensure that `iter` reports an exact size via `size_hint`.
391    ///
392    #[inline]
393    pub unsafe fn extend_trusted_len<I>(&mut self, iterator: I)
394    where
395        I: Iterator<Item = bool>,
396    {
397        let len = iterator.size_hint().0;
398        unsafe { self.buffer.extend_bool_trusted_len(iterator, self.len) };
399        self.len += len;
400    }
401}
402
403impl From<BooleanBufferBuilder> for Buffer {
404    #[inline]
405    fn from(builder: BooleanBufferBuilder) -> Self {
406        builder.buffer.into()
407    }
408}
409
410impl From<BooleanBufferBuilder> for BooleanBuffer {
411    #[inline]
412    fn from(builder: BooleanBufferBuilder) -> Self {
413        builder.build()
414    }
415}
416
417impl From<BooleanBufferBuilder> for NullBuffer {
418    #[inline]
419    fn from(builder: BooleanBufferBuilder) -> Self {
420        let boolean_buffer = BooleanBuffer::from(builder);
421        NullBuffer::new(boolean_buffer)
422    }
423}
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428
429    #[test]
430    fn test_boolean_buffer_builder_write_bytes() {
431        let mut b = BooleanBufferBuilder::new(4);
432        b.append(false);
433        b.append(true);
434        b.append(false);
435        b.append(true);
436        assert_eq!(4, b.len());
437        assert_eq!(512, b.capacity());
438        let buffer = b.finish();
439        assert_eq!(4, buffer.len());
440
441        // Overallocate capacity
442        let mut b = BooleanBufferBuilder::new(8);
443        b.append_slice(&[false, true, false, true]);
444        assert_eq!(4, b.len());
445        assert_eq!(512, b.capacity());
446        let buffer = b.finish();
447        assert_eq!(4, buffer.len());
448    }
449
450    #[test]
451    fn test_boolean_buffer_builder_unset_first_bit() {
452        let mut buffer = BooleanBufferBuilder::new(4);
453        buffer.append(true);
454        buffer.append(true);
455        buffer.append(false);
456        buffer.append(true);
457        buffer.set_bit(0, false);
458        assert_eq!(buffer.len(), 4);
459        assert_eq!(buffer.finish().values(), &[0b1010_u8]);
460    }
461
462    #[test]
463    fn test_boolean_buffer_builder_unset_last_bit() {
464        let mut buffer = BooleanBufferBuilder::new(4);
465        buffer.append(true);
466        buffer.append(true);
467        buffer.append(false);
468        buffer.append(true);
469        buffer.set_bit(3, false);
470        assert_eq!(buffer.len(), 4);
471        assert_eq!(buffer.finish().values(), &[0b0011_u8]);
472    }
473
474    #[test]
475    fn test_boolean_buffer_builder_unset_an_inner_bit() {
476        let mut buffer = BooleanBufferBuilder::new(5);
477        buffer.append(true);
478        buffer.append(true);
479        buffer.append(false);
480        buffer.append(true);
481        buffer.set_bit(1, false);
482        assert_eq!(buffer.len(), 4);
483        assert_eq!(buffer.finish().values(), &[0b1001_u8]);
484    }
485
486    #[test]
487    fn test_boolean_buffer_builder_unset_several_bits() {
488        let mut buffer = BooleanBufferBuilder::new(5);
489        buffer.append(true);
490        buffer.append(true);
491        buffer.append(true);
492        buffer.append(false);
493        buffer.append(true);
494        buffer.set_bit(1, false);
495        buffer.set_bit(2, false);
496        assert_eq!(buffer.len(), 5);
497        assert_eq!(buffer.finish().values(), &[0b10001_u8]);
498    }
499
500    #[test]
501    fn test_boolean_buffer_builder_unset_several_bits_bigger_than_one_byte() {
502        let mut buffer = BooleanBufferBuilder::new(16);
503        buffer.append_n(10, true);
504        buffer.set_bit(0, false);
505        buffer.set_bit(3, false);
506        buffer.set_bit(9, false);
507        assert_eq!(buffer.len(), 10);
508        assert_eq!(buffer.finish().values(), &[0b11110110_u8, 0b01_u8]);
509    }
510
511    #[test]
512    fn test_boolean_buffer_builder_flip_several_bits_bigger_than_one_byte() {
513        let mut buffer = BooleanBufferBuilder::new(16);
514        buffer.append_n(5, true);
515        buffer.append_n(5, false);
516        buffer.append_n(5, true);
517        buffer.set_bit(0, false);
518        buffer.set_bit(3, false);
519        buffer.set_bit(9, false);
520        buffer.set_bit(6, true);
521        buffer.set_bit(14, true);
522        buffer.set_bit(13, false);
523        assert_eq!(buffer.len(), 15);
524        assert_eq!(buffer.finish().values(), &[0b01010110_u8, 0b1011100_u8]);
525    }
526
527    #[test]
528    fn test_bool_buffer_builder_get_first_bit() {
529        let mut buffer = BooleanBufferBuilder::new(16);
530        buffer.append_n(8, true);
531        buffer.append_n(8, false);
532        assert!(buffer.get_bit(0));
533    }
534
535    #[test]
536    fn test_bool_buffer_builder_get_first_bit_not_requires_mutability() {
537        let buffer = {
538            let mut buffer = BooleanBufferBuilder::new(16);
539            buffer.append_n(8, true);
540            buffer
541        };
542
543        assert!(buffer.get_bit(0));
544    }
545
546    #[test]
547    fn test_bool_buffer_builder_get_last_bit() {
548        let mut buffer = BooleanBufferBuilder::new(16);
549        buffer.append_n(8, true);
550        buffer.append_n(8, false);
551        assert!(!buffer.get_bit(15));
552    }
553
554    #[test]
555    fn test_bool_buffer_builder_get_an_inner_bit() {
556        let mut buffer = BooleanBufferBuilder::new(16);
557        buffer.append_n(4, false);
558        buffer.append_n(8, true);
559        buffer.append_n(4, false);
560        assert!(buffer.get_bit(11));
561    }
562
563    #[test]
564    fn test_bool_buffer_fuzz() {
565        use rand::prelude::*;
566
567        let mut buffer = BooleanBufferBuilder::new(12);
568        let mut all_bools = vec![];
569        let mut rng = rand::rng();
570
571        let src_len = 32;
572        let (src, compacted_src) = {
573            let src: Vec<_> = std::iter::from_fn(|| Some(rng.next_u32() & 1 == 0))
574                .take(src_len)
575                .collect();
576
577            let mut compacted_src = BooleanBufferBuilder::new(src_len);
578            compacted_src.append_slice(&src);
579            (src, compacted_src.finish())
580        };
581
582        for _ in 0..100 {
583            let a = rng.next_u32() as usize % src_len;
584            let b = rng.next_u32() as usize % src_len;
585
586            let start = a.min(b);
587            let end = a.max(b);
588
589            buffer.append_packed_range(start..end, compacted_src.values());
590            all_bools.extend_from_slice(&src[start..end]);
591        }
592
593        let mut compacted = BooleanBufferBuilder::new(all_bools.len());
594        compacted.append_slice(&all_bools);
595
596        assert_eq!(buffer.finish(), compacted.finish())
597    }
598
599    #[test]
600    fn test_boolean_array_builder_resize() {
601        let mut builder = BooleanBufferBuilder::new(20);
602        builder.append_n(4, true);
603        builder.append_n(7, false);
604        builder.append_n(2, true);
605        builder.resize(20);
606
607        assert_eq!(builder.len(), 20);
608        assert_eq!(builder.as_slice(), &[0b00001111, 0b00011000, 0b00000000]);
609
610        builder.resize(5);
611        assert_eq!(builder.len(), 5);
612        assert_eq!(builder.as_slice(), &[0b00001111]);
613
614        builder.append_n(4, true);
615        assert_eq!(builder.len(), 9);
616        assert_eq!(builder.as_slice(), &[0b11101111, 0b00000001]);
617    }
618
619    #[test]
620    fn test_truncate() {
621        let b = MutableBuffer::from_iter([true, true, true, true]);
622        let mut builder = BooleanBufferBuilder::new_from_buffer(b, 2);
623        builder.advance(2);
624        let finished = builder.finish();
625        assert_eq!(finished.values(), &[0b00000011]);
626
627        let mut builder = BooleanBufferBuilder::new(10);
628        builder.append_n(5, true);
629        builder.resize(3);
630        builder.advance(2);
631        let finished = builder.finish();
632        assert_eq!(finished.values(), &[0b00000111]);
633
634        let mut builder = BooleanBufferBuilder::new(10);
635        builder.append_n(16, true);
636        assert_eq!(builder.as_slice(), &[0xFF, 0xFF]);
637        builder.truncate(20);
638        assert_eq!(builder.as_slice(), &[0xFF, 0xFF]);
639        builder.truncate(14);
640        assert_eq!(builder.as_slice(), &[0xFF, 0b00111111]);
641        builder.append(false);
642        builder.append(true);
643        assert_eq!(builder.as_slice(), &[0xFF, 0b10111111]);
644        builder.append_packed_range(0..3, &[0xFF]);
645        assert_eq!(builder.as_slice(), &[0xFF, 0b10111111, 0b00000111]);
646        builder.truncate(17);
647        assert_eq!(builder.as_slice(), &[0xFF, 0b10111111, 0b00000001]);
648        builder.append_packed_range(0..2, &[2]);
649        assert_eq!(builder.as_slice(), &[0xFF, 0b10111111, 0b0000101]);
650        builder.truncate(8);
651        assert_eq!(builder.as_slice(), &[0xFF]);
652        builder.resize(14);
653        assert_eq!(builder.as_slice(), &[0xFF, 0x00]);
654        builder.truncate(0);
655        assert_eq!(builder.as_slice(), &[]);
656    }
657
658    #[test]
659    fn test_boolean_builder_increases_buffer_len() {
660        // 00000010 01001000
661        let buf = Buffer::from([72_u8, 2_u8]);
662        let mut builder = BooleanBufferBuilder::new(8);
663
664        for i in 0..16 {
665            if i == 3 || i == 6 || i == 9 {
666                builder.append(true);
667            } else {
668                builder.append(false);
669            }
670        }
671        let buf2 = builder.finish();
672
673        assert_eq!(buf.len(), buf2.inner().len());
674        assert_eq!(buf.as_slice(), buf2.values());
675    }
676
677    #[test]
678    fn test_extend() {
679        let mut builder = BooleanBufferBuilder::new(0);
680        let bools = vec![true, false, true, true, false, true, true, true, false];
681        unsafe { builder.extend_trusted_len(bools.clone().into_iter()) };
682        assert_eq!(builder.len(), 9);
683        let finished = builder.finish();
684        for (i, v) in bools.into_iter().enumerate() {
685            assert_eq!(finished.value(i), v);
686        }
687
688        // Test > 64 bits
689        let mut builder = BooleanBufferBuilder::new(0);
690        let bools: Vec<_> = (0..100).map(|i| i % 3 == 0 || i % 7 == 0).collect();
691        unsafe { builder.extend_trusted_len(bools.clone().into_iter()) };
692        assert_eq!(builder.len(), 100);
693        let finished = builder.finish();
694        for (i, v) in bools.into_iter().enumerate() {
695            assert_eq!(finished.value(i), v, "at index {i}");
696        }
697    }
698
699    #[test]
700    fn test_extend_misaligned() {
701        // Test misaligned start
702        for offset in 1..65 {
703            let mut builder = BooleanBufferBuilder::new(0);
704            builder.append_n(offset, false);
705
706            let bools: Vec<_> = (0..100).map(|i| i % 3 == 0 || i % 7 == 0).collect();
707            unsafe { builder.extend_trusted_len(bools.clone().into_iter()) };
708            assert_eq!(builder.len(), offset + 100);
709
710            let finished = builder.finish();
711            for i in 0..offset {
712                assert!(!finished.value(i));
713            }
714            for (i, v) in bools.into_iter().enumerate() {
715                assert_eq!(finished.value(offset + i), v, "at index {}", offset + i);
716            }
717        }
718    }
719
720    /// Helper: build via append_word and verify against bit-by-bit append.
721    fn check_append_word(initial_bits: usize, word: u64, count: usize) {
722        let mut got = BooleanBufferBuilder::new(0);
723        let mut expected = BooleanBufferBuilder::new(0);
724        got.append_n(initial_bits, true);
725        expected.append_n(initial_bits, true);
726        got.append_word(word, count);
727        for i in 0..count {
728            expected.append(word & (1 << i) != 0);
729        }
730        assert_eq!(got.len(), expected.len());
731        assert_eq!(got.finish(), expected.finish());
732    }
733
734    #[test]
735    fn test_append_word_zero_count() {
736        check_append_word(0, u64::MAX, 0);
737        check_append_word(3, u64::MAX, 0);
738    }
739
740    #[test]
741    fn test_append_word_aligned() {
742        for count in [1, 5, 8, 17, 64] {
743            check_append_word(0, 0xDEAD_BEEF_CAFE_BABE, count);
744        }
745    }
746
747    #[test]
748    fn test_append_word_unaligned() {
749        for offset in 1..=7 {
750            check_append_word(offset, 0xDEAD_BEEF_CAFE_BABE, 13);
751        }
752    }
753
754    #[test]
755    fn test_append_word_overflow_9th_byte() {
756        check_append_word(3, u64::MAX, 64);
757        check_append_word(7, 0xA5A5_A5A5_A5A5_A5A5, 64);
758    }
759
760    #[test]
761    fn test_append_word_small_counts() {
762        check_append_word(0, 0b1, 1);
763        check_append_word(0, 0b1010101, 7);
764        check_append_word(3, 0b1, 1);
765        check_append_word(3, 0b1010101, 7);
766    }
767
768    #[test]
769    fn test_append_word_ignores_high_bits_before_later_appends() {
770        let mut builder = BooleanBufferBuilder::new(0);
771        builder.append_word(0b10, 1);
772        builder.append(false);
773
774        let finished = builder.finish();
775        assert_eq!(finished.len(), 2);
776        assert!(!finished.value(0));
777        assert!(!finished.value(1));
778    }
779
780    #[test]
781    fn test_append_word_full_word() {
782        check_append_word(0, u64::MAX, 64);
783        check_append_word(0, 0, 64);
784    }
785
786    #[test]
787    fn test_append_word_sequential() {
788        let mut got = BooleanBufferBuilder::new(0);
789        let mut expected = BooleanBufferBuilder::new(0);
790        for (word, count) in [(0b1010u64, 4), (0b111u64, 3), (0u64, 5), (u64::MAX, 64)] {
791            got.append_word(word, count);
792            for i in 0..count {
793                expected.append(word & (1 << i) != 0);
794            }
795        }
796        assert_eq!(got.finish(), expected.finish());
797    }
798
799    #[test]
800    fn test_append_word_mixed_with_append() {
801        let mut got = BooleanBufferBuilder::new(0);
802        let mut expected = BooleanBufferBuilder::new(0);
803        got.append(true);
804        expected.append(true);
805        got.append_word(0b1100, 4);
806        for i in 0..4 {
807            expected.append(0b1100u64 & (1 << i) != 0);
808        }
809        got.append(false);
810        expected.append(false);
811        got.append_word(0xFF, 8);
812        expected.append_n(8, true);
813        assert_eq!(got.finish(), expected.finish());
814    }
815
816    #[test]
817    fn test_extend_misaligned_end() {
818        for len in 1..130 {
819            let mut builder = BooleanBufferBuilder::new(0);
820            let mut bools: Vec<_> = (0..len).map(|i| i % 2 == 0).collect();
821            unsafe { builder.extend_trusted_len(bools.clone().into_iter()) };
822            unsafe { builder.extend_trusted_len(bools.clone().into_iter()) };
823            let copy = bools.clone();
824            bools.extend(copy);
825            assert_eq!(builder.len(), 2 * len);
826
827            let finished = builder.finish();
828            for (i, &v) in bools.iter().enumerate() {
829                assert_eq!(finished.value(i), v, "at index {i} for len {len}");
830            }
831        }
832    }
833
834    #[test]
835    fn test_append_slice_all_byte_patterns() {
836        for byte in 0..=u8::MAX {
837            // 8 values take the short path, 16 values are packed with the multiply
838            let values: Vec<bool> = (0..16).map(|i| byte & (1 << (i % 8)) != 0).collect();
839            let mut b = BooleanBufferBuilder::new(8);
840            b.append_slice(&values[..8]);
841            assert_eq!(b.as_slice(), &[byte]);
842
843            let mut b = BooleanBufferBuilder::new(16);
844            b.append_slice(&values);
845            assert_eq!(b.as_slice(), &[byte, byte]);
846        }
847    }
848
849    #[test]
850    fn test_append_slice_offsets() {
851        // Exercise every combination of starting bit offset and slice length
852        // around the byte boundaries, including after a truncate.
853        let pattern: Vec<bool> = (0..100).map(|i| (i * 7 + i / 3) % 5 < 2).collect();
854        for start in 0..17 {
855            for len in 0..pattern.len() {
856                let mut b = BooleanBufferBuilder::new(0);
857                b.append_n(start + 3, true);
858                b.truncate(start);
859                b.append_slice(&pattern[..len]);
860                b.append(true);
861                let finished = b.finish();
862                assert_eq!(finished.len(), start + len + 1);
863                for i in 0..start {
864                    assert!(finished.value(i), "start {start} len {len} index {i}");
865                }
866                for (i, v) in pattern[..len].iter().enumerate() {
867                    assert_eq!(
868                        finished.value(start + i),
869                        *v,
870                        "start {start} len {len} index {i}"
871                    );
872                }
873                assert!(finished.value(start + len));
874            }
875        }
876    }
877}
arrow_buffer/builder/boolean.rs

arrow_buffer/builder/
boolean.rs
