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

arrow_buffer/builder/
boolean.rs
