Function test_skip_rep_levels

parquet/src/column/reader/decoder.rs:517–564

()

    #[test]
    fn test_skip_rep_levels() {
        for _ in 0..10 {
            let mut rng = rng();
            let total_len = 10000_usize;
            let mut encoded: Vec<i16> = (0..total_len).map(|_| rng.random_range(0..5)).collect();
            encoded[0] = 0;
            let mut encoder = RleEncoder::new(3, 1024);
            for v in &encoded {
                encoder.put(*v as _)
            }
            let data = Bytes::from(encoder.consume());

            let mut decoder = RepetitionLevelDecoderImpl::new(5);
            decoder.set_data(Encoding::RLE, data).unwrap();

            let total_records = encoded.iter().filter(|x| **x == 0).count();
            let mut remaining_records = total_records;
            let mut remaining_levels = encoded.len();
            loop {
                let skip = rng.random_bool(0.5);
                let records = rng.random_range(1..=remaining_records.min(5));
                let (records_read, levels_read) = if skip {
                    decoder.skip_rep_levels(records, remaining_levels).unwrap()
                } else {
                    let mut decoded = Vec::new();
                    let (records_read, levels_read) = decoder
                        .read_rep_levels(&mut decoded, records, remaining_levels)
                        .unwrap();

                    assert_eq!(
                        decoded,
                        encoded[encoded.len() - remaining_levels..][..levels_read]
                    );
                    (records_read, levels_read)
                };

                remaining_levels = remaining_levels.checked_sub(levels_read).unwrap();
                if remaining_levels == 0 {
                    assert_eq!(records_read + 1, records);
                    assert_eq!(records, remaining_records);
                    break;
                }
                assert_eq!(records_read, records);
                remaining_records -= records;
                assert_ne!(remaining_records, 0);
            }
        }
    }
}