Skip to main content

reifydb_value/util/
bitmap.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 ReifyDB
3
4use 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		// Arrow's slice(offset, len) panics past the end; ours must clamp end to len like the old bitvec.
99		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		// Reading the second argument as a length would give 5 rows here instead of 3.
109		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		// start > end must clamp start to end instead of underflowing end - start.
118		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		// A mask shorter than the bits must drop the rows past it; a longer one must not read past len.
126		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		// An out of range index must read as false, as the old containers did, never panic.
136		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		// A view past bit 0 must repack from bit 0 and zero every bit past len.
143		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		// Arrow's `!` sets the bits past len; a raw byte reader must never see them.
153		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		// A fresh buffer already holds the packed form, so it must not be copied.
161		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}