parquet/arrow/record_reader/
definition_levels.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow_array::builder::BooleanBufferBuilder;
19use arrow_buffer::Buffer;
20use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk;
21use arrow_buffer::bit_util::compress;
22use bytes::Bytes;
23
24use crate::arrow::buffer::bit_util::count_set_bits;
25use crate::basic::Encoding;
26use crate::column::reader::decoder::{
27    ColumnLevelDecoder, DefinitionLevelDecoder, DefinitionLevelDecoderImpl,
28};
29use crate::encodings::rle::run_len;
30use crate::errors::{ParquetError, Result};
31use crate::schema::types::ColumnDescPtr;
32
33enum BufferInner {
34    /// Compute levels and null mask
35    Full {
36        levels: Vec<i16>,
37        nulls: BooleanBufferBuilder,
38        max_level: i16,
39    },
40    /// Only compute null bitmask - requires max level to be 1
41    ///
42    /// This is an optimisation for the common case of a nullable scalar column, as decoding
43    /// the definition level data is only required when decoding nested structures
44    ///
45    Mask { nulls: BooleanBufferBuilder },
46}
47
48pub struct DefinitionLevelBuffer {
49    inner: BufferInner,
50
51    /// The length of this buffer
52    ///
53    /// Note: `buffer` and `builder` may contain more elements
54    len: usize,
55}
56
57impl DefinitionLevelBuffer {
58    pub fn new(desc: &ColumnDescPtr, null_mask_only: bool) -> Self {
59        let inner = match null_mask_only {
60            true => {
61                assert_eq!(
62                    desc.max_def_level(),
63                    1,
64                    "max definition level must be 1 to only compute null bitmask"
65                );
66
67                assert_eq!(
68                    desc.max_rep_level(),
69                    0,
70                    "max repetition level must be 0 to only compute null bitmask"
71                );
72
73                BufferInner::Mask {
74                    nulls: BooleanBufferBuilder::new(0),
75                }
76            }
77            false => BufferInner::Full {
78                levels: Vec::new(),
79                nulls: BooleanBufferBuilder::new(0),
80                max_level: desc.max_def_level(),
81            },
82        };
83
84        Self { inner, len: 0 }
85    }
86
87    /// Returns the built level data
88    pub fn consume_levels(&mut self) -> Option<Vec<i16>> {
89        match &mut self.inner {
90            BufferInner::Full { levels, .. } => Some(std::mem::take(levels)),
91            BufferInner::Mask { .. } => None,
92        }
93    }
94
95    /// Returns the built null bitmask, or None if all values are valid
96    pub fn consume_bitmask(&mut self) -> Option<Buffer> {
97        self.len = 0;
98        let nulls = match &mut self.inner {
99            BufferInner::Full { nulls, .. } => nulls,
100            BufferInner::Mask { nulls } => nulls,
101        };
102
103        // Always call finish() to reset the builder state for the next batch.
104        let buffer = nulls.finish().into_inner();
105
106        // If no bitmap was constructed, return None
107        if buffer.is_empty() {
108            return None;
109        }
110
111        Some(buffer)
112    }
113
114    pub fn nulls(&self) -> &BooleanBufferBuilder {
115        match &self.inner {
116            BufferInner::Full { nulls, .. } => nulls,
117            BufferInner::Mask { nulls } => nulls,
118        }
119    }
120
121    /// Returns the raw definition levels accumulated so far, if available.
122    /// Only available when the buffer is in Full mode (nested columns).
123    pub fn levels(&self) -> Option<&[i16]> {
124        match &self.inner {
125            BufferInner::Full { levels, .. } => Some(levels.as_slice()),
126            BufferInner::Mask { .. } => None,
127        }
128    }
129}
130
131/// Build a filtered validity bitmap from definition/repetition levels.
132///
133/// For each level entry where `d >= include_threshold` (when set) and
134/// `r <= max_rep` (when `rep_filter` is provided), appends one bit to
135/// `bitmap`: set when `d >= value_level`, unset otherwise.
136///
137/// Returns the number of bits appended.
138///
139/// This is the shared implementation behind the compact bitmap (leaf
140/// readers with selective padding) and the struct validity bitmap.
141/// Processing uses 64-level chunks with [`compress`] for word-at-a-time
142/// packing.
143pub(crate) fn build_filtered_validity_bitmap(
144    def_levels: &[i16],
145    rep_filter: Option<(&[i16], i16)>,
146    include_threshold: Option<i16>,
147    value_level: i16,
148    bitmap: &mut BooleanBufferBuilder,
149) -> usize {
150    const CHUNKS: usize = u64::BITS as usize;
151    // Fast path: no filtering — every def level produces a bit.
152    if include_threshold.is_none() && rep_filter.is_none() {
153        let (chunks, remainder) = def_levels.as_chunks::<CHUNKS>();
154        for chunk in chunks {
155            let mut word: u64 = 0;
156            for (i, &d) in chunk.iter().enumerate() {
157                word |= ((d >= value_level) as u64) << i;
158            }
159            bitmap.append_word(word, u64::BITS as usize);
160        }
161        for &d in remainder {
162            bitmap.append(d >= value_level);
163        }
164        return def_levels.len();
165    }
166
167    // Filtered path: build include mask + value mask per chunk, compress.
168    let mut item_count: usize = 0;
169    let (chunks, remainder) = def_levels.as_chunks::<CHUNKS>();
170    let remainder_offset = def_levels.len() - remainder.len();
171
172    for (chunk_idx, chunk) in chunks.iter().enumerate() {
173        let base = chunk_idx * u64::BITS as usize;
174        let mut include_mask: u64 = 0;
175        let mut value_mask: u64 = 0;
176        for (i, &d) in chunk.iter().enumerate() {
177            let include = include_threshold.is_none_or(|threshold| threshold <= d)
178                && rep_filter.is_none_or(|(reps, max_rep)| reps[base + i] <= max_rep);
179            include_mask |= (include as u64) << i;
180            value_mask |= ((d >= value_level) as u64) << i;
181        }
182
183        let count = include_mask.count_ones() as usize;
184        let compact = compress(value_mask, include_mask);
185        bitmap.append_word(compact, count);
186        item_count += count;
187    }
188
189    for idx in remainder_offset..def_levels.len() {
190        let d = def_levels[idx];
191        if let Some(threshold) = include_threshold
192            && d < threshold
193        {
194            continue;
195        }
196        if let Some((reps, max_rep)) = rep_filter
197            && reps[idx] > max_rep
198        {
199            continue;
200        }
201        bitmap.append(d >= value_level);
202        item_count += 1;
203    }
204
205    item_count
206}
207
208enum MaybePacked {
209    Packed(PackedDecoder),
210    Fallback(DefinitionLevelDecoderImpl),
211}
212
213pub struct DefinitionLevelBufferDecoder {
214    max_level: i16,
215    decoder: MaybePacked,
216}
217
218impl DefinitionLevelBufferDecoder {
219    pub fn new(max_level: i16, packed: bool) -> Self {
220        let decoder = match packed {
221            true => MaybePacked::Packed(PackedDecoder::new()),
222            false => MaybePacked::Fallback(DefinitionLevelDecoderImpl::new(max_level)),
223        };
224
225        Self { max_level, decoder }
226    }
227}
228
229impl ColumnLevelDecoder for DefinitionLevelBufferDecoder {
230    type Buffer = DefinitionLevelBuffer;
231
232    fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> {
233        match &mut self.decoder {
234            MaybePacked::Packed(d) => d.set_data(encoding, data),
235            MaybePacked::Fallback(d) => d.set_data(encoding, data)?,
236        }
237        Ok(())
238    }
239}
240
241impl DefinitionLevelDecoder for DefinitionLevelBufferDecoder {
242    fn read_def_levels(
243        &mut self,
244        writer: &mut Self::Buffer,
245        num_levels: usize,
246    ) -> Result<(usize, usize)> {
247        match (&mut writer.inner, &mut self.decoder) {
248            (
249                BufferInner::Full {
250                    levels,
251                    nulls,
252                    max_level,
253                },
254                MaybePacked::Fallback(decoder),
255            ) => {
256                assert_eq!(self.max_level, *max_level);
257
258                let start = levels.len();
259                let (values_read, levels_read) = decoder.read_def_levels(levels, num_levels)?;
260
261                if values_read == levels_read {
262                    // Every level is at `max_level`
263                    nulls.append_n(levels_read, true);
264                } else if values_read == 0 {
265                    // No level is at `max_level`
266                    nulls.append_n(levels_read, false);
267                } else {
268                    // Safety: slice iterator has a trusted length
269                    unsafe {
270                        nulls.extend_trusted_len(
271                            levels[start..].iter().map(|level| level == max_level),
272                        );
273                    }
274                }
275
276                Ok((values_read, levels_read))
277            }
278            (BufferInner::Mask { nulls }, MaybePacked::Packed(decoder)) => {
279                assert_eq!(self.max_level, 1);
280
281                // Fast path: if all requested levels are valid (max definition level),
282                // we can skip RLE decoding and just append all-ones to the bitmap.
283                // This is faster than decoding RLE data.
284                if let Some(count) = decoder.try_consume_all_valid(num_levels)? {
285                    nulls.append_n(count, true);
286                    return Ok((count, count)); // values_read == levels_read when all valid
287                }
288
289                // Normal path: decode RLE data into the bitmap
290                let start = nulls.len();
291                let levels_read = decoder.read(nulls, num_levels)?;
292
293                let values_read = count_set_bits(nulls.as_slice(), start..start + levels_read);
294                Ok((values_read, levels_read))
295            }
296            _ => unreachable!("inconsistent null mask"),
297        }
298    }
299
300    fn skip_def_levels(&mut self, num_levels: usize) -> Result<(usize, usize)> {
301        match &mut self.decoder {
302            MaybePacked::Fallback(decoder) => decoder.skip_def_levels(num_levels),
303            MaybePacked::Packed(decoder) => decoder.skip(num_levels),
304        }
305    }
306}
307
308/// An optimized decoder for decoding [RLE] and [BIT_PACKED] data with a bit width of 1
309/// directly into a bitmask
310///
311/// This is significantly faster than decoding the data into `[i16]` and then computing
312/// a bitmask from this, as not only can it skip this buffer allocation and construction,
313/// but it can exploit properties of the encoded data to reduce work further
314///
315/// In particular:
316///
317/// * Packed runs are already bitmask encoded and can simply be appended
318/// * Runs of 1 or 0 bits can be efficiently appended with byte (or larger) operations
319///
320/// [RLE]: https://github.com/apache/parquet-format/blob/master/Encodings.md#run-length-encoding--bit-packing-hybrid-rle--3
321/// [BIT_PACKED]: https://github.com/apache/parquet-format/blob/master/Encodings.md#bit-packed-deprecated-bit_packed--4
322struct PackedDecoder {
323    data: Bytes,
324    data_offset: usize,
325    rle_left: usize,
326    rle_value: bool,
327    packed_count: usize,
328    packed_offset: usize,
329}
330
331impl PackedDecoder {
332    fn next_rle_block(&mut self) -> Result<()> {
333        let indicator_value = self.decode_header()?;
334        let num_values = run_len(indicator_value as u64)? as usize;
335        if indicator_value & 1 == 1 {
336            self.packed_count = num_values;
337            self.packed_offset = 0;
338        } else {
339            self.rle_left = num_values;
340            let byte = *self.data.as_ref().get(self.data_offset).ok_or_else(|| {
341                ParquetError::EOF(
342                    "unexpected end of file whilst decoding definition levels rle value".into(),
343                )
344            })?;
345
346            self.data_offset += 1;
347            self.rle_value = byte != 0;
348        }
349        Ok(())
350    }
351
352    /// Decodes a VLQ encoded little endian integer and returns it
353    fn decode_header(&mut self) -> Result<i64> {
354        let mut offset = 0;
355        let mut v: i64 = 0;
356        while offset < 10 {
357            let byte = *self
358                .data
359                .as_ref()
360                .get(self.data_offset + offset)
361                .ok_or_else(|| {
362                    ParquetError::EOF(
363                        "unexpected end of file whilst decoding definition levels rle header"
364                            .into(),
365                    )
366                })?;
367
368            v |= ((byte & 0x7F) as i64) << (offset * 7);
369            offset += 1;
370            if byte & 0x80 == 0 {
371                self.data_offset += offset;
372                return Ok(v);
373            }
374        }
375        Err(general_err!("too many bytes for VLQ"))
376    }
377}
378
379impl PackedDecoder {
380    fn new() -> Self {
381        Self {
382            data: Bytes::from(vec![]),
383            data_offset: 0,
384            rle_left: 0,
385            rle_value: false,
386            packed_count: 0,
387            packed_offset: 0,
388        }
389    }
390
391    fn set_data(&mut self, encoding: Encoding, data: Bytes) {
392        self.rle_left = 0;
393        self.rle_value = false;
394        self.packed_offset = 0;
395        self.packed_count = match encoding {
396            Encoding::RLE => 0,
397            #[expect(deprecated)]
398            Encoding::BIT_PACKED => data.len() * 8,
399            _ => unreachable!("invalid level encoding: {}", encoding),
400        };
401        self.data = data;
402        self.data_offset = 0;
403    }
404
405    /// Try to consume `len` levels if all are valid (max definition level).
406    ///
407    /// Returns `Ok(Some(count))` if successfully consumed `count` all-valid levels.
408    /// Returns `Ok(None)` if there are any nulls or packed data that prevents fast path.
409    ///
410    /// Note: On `None`, the decoder state may have advanced to the next RLE block,
411    /// but only if `rle_left` was zero (i.e., the block would have been loaded
412    /// on the next read anyway).
413    fn try_consume_all_valid(&mut self, len: usize) -> Result<Option<usize>> {
414        // If no active run and no packed data pending, try to parse the next RLE block
415        if self.rle_left == 0 && self.packed_count == self.packed_offset {
416            if self.data_offset < self.data.len() {
417                self.next_rle_block()?;
418            } else {
419                // No more data available
420                return Ok(None);
421            }
422        }
423
424        // Fast path only works when we have an active RLE run of true values
425        // that covers the entire requested length.
426        if self.rle_left >= len && self.rle_value {
427            self.rle_left -= len;
428            return Ok(Some(len));
429        }
430
431        // Any other case (null run, packed data, or insufficient length)
432        // falls back to normal path
433        Ok(None)
434    }
435
436    fn read(&mut self, buffer: &mut BooleanBufferBuilder, len: usize) -> Result<usize> {
437        let mut read = 0;
438        while read != len {
439            if self.rle_left != 0 {
440                let to_read = self.rle_left.min(len - read);
441                buffer.append_n(to_read, self.rle_value);
442                self.rle_left -= to_read;
443                read += to_read;
444            } else if self.packed_count != self.packed_offset {
445                let to_read = (self.packed_count - self.packed_offset).min(len - read);
446                let offset = self.data_offset * 8 + self.packed_offset;
447                buffer.append_packed_range(offset..offset + to_read, self.data.as_ref());
448                self.packed_offset += to_read;
449                read += to_read;
450
451                if self.packed_offset == self.packed_count {
452                    self.data_offset += self.packed_count / 8;
453                }
454            } else if self.data_offset == self.data.len() {
455                break;
456            } else {
457                self.next_rle_block()?
458            }
459        }
460        Ok(read)
461    }
462
463    /// Skips `level_num` definition levels
464    ///
465    /// Returns the number of values skipped and the number of levels skipped
466    fn skip(&mut self, level_num: usize) -> Result<(usize, usize)> {
467        let mut skipped_value = 0;
468        let mut skipped_level = 0;
469        while skipped_level != level_num {
470            if self.rle_left != 0 {
471                let to_skip = self.rle_left.min(level_num - skipped_level);
472                self.rle_left -= to_skip;
473                skipped_level += to_skip;
474                if self.rle_value {
475                    skipped_value += to_skip;
476                }
477            } else if self.packed_count != self.packed_offset {
478                let to_skip =
479                    (self.packed_count - self.packed_offset).min(level_num - skipped_level);
480                let offset = self.data_offset * 8 + self.packed_offset;
481                let bit_chunk = UnalignedBitChunk::new(self.data.as_ref(), offset, to_skip);
482                skipped_value += bit_chunk.count_ones();
483                self.packed_offset += to_skip;
484                skipped_level += to_skip;
485                if self.packed_offset == self.packed_count {
486                    self.data_offset += self.packed_count / 8;
487                }
488            } else if self.data_offset == self.data.len() {
489                break;
490            } else {
491                self.next_rle_block()?
492            }
493        }
494        Ok((skipped_value, skipped_level))
495    }
496}
497
498#[cfg(test)]
499mod tests {
500    use super::*;
501
502    use crate::encodings::rle::RleEncoder;
503    use crate::util::bit_util::BitWriter;
504    use rand::{RngExt, rng};
505
506    #[test]
507    fn test_build_validity_bitmap_unfiltered_word_chunk() {
508        // 65 levels forces the unfiltered path to process one full u64 word
509        // with append_word, plus a remainder bit.
510        let def_levels = (0..65)
511            .map(|i| if i % 3 == 0 { 2 } else { 1 })
512            .collect::<Vec<_>>();
513        let mut bitmap = BooleanBufferBuilder::new(0);
514
515        assert_eq!(
516            build_filtered_validity_bitmap(&def_levels, None, None, 2, &mut bitmap),
517            def_levels.len()
518        );
519
520        let bitmap = bitmap.finish();
521        for (idx, def) in def_levels.iter().enumerate() {
522            assert_eq!(bitmap.value(idx), *def >= 2);
523        }
524    }
525
526    #[test]
527    fn test_packed_decoder() {
528        let mut rng = rng();
529        let len: usize = rng.random_range(512..1024);
530
531        let mut expected = BooleanBufferBuilder::new(len);
532        let mut encoder = RleEncoder::new(1, 1024);
533        for _ in 0..len {
534            let bool = rng.random_bool(0.8);
535            encoder.put(bool as u64);
536            expected.append(bool);
537        }
538        assert_eq!(expected.len(), len);
539
540        let encoded = encoder.consume();
541        let mut decoder = PackedDecoder::new();
542        decoder.set_data(Encoding::RLE, encoded.into());
543
544        // Decode data in random length intervals
545        let mut decoded = BooleanBufferBuilder::new(len);
546        loop {
547            let remaining = len - decoded.len();
548            if remaining == 0 {
549                break;
550            }
551
552            let to_read = rng.random_range(1..=remaining);
553            decoder.read(&mut decoded, to_read).unwrap();
554        }
555
556        assert_eq!(decoded.len(), len);
557        assert_eq!(decoded.as_slice(), expected.as_slice());
558    }
559
560    #[test]
561    fn test_packed_decoder_skip() {
562        let mut rng = rng();
563        let len: usize = rng.random_range(512..1024);
564
565        let mut expected = BooleanBufferBuilder::new(len);
566        let mut encoder = RleEncoder::new(1, 1024);
567
568        let mut total_value = 0;
569        for _ in 0..len {
570            let bool = rng.random_bool(0.8);
571            encoder.put(bool as u64);
572            expected.append(bool);
573            if bool {
574                total_value += 1;
575            }
576        }
577        assert_eq!(expected.len(), len);
578
579        let encoded = encoder.consume();
580        let mut decoder = PackedDecoder::new();
581        decoder.set_data(Encoding::RLE, encoded.into());
582
583        let mut skip_value = 0;
584        let mut read_value = 0;
585        let mut skip_level = 0;
586        let mut read_level = 0;
587
588        loop {
589            let offset = skip_level + read_level;
590            let remaining_levels = len - offset;
591            if remaining_levels == 0 {
592                break;
593            }
594            let to_read_or_skip_level = rng.random_range(1..=remaining_levels);
595            if rng.random_bool(0.5) {
596                let (skip_val_num, skip_level_num) = decoder.skip(to_read_or_skip_level).unwrap();
597                skip_value += skip_val_num;
598                skip_level += skip_level_num
599            } else {
600                let mut decoded = BooleanBufferBuilder::new(to_read_or_skip_level);
601                let read_level_num = decoder.read(&mut decoded, to_read_or_skip_level).unwrap();
602                read_level += read_level_num;
603                for i in 0..read_level_num {
604                    assert!(!decoded.is_empty());
605                    //check each read bit
606                    let read_bit = decoded.get_bit(i);
607                    if read_bit {
608                        read_value += 1;
609                    }
610                    let expect_bit = expected.get_bit(i + offset);
611                    assert_eq!(read_bit, expect_bit);
612                }
613            }
614        }
615        assert_eq!(read_level + skip_level, len);
616        assert_eq!(read_value + skip_value, total_value);
617    }
618
619    #[test]
620    fn test_try_consume_all_valid() {
621        // Test with all-valid data (all 1s) - single RLE run
622        let len = 100;
623        let mut encoder = RleEncoder::new(1, 1024);
624        for _ in 0..len {
625            encoder.put(1); // all valid
626        }
627        let encoded = encoder.consume();
628        let mut decoder = PackedDecoder::new();
629        decoder.set_data(Encoding::RLE, encoded.into());
630
631        // try_consume_all_valid now parses the RLE block itself, no need to read first
632        let result = decoder.try_consume_all_valid(len).unwrap();
633        assert_eq!(result, Some(len));
634
635        // Test with all-null data (all 0s)
636        let mut encoder = RleEncoder::new(1, 1024);
637        for _ in 0..len {
638            encoder.put(0); // all null
639        }
640        let encoded = encoder.consume();
641        let mut decoder = PackedDecoder::new();
642        decoder.set_data(Encoding::RLE, encoded.into());
643
644        // Should return None because rle_value is false (all nulls)
645        let result = decoder.try_consume_all_valid(len).unwrap();
646        assert_eq!(result, None);
647
648        // Test when requesting more than available in current RLE run
649        let mut encoder = RleEncoder::new(1, 1024);
650        for _ in 0..10 {
651            encoder.put(1); // small run of valid
652        }
653        for _ in 0..10 {
654            encoder.put(0); // followed by nulls
655        }
656        let encoded = encoder.consume();
657        let mut decoder = PackedDecoder::new();
658        decoder.set_data(Encoding::RLE, encoded.into());
659
660        // Request more than the valid run - should return None
661        // (because we don't look ahead to next block)
662        let result = decoder.try_consume_all_valid(20).unwrap();
663        assert_eq!(result, None);
664
665        // Reset decoder and try requesting within the run
666        decoder.set_data(Encoding::RLE, {
667            let mut encoder = RleEncoder::new(1, 1024);
668            for _ in 0..10 {
669                encoder.put(1);
670            }
671            for _ in 0..10 {
672                encoder.put(0);
673            }
674            encoder.consume().into()
675        });
676
677        let result = decoder.try_consume_all_valid(5).unwrap();
678        assert_eq!(result, Some(5));
679
680        // After skipping 5, we should have 5 left in the valid RLE run
681        let result = decoder.try_consume_all_valid(5).unwrap();
682        assert_eq!(result, Some(5));
683
684        // Now the valid run is exhausted, next call should parse the null run and return None
685        let result = decoder.try_consume_all_valid(5).unwrap();
686        assert_eq!(result, None);
687    }
688
689    #[test]
690    fn test_packed_decoder_invalid_run_len() {
691        let mut writer = BitWriter::new(16);
692        writer.put_vlq_int(u64::MAX);
693        let mut decoder = PackedDecoder::new();
694        decoder.set_data(Encoding::RLE, writer.consume().into());
695
696        let mut decoded = BooleanBufferBuilder::new(8);
697        let err = decoder.read(&mut decoded, 8).unwrap_err();
698        assert_eq!(
699            err.to_string(),
700            "Parquet error: Invalid or corrupted RLE run length 18446744073709551615. Max allowed is 2147483647"
701        );
702    }
703}
parquet/arrow/record_reader/definition_levels.rs

parquet/arrow/record_reader/
definition_levels.rs
