arrow_buffer/builder/
boolean.rs

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

arrow_buffer/builder/
boolean.rs
