arrow_buffer/buffer/
boolean.rs

Help
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_chunk_iterator::BitChunks;
19use crate::bit_iterator::{BitIndexIterator, BitIterator, BitSliceIterator};
20use crate::{
21    bit_util, buffer_bin_and, buffer_bin_or, buffer_bin_xor, buffer_unary_not,
22    BooleanBufferBuilder, Buffer, MutableBuffer,
23};
24
25use std::ops::{BitAnd, BitOr, BitXor, Not};
26
27/// A slice-able [`Buffer`] containing bit-packed booleans
28#[derive(Debug, Clone, Eq)]
29pub struct BooleanBuffer {
30    buffer: Buffer,
31    offset: usize,
32    len: usize,
33}
34
35impl PartialEq for BooleanBuffer {
36    fn eq(&self, other: &Self) -> bool {
37        if self.len != other.len {
38            return false;
39        }
40
41        let lhs = self.bit_chunks().iter_padded();
42        let rhs = other.bit_chunks().iter_padded();
43        lhs.zip(rhs).all(|(a, b)| a == b)
44    }
45}
46
47impl BooleanBuffer {
48    /// Create a new [`BooleanBuffer`] from a [`Buffer`], an `offset` and `length` in bits
49    ///
50    /// # Panics
51    ///
52    /// This method will panic if `buffer` is not large enough
53    pub fn new(buffer: Buffer, offset: usize, len: usize) -> Self {
54        let total_len = offset.saturating_add(len);
55        let bit_len = buffer.len().saturating_mul(8);
56        assert!(total_len <= bit_len);
57        Self {
58            buffer,
59            offset,
60            len,
61        }
62    }
63
64    /// Create a new [`BooleanBuffer`] of `length` where all values are `true`
65    pub fn new_set(length: usize) -> Self {
66        let mut builder = BooleanBufferBuilder::new(length);
67        builder.append_n(length, true);
68        builder.finish()
69    }
70
71    /// Create a new [`BooleanBuffer`] of `length` where all values are `false`
72    pub fn new_unset(length: usize) -> Self {
73        let buffer = MutableBuffer::new_null(length).into_buffer();
74        Self {
75            buffer,
76            offset: 0,
77            len: length,
78        }
79    }
80
81    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BooleanBuffer`
82    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
83        let buffer = MutableBuffer::collect_bool(len, f);
84        Self::new(buffer.into(), 0, len)
85    }
86
87    /// Returns the number of set bits in this buffer
88    pub fn count_set_bits(&self) -> usize {
89        self.buffer.count_set_bits_offset(self.offset, self.len)
90    }
91
92    /// Returns a `BitChunks` instance which can be used to iterate over
93    /// this buffer's bits in `u64` chunks
94    #[inline]
95    pub fn bit_chunks(&self) -> BitChunks {
96        BitChunks::new(self.values(), self.offset, self.len)
97    }
98
99    /// Returns `true` if the bit at index `i` is set
100    ///
101    /// # Panics
102    ///
103    /// Panics if `i >= self.len()`
104    #[inline]
105    #[deprecated(note = "use BooleanBuffer::value")]
106    pub fn is_set(&self, i: usize) -> bool {
107        self.value(i)
108    }
109
110    /// Returns the offset of this [`BooleanBuffer`] in bits
111    #[inline]
112    pub fn offset(&self) -> usize {
113        self.offset
114    }
115
116    /// Returns the length of this [`BooleanBuffer`] in bits
117    #[inline]
118    pub fn len(&self) -> usize {
119        self.len
120    }
121
122    /// Returns true if this [`BooleanBuffer`] is empty
123    #[inline]
124    pub fn is_empty(&self) -> bool {
125        self.len == 0
126    }
127
128    /// Returns the boolean value at index `i`.
129    ///
130    /// # Panics
131    ///
132    /// Panics if `i >= self.len()`
133    #[inline]
134    pub fn value(&self, idx: usize) -> bool {
135        assert!(idx < self.len);
136        unsafe { self.value_unchecked(idx) }
137    }
138
139    /// Returns the boolean value at index `i`.
140    ///
141    /// # Safety
142    /// This doesn't check bounds, the caller must ensure that index < self.len()
143    #[inline]
144    pub unsafe fn value_unchecked(&self, i: usize) -> bool {
145        unsafe { bit_util::get_bit_raw(self.buffer.as_ptr(), i + self.offset) }
146    }
147
148    /// Returns the packed values of this [`BooleanBuffer`] not including any offset
149    #[inline]
150    pub fn values(&self) -> &[u8] {
151        &self.buffer
152    }
153
154    /// Slices this [`BooleanBuffer`] by the provided `offset` and `length`
155    pub fn slice(&self, offset: usize, len: usize) -> Self {
156        assert!(
157            offset.saturating_add(len) <= self.len,
158            "the length + offset of the sliced BooleanBuffer cannot exceed the existing length"
159        );
160        Self {
161            buffer: self.buffer.clone(),
162            offset: self.offset + offset,
163            len,
164        }
165    }
166
167    /// Returns a [`Buffer`] containing the sliced contents of this [`BooleanBuffer`]
168    ///
169    /// Equivalent to `self.buffer.bit_slice(self.offset, self.len)`
170    pub fn sliced(&self) -> Buffer {
171        self.buffer.bit_slice(self.offset, self.len)
172    }
173
174    /// Returns true if this [`BooleanBuffer`] is equal to `other`, using pointer comparisons
175    /// to determine buffer equality. This is cheaper than `PartialEq::eq` but may
176    /// return false when the arrays are logically equal
177    pub fn ptr_eq(&self, other: &Self) -> bool {
178        self.buffer.as_ptr() == other.buffer.as_ptr()
179            && self.offset == other.offset
180            && self.len == other.len
181    }
182
183    /// Returns the inner [`Buffer`]
184    #[inline]
185    pub fn inner(&self) -> &Buffer {
186        &self.buffer
187    }
188
189    /// Returns the inner [`Buffer`], consuming self
190    pub fn into_inner(self) -> Buffer {
191        self.buffer
192    }
193
194    /// Returns an iterator over the bits in this [`BooleanBuffer`]
195    pub fn iter(&self) -> BitIterator<'_> {
196        self.into_iter()
197    }
198
199    /// Returns an iterator over the set bit positions in this [`BooleanBuffer`]
200    pub fn set_indices(&self) -> BitIndexIterator<'_> {
201        BitIndexIterator::new(self.values(), self.offset, self.len)
202    }
203
204    /// Returns a [`BitSliceIterator`] yielding contiguous ranges of set bits
205    pub fn set_slices(&self) -> BitSliceIterator<'_> {
206        BitSliceIterator::new(self.values(), self.offset, self.len)
207    }
208}
209
210impl Not for &BooleanBuffer {
211    type Output = BooleanBuffer;
212
213    fn not(self) -> Self::Output {
214        BooleanBuffer {
215            buffer: buffer_unary_not(&self.buffer, self.offset, self.len),
216            offset: 0,
217            len: self.len,
218        }
219    }
220}
221
222impl BitAnd<&BooleanBuffer> for &BooleanBuffer {
223    type Output = BooleanBuffer;
224
225    fn bitand(self, rhs: &BooleanBuffer) -> Self::Output {
226        assert_eq!(self.len, rhs.len);
227        BooleanBuffer {
228            buffer: buffer_bin_and(&self.buffer, self.offset, &rhs.buffer, rhs.offset, self.len),
229            offset: 0,
230            len: self.len,
231        }
232    }
233}
234
235impl BitOr<&BooleanBuffer> for &BooleanBuffer {
236    type Output = BooleanBuffer;
237
238    fn bitor(self, rhs: &BooleanBuffer) -> Self::Output {
239        assert_eq!(self.len, rhs.len);
240        BooleanBuffer {
241            buffer: buffer_bin_or(&self.buffer, self.offset, &rhs.buffer, rhs.offset, self.len),
242            offset: 0,
243            len: self.len,
244        }
245    }
246}
247
248impl BitXor<&BooleanBuffer> for &BooleanBuffer {
249    type Output = BooleanBuffer;
250
251    fn bitxor(self, rhs: &BooleanBuffer) -> Self::Output {
252        assert_eq!(self.len, rhs.len);
253        BooleanBuffer {
254            buffer: buffer_bin_xor(&self.buffer, self.offset, &rhs.buffer, rhs.offset, self.len),
255            offset: 0,
256            len: self.len,
257        }
258    }
259}
260
261impl<'a> IntoIterator for &'a BooleanBuffer {
262    type Item = bool;
263    type IntoIter = BitIterator<'a>;
264
265    fn into_iter(self) -> Self::IntoIter {
266        BitIterator::new(self.values(), self.offset, self.len)
267    }
268}
269
270impl From<&[bool]> for BooleanBuffer {
271    fn from(value: &[bool]) -> Self {
272        let mut builder = BooleanBufferBuilder::new(value.len());
273        builder.append_slice(value);
274        builder.finish()
275    }
276}
277
278impl From<Vec<bool>> for BooleanBuffer {
279    fn from(value: Vec<bool>) -> Self {
280        value.as_slice().into()
281    }
282}
283
284impl FromIterator<bool> for BooleanBuffer {
285    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> Self {
286        let iter = iter.into_iter();
287        let (hint, _) = iter.size_hint();
288        let mut builder = BooleanBufferBuilder::new(hint);
289        iter.for_each(|b| builder.append(b));
290        builder.finish()
291    }
292}
293
294#[cfg(test)]
295mod tests {
296    use super::*;
297
298    #[test]
299    fn test_boolean_new() {
300        let bytes = &[0, 1, 2, 3, 4];
301        let buf = Buffer::from(bytes);
302        let offset = 0;
303        let len = 24;
304
305        let boolean_buf = BooleanBuffer::new(buf.clone(), offset, len);
306        assert_eq!(bytes, boolean_buf.values());
307        assert_eq!(offset, boolean_buf.offset());
308        assert_eq!(len, boolean_buf.len());
309
310        assert_eq!(2, boolean_buf.count_set_bits());
311        assert_eq!(&buf, boolean_buf.inner());
312        assert_eq!(buf, boolean_buf.clone().into_inner());
313
314        assert!(!boolean_buf.is_empty())
315    }
316
317    #[test]
318    fn test_boolean_data_equality() {
319        let boolean_buf1 = BooleanBuffer::new(Buffer::from(&[0, 1, 4, 3, 5]), 0, 32);
320        let boolean_buf2 = BooleanBuffer::new(Buffer::from(&[0, 1, 4, 3, 5]), 0, 32);
321        assert_eq!(boolean_buf1, boolean_buf2);
322
323        // slice with same offset and same length should still preserve equality
324        let boolean_buf3 = boolean_buf1.slice(8, 16);
325        assert_ne!(boolean_buf1, boolean_buf3);
326        let boolean_buf4 = boolean_buf1.slice(0, 32);
327        assert_eq!(boolean_buf1, boolean_buf4);
328
329        // unequal because of different elements
330        let boolean_buf2 = BooleanBuffer::new(Buffer::from(&[0, 0, 2, 3, 4]), 0, 32);
331        assert_ne!(boolean_buf1, boolean_buf2);
332
333        // unequal because of different length
334        let boolean_buf2 = BooleanBuffer::new(Buffer::from(&[0, 1, 4, 3, 5]), 0, 24);
335        assert_ne!(boolean_buf1, boolean_buf2);
336
337        // ptr_eq
338        assert!(boolean_buf1.ptr_eq(&boolean_buf1));
339        assert!(boolean_buf2.ptr_eq(&boolean_buf2));
340        assert!(!boolean_buf1.ptr_eq(&boolean_buf2));
341    }
342
343    #[test]
344    fn test_boolean_slice() {
345        let bytes = &[0, 3, 2, 6, 2];
346        let boolean_buf1 = BooleanBuffer::new(Buffer::from(bytes), 0, 32);
347        let boolean_buf2 = BooleanBuffer::new(Buffer::from(bytes), 0, 32);
348
349        let boolean_slice1 = boolean_buf1.slice(16, 16);
350        let boolean_slice2 = boolean_buf2.slice(0, 16);
351        assert_eq!(boolean_slice1.values(), boolean_slice2.values());
352
353        assert_eq!(bytes, boolean_slice1.values());
354        assert_eq!(16, boolean_slice1.offset);
355        assert_eq!(16, boolean_slice1.len);
356
357        assert_eq!(bytes, boolean_slice2.values());
358        assert_eq!(0, boolean_slice2.offset);
359        assert_eq!(16, boolean_slice2.len);
360    }
361
362    #[test]
363    fn test_boolean_bitand() {
364        let offset = 0;
365        let len = 40;
366
367        let buf1 = Buffer::from(&[0, 1, 1, 0, 0]);
368        let boolean_buf1 = &BooleanBuffer::new(buf1, offset, len);
369
370        let buf2 = Buffer::from(&[0, 1, 1, 1, 0]);
371        let boolean_buf2 = &BooleanBuffer::new(buf2, offset, len);
372
373        let expected = BooleanBuffer::new(Buffer::from(&[0, 1, 1, 0, 0]), offset, len);
374        assert_eq!(boolean_buf1 & boolean_buf2, expected);
375    }
376
377    #[test]
378    fn test_boolean_bitor() {
379        let offset = 0;
380        let len = 40;
381
382        let buf1 = Buffer::from(&[0, 1, 1, 0, 0]);
383        let boolean_buf1 = &BooleanBuffer::new(buf1, offset, len);
384
385        let buf2 = Buffer::from(&[0, 1, 1, 1, 0]);
386        let boolean_buf2 = &BooleanBuffer::new(buf2, offset, len);
387
388        let expected = BooleanBuffer::new(Buffer::from(&[0, 1, 1, 1, 0]), offset, len);
389        assert_eq!(boolean_buf1 | boolean_buf2, expected);
390    }
391
392    #[test]
393    fn test_boolean_bitxor() {
394        let offset = 0;
395        let len = 40;
396
397        let buf1 = Buffer::from(&[0, 1, 1, 0, 0]);
398        let boolean_buf1 = &BooleanBuffer::new(buf1, offset, len);
399
400        let buf2 = Buffer::from(&[0, 1, 1, 1, 0]);
401        let boolean_buf2 = &BooleanBuffer::new(buf2, offset, len);
402
403        let expected = BooleanBuffer::new(Buffer::from(&[0, 0, 0, 1, 0]), offset, len);
404        assert_eq!(boolean_buf1 ^ boolean_buf2, expected);
405    }
406
407    #[test]
408    fn test_boolean_not() {
409        let offset = 0;
410        let len = 40;
411
412        let buf = Buffer::from(&[0, 1, 1, 0, 0]);
413        let boolean_buf = &BooleanBuffer::new(buf, offset, len);
414
415        let expected = BooleanBuffer::new(Buffer::from(&[255, 254, 254, 255, 255]), offset, len);
416        assert_eq!(!boolean_buf, expected);
417    }
418
419    #[test]
420    fn test_boolean_from_slice_bool() {
421        let v = [true, false, false];
422        let buf = BooleanBuffer::from(&v[..]);
423        assert_eq!(buf.offset(), 0);
424        assert_eq!(buf.len(), 3);
425        assert_eq!(buf.values().len(), 1);
426        assert!(buf.value(0));
427    }
428}
arrow_buffer/buffer/boolean.rs

arrow_buffer/buffer/
boolean.rs
