Skip to content

Commit

Permalink
Add packed decoder test
Browse files Browse the repository at this point in the history
  • Loading branch information
tustvold committed Jan 12, 2022
1 parent 59846eb commit b001f11
Showing 1 changed file with 34 additions and 0 deletions.
34 changes: 34 additions & 0 deletions parquet/src/arrow/record_reader/definition_levels.rs
Original file line number Diff line number Diff line change
Expand Up @@ -392,8 +392,42 @@ fn iter_set_bits_rev(bytes: &[u8]) -> impl Iterator<Item = usize> + '_ {
mod tests {
use super::*;

use crate::encodings::rle::RleEncoder;
use rand::{thread_rng, Rng, RngCore};

#[test]
fn test_packed_decoder() {
let mut rng = thread_rng();
let len: usize = rng.gen_range(512..1024);

let mut expected = BooleanBufferBuilder::new(len);
let mut encoder = RleEncoder::new(1, 1024);
for _ in 0..len {
let bool = rng.gen_bool(0.8);
assert!(encoder.put(bool as u64).unwrap());
expected.append(bool);
}
assert_eq!(expected.len(), len);

let encoded = encoder.consume().unwrap();
let mut decoder = PackedDecoder::new(Encoding::RLE, ByteBufferPtr::new(encoded));

// Decode data in random length intervals
let mut decoded = BooleanBufferBuilder::new(len);
loop {
let remaining = len - decoded.len();
if remaining == 0 {
break;
}

let to_read = rng.gen_range(1..=remaining);
decoder.read(&mut decoded, to_read).unwrap();
}

assert_eq!(decoded.len(), len);
assert_eq!(decoded.as_slice(), expected.as_slice());
}

#[test]
fn test_bit_fns() {
let mut rng = thread_rng();
Expand Down

0 comments on commit b001f11

Please sign in to comment.