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

arrow_buffer/builder/
boolean.rs
