reifydb_value/util/
bitmap.rs1use std::borrow::Cow;
5
6use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer};
7
8pub fn slice(bits: &BooleanBuffer, start: usize, end: usize) -> BooleanBuffer {
9 let end = end.min(bits.len());
10 let start = start.min(end);
11 bits.slice(start, end - start)
12}
13
14pub fn filter(bits: &BooleanBuffer, mask: &BooleanBuffer) -> BooleanBuffer {
15 let mut kept = BooleanBufferBuilder::new(mask.count_set_bits());
16 for i in mask.set_indices().take_while(|&i| i < bits.len()) {
17 kept.append(bits.value(i));
18 }
19 kept.finish()
20}
21
22pub fn reorder(bits: &BooleanBuffer, indices: &[usize]) -> BooleanBuffer {
23 BooleanBuffer::collect_bool(indices.len(), |i| indices[i] < bits.len() && bits.value(indices[i]))
24}
25
26pub fn resize(bits: &BooleanBuffer, len: usize) -> BooleanBuffer {
27 if bits.len() >= len {
28 return bits.slice(0, len);
29 }
30 let mut resized = BooleanBufferBuilder::new(len);
31 resized.append_buffer(bits);
32 resized.append_n(len - bits.len(), false);
33 resized.finish()
34}
35
36pub fn assert_nulls_len(nulls: Option<&NullBuffer>, len: usize) {
37 if let Some(nulls) = nulls {
38 assert_eq!(nulls.len(), len, "validity of {} rows does not match an array of {len} rows", nulls.len());
39 }
40}
41
42pub fn and_nulls(left: &NullBuffer, right: &NullBuffer) -> NullBuffer {
43 assert_eq!(left.len(), right.len(), "validity of {} and {} rows cannot be combined", left.len(), right.len());
44 NullBuffer::new(left.inner() & right.inner())
45}
46
47pub fn slice_nulls(nulls: Option<&NullBuffer>, start: usize, end: usize) -> Option<NullBuffer> {
48 nulls.map(|nulls| NullBuffer::new(slice(nulls.inner(), start, end)))
49}
50
51pub fn reorder_nulls(nulls: Option<&NullBuffer>, indices: &[usize]) -> Option<NullBuffer> {
52 nulls.map(|nulls| NullBuffer::new(reorder(nulls.inner(), indices)))
53}
54
55fn clear_tail(bytes: &mut [u8], len: usize) {
56 let used = len % 8;
57 if used != 0
58 && let Some(last) = bytes.last_mut()
59 {
60 *last &= (1u8 << used) - 1;
61 }
62}
63
64pub fn packed_bytes(bits: &BooleanBuffer) -> Cow<'_, [u8]> {
65 let len = bits.len();
66 let byte_len = len.div_ceil(8);
67 if bits.offset().is_multiple_of(8) {
68 let start = bits.offset() / 8;
69 let bytes = &bits.values()[start..start + byte_len];
70 let used = len % 8;
71 if used == 0 || bytes[byte_len - 1] >> used == 0 {
72 return Cow::Borrowed(bytes);
73 }
74 }
75 let mut packed = bits.sliced().as_slice()[..byte_len].to_vec();
76 clear_tail(&mut packed, len);
77 Cow::Owned(packed)
78}
79
80#[cfg(test)]
81mod tests {
82 use std::borrow::Cow;
83
84 use arrow_buffer::BooleanBuffer;
85
86 use super::{filter, packed_bytes, reorder, slice};
87
88 fn pattern(len: usize) -> Vec<bool> {
89 (0..len).map(|i| i % 3 != 1).collect()
90 }
91
92 fn bits_of(bits: &BooleanBuffer) -> Vec<bool> {
93 bits.iter().collect()
94 }
95
96 #[test]
97 fn slice_clamps_out_of_range_bounds() {
98 let bits = BooleanBuffer::from(pattern(10));
100 let tail = slice(&bits, 4, 100);
101 assert_eq!(tail.len(), 6);
102 assert_eq!(bits_of(&tail), pattern(10)[4..].to_vec());
103 assert_eq!(slice(&bits, 12, 20).len(), 0);
104 }
105
106 #[test]
107 fn slice_takes_end_not_length_as_its_second_argument() {
108 let bits = BooleanBuffer::from(pattern(10));
110 let mid = slice(&bits, 2, 5);
111 assert_eq!(mid.len(), 3);
112 assert_eq!(bits_of(&mid), pattern(10)[2..5].to_vec());
113 }
114
115 #[test]
116 fn slice_with_start_past_end_is_empty() {
117 let bits = BooleanBuffer::from(pattern(10));
119 assert_eq!(slice(&bits, 7, 3).len(), 0);
120 assert_eq!(slice(&bits, 30, 20).len(), 0);
121 }
122
123 #[test]
124 fn filter_short_mask_drops_the_tail_and_long_mask_is_ignored_past_len() {
125 let bits = BooleanBuffer::from(vec![true, false, true, true]);
127 let short = filter(&bits, &BooleanBuffer::from(vec![true, true]));
128 assert_eq!(bits_of(&short), vec![true, false]);
129 let long = filter(&bits, &BooleanBuffer::from(vec![false, true, false, true, true, true]));
130 assert_eq!(bits_of(&long), vec![false, true]);
131 }
132
133 #[test]
134 fn reorder_out_of_range_gives_false() {
135 let bits = BooleanBuffer::from(vec![true, true]);
137 assert_eq!(bits_of(&reorder(&bits, &[1, 5, 0])), vec![true, false, true]);
138 }
139
140 #[test]
141 fn packed_bytes_of_a_view_have_a_zero_tail() {
142 let bits = BooleanBuffer::new_set(16);
144 assert_eq!(packed_bytes(&slice(&bits, 3, 9)).as_ref(), &[0b0011_1111]);
145 let aligned = slice(&bits, 8, 13);
146 assert_eq!(packed_bytes(&aligned).as_ref(), &[0b0001_1111]);
147 assert!(matches!(packed_bytes(&aligned), Cow::Owned(_)));
148 }
149
150 #[test]
151 fn packed_bytes_after_not_have_a_zero_tail() {
152 let bits = BooleanBuffer::from(vec![true, false, true]);
154 let inverted = !&bits;
155 assert_eq!(packed_bytes(&inverted).as_ref(), &[0b0000_0010]);
156 }
157
158 #[test]
159 fn packed_bytes_borrow_a_clean_buffer() {
160 let bits = BooleanBuffer::from(pattern(12));
162 let packed = packed_bytes(&bits);
163 assert!(matches!(packed, Cow::Borrowed(_)));
164 assert_eq!(packed.len(), 2);
165 }
166}