use std::iter::Peekable;
use crate::decoder::DecodeError;
pub(crate) struct BitmapDecodeIterator<B: Iterator, I> {
bitmap: Peekable<B>,
values: I,
len: usize,
offset: usize,
}
impl<'b, B, I> BitmapDecodeIterator<B, I>
where
B: Iterator<Item = &'b u8>,
{
pub(crate) fn new(bitmap: B, values: I, len: usize) -> Result<Self, DecodeError> {
let (bitmap_len, _) = bitmap.size_hint();
if bitmap_len * 8 < len {
return Err(DecodeError::LengthMismatch);
}
Ok(Self {
bitmap: bitmap.peekable(),
values,
len,
offset: 0,
})
}
}
impl<'b, B, I> Iterator for BitmapDecodeIterator<B, I>
where
B: Iterator<Item = &'b u8>,
I: Iterator<Item = f32>,
{
type Item = f32;
fn next(&mut self) -> Option<Self::Item> {
if self.len == 0 {
return None;
}
self.len -= 1;
let offset = self.offset;
let byte = if self.offset < 7 {
self.offset += 1;
self.bitmap.peek()?
} else {
self.offset = 0;
self.bitmap.next()?
};
if has_zero_at_offset(byte, &offset) {
Some(f32::NAN)
} else {
self.values.next()
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let size = self.len;
(size, Some(size))
}
}
const MASK: u8 = 0b10000000;
fn has_zero_at_offset(byte: &u8, offset: &usize) -> bool {
let masked = byte & (MASK >> offset);
masked == 0
}
pub(crate) fn dummy_bitmap_for_nonnullable_data(num_points: usize) -> Vec<u8> {
let size = super::helpers::num_octets(num_points);
vec![0b11111111u8; size]
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn bitmap_iterator_works() {
let bitmap = [0b01001100u8, 0b01110000, 0b11110000];
let values = (0..10).map(|n| n as f32).collect::<Vec<_>>();
let values = values.into_iter();
let iter = BitmapDecodeIterator::new(bitmap.iter(), values, 24).unwrap();
let actual = iter.collect::<Vec<_>>();
let expected = [
f32::NAN,
0.0,
f32::NAN,
f32::NAN,
1.0,
2.0,
f32::NAN,
f32::NAN,
f32::NAN,
3.0,
4.0,
5.0,
f32::NAN,
f32::NAN,
f32::NAN,
f32::NAN,
6.0,
7.0,
8.0,
9.0,
f32::NAN,
f32::NAN,
f32::NAN,
f32::NAN,
];
assert_eq!(actual.len(), expected.len());
actual
.iter()
.zip(expected.iter())
.all(|(a, b)| (a.is_nan() && b.is_nan()) || (a == b));
}
#[test]
fn bitmap_iterator_length() {
let bitmap = [0b01001100u8, 0b01110000];
let values = (0..6).map(|n| n as f32).collect::<Vec<_>>();
let values = values.into_iter();
let iter = BitmapDecodeIterator::new(bitmap.iter(), values, 12).unwrap();
let actual = iter.collect::<Vec<_>>();
let expected = [
f32::NAN,
0.0,
f32::NAN,
f32::NAN,
1.0,
2.0,
f32::NAN,
f32::NAN,
f32::NAN,
3.0,
4.0,
5.0,
];
assert_eq!(actual.len(), expected.len());
actual
.iter()
.zip(expected.iter())
.all(|(a, b)| (a.is_nan() && b.is_nan()) || (a == b));
}
#[test]
fn bitmap_iterator_size_hint() {
let bitmap = [0b01001100u8, 0b01110000, 0b11110000];
let values = (0..10).map(|n| n as f32).collect::<Vec<_>>();
let values = values.into_iter();
let mut iter = BitmapDecodeIterator::new(bitmap.iter(), values, 24).unwrap();
assert_eq!(iter.size_hint(), (24, Some(24)));
let _ = iter.next();
assert_eq!(iter.size_hint(), (23, Some(23)));
}
}