parquet/arrow/record_reader/

definition_levels.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use arrow_array::builder::BooleanBufferBuilder;
use arrow_buffer::Buffer;
use arrow_buffer::bit_chunk_iterator::UnalignedBitChunk;
use bytes::Bytes;

use crate::arrow::buffer::bit_util::count_set_bits;
use crate::basic::Encoding;
use crate::column::reader::decoder::{
    ColumnLevelDecoder, DefinitionLevelDecoder, DefinitionLevelDecoderImpl,
};
use crate::errors::{ParquetError, Result};
use crate::schema::types::ColumnDescPtr;

enum BufferInner {
    /// Compute levels and null mask
    Full {
        levels: Vec<i16>,
        nulls: BooleanBufferBuilder,
        max_level: i16,
    },
    /// Only compute null bitmask - requires max level to be 1
    ///
    /// This is an optimisation for the common case of a nullable scalar column, as decoding
    /// the definition level data is only required when decoding nested structures
    ///
    Mask { nulls: BooleanBufferBuilder },
}

pub struct DefinitionLevelBuffer {
    inner: BufferInner,

    /// The length of this buffer
    ///
    /// Note: `buffer` and `builder` may contain more elements
    len: usize,
}

impl DefinitionLevelBuffer {
    pub fn new(desc: &ColumnDescPtr, null_mask_only: bool) -> Self {
        let inner = match null_mask_only {
            true => {
                assert_eq!(
                    desc.max_def_level(),
                    1,
                    "max definition level must be 1 to only compute null bitmask"
                );

                assert_eq!(
                    desc.max_rep_level(),
                    0,
                    "max repetition level must be 0 to only compute null bitmask"
                );

                BufferInner::Mask {
                    nulls: BooleanBufferBuilder::new(0),
                }
            }
            false => BufferInner::Full {
                levels: Vec::new(),
                nulls: BooleanBufferBuilder::new(0),
                max_level: desc.max_def_level(),
            },
        };

        Self { inner, len: 0 }
    }

    /// Returns the built level data
    pub fn consume_levels(&mut self) -> Option<Vec<i16>> {
        match &mut self.inner {
            BufferInner::Full { levels, .. } => Some(std::mem::take(levels)),
            BufferInner::Mask { .. } => None,
        }
    }

    /// Returns the built null bitmask, or None if all values are valid
    pub fn consume_bitmask(&mut self) -> Option<Buffer> {
        self.len = 0;
        let nulls = match &mut self.inner {
            BufferInner::Full { nulls, .. } => nulls,
            BufferInner::Mask { nulls } => nulls,
        };

        // Always call finish() to reset the builder state for the next batch.
        let buffer = nulls.finish().into_inner();

        // If no bitmap was constructed, return None
        if buffer.is_empty() {
            return None;
        }

        Some(buffer)
    }

    pub fn nulls(&self) -> &BooleanBufferBuilder {
        match &self.inner {
            BufferInner::Full { nulls, .. } => nulls,
            BufferInner::Mask { nulls } => nulls,
        }
    }

    /// Returns the raw definition levels accumulated so far, if available.
    /// Only available when the buffer is in Full mode (nested columns).
    pub fn levels(&self) -> Option<&[i16]> {
        match &self.inner {
            BufferInner::Full { levels, .. } => Some(levels.as_slice()),
            BufferInner::Mask { .. } => None,
        }
    }
}

/// Build a filtered validity bitmap from definition/repetition levels.
///
/// For each level entry where `d >= include_threshold` (when set) and
/// `r <= max_rep` (when `rep_filter` is provided), appends one bit to
/// `bitmap`: set when `d >= value_level`, unset otherwise.
///
/// Returns the number of bits appended.
///
/// This is the shared implementation behind the compact bitmap (leaf
/// readers with selective padding) and the struct validity bitmap.
/// Processing uses 64-level chunks with [`compress`] for word-at-a-time
/// packing.
pub(crate) fn build_filtered_validity_bitmap(
    def_levels: &[i16],
    rep_filter: Option<(&[i16], i16)>,
    include_threshold: Option<i16>,
    value_level: i16,
    bitmap: &mut BooleanBufferBuilder,
) -> usize {
    // Fast path: no filtering — every def level produces a bit.
    if include_threshold.is_none() && rep_filter.is_none() {
        let chunks = def_levels.chunks_exact(u64::BITS as usize);
        let remainder = chunks.remainder();
        for chunk in chunks {
            let mut word: u64 = 0;
            for (i, &d) in chunk.iter().enumerate() {
                word |= ((d >= value_level) as u64) << i;
            }
            bitmap.append_word(word, u64::BITS as usize);
        }
        for &d in remainder {
            bitmap.append(d >= value_level);
        }
        return def_levels.len();
    }

    // Filtered path: build include mask + value mask per chunk, compress.
    let mut item_count: usize = 0;
    let chunks = def_levels.chunks_exact(u64::BITS as usize);
    let remainder_offset = def_levels.len() - chunks.remainder().len();

    for (chunk_idx, chunk) in chunks.enumerate() {
        let base = chunk_idx * u64::BITS as usize;
        let mut include_mask: u64 = 0;
        let mut value_mask: u64 = 0;
        for (i, &d) in chunk.iter().enumerate() {
            let mut include = true;
            if let Some(threshold) = include_threshold {
                if d < threshold {
                    include = false;
                }
            }
            if include {
                if let Some((reps, max_rep)) = rep_filter {
                    if reps[base + i] > max_rep {
                        include = false;
                    }
                }
            }
            include_mask |= (include as u64) << i;
            value_mask |= ((d >= value_level) as u64) << i;
        }

        let count = include_mask.count_ones() as usize;
        let compact = compress(value_mask, include_mask);
        bitmap.append_word(compact, count);
        item_count += count;
    }

    for idx in remainder_offset..def_levels.len() {
        let d = def_levels[idx];
        if let Some(threshold) = include_threshold {
            if d < threshold {
                continue;
            }
        }
        if let Some((reps, max_rep)) = rep_filter {
            if reps[idx] > max_rep {
                continue;
            }
        }
        bitmap.append(d >= value_level);
        item_count += 1;
    }

    item_count
}

use crate::util::bit_util::compress;

enum MaybePacked {
    Packed(PackedDecoder),
    Fallback(DefinitionLevelDecoderImpl),
}

pub struct DefinitionLevelBufferDecoder {
    max_level: i16,
    decoder: MaybePacked,
}

impl DefinitionLevelBufferDecoder {
    pub fn new(max_level: i16, packed: bool) -> Self {
        let decoder = match packed {
            true => MaybePacked::Packed(PackedDecoder::new()),
            false => MaybePacked::Fallback(DefinitionLevelDecoderImpl::new(max_level)),
        };

        Self { max_level, decoder }
    }
}

impl ColumnLevelDecoder for DefinitionLevelBufferDecoder {
    type Buffer = DefinitionLevelBuffer;

    fn set_data(&mut self, encoding: Encoding, data: Bytes) -> Result<()> {
        match &mut self.decoder {
            MaybePacked::Packed(d) => d.set_data(encoding, data),
            MaybePacked::Fallback(d) => d.set_data(encoding, data)?,
        };
        Ok(())
    }
}

impl DefinitionLevelDecoder for DefinitionLevelBufferDecoder {
    fn read_def_levels(
        &mut self,
        writer: &mut Self::Buffer,
        num_levels: usize,
    ) -> Result<(usize, usize)> {
        match (&mut writer.inner, &mut self.decoder) {
            (
                BufferInner::Full {
                    levels,
                    nulls,
                    max_level,
                },
                MaybePacked::Fallback(decoder),
            ) => {
                assert_eq!(self.max_level, *max_level);

                let start = levels.len();
                let (values_read, levels_read) = decoder.read_def_levels(levels, num_levels)?;

                // Safety: slice iterator has a trusted length
                unsafe {
                    nulls
                        .extend_trusted_len(levels[start..].iter().map(|level| level == max_level));
                }

                Ok((values_read, levels_read))
            }
            (BufferInner::Mask { nulls }, MaybePacked::Packed(decoder)) => {
                assert_eq!(self.max_level, 1);

                // Fast path: if all requested levels are valid (max definition level),
                // we can skip RLE decoding and just append all-ones to the bitmap.
                // This is faster than decoding RLE data.
                if let Some(count) = decoder.try_consume_all_valid(num_levels)? {
                    nulls.append_n(count, true);
                    return Ok((count, count)); // values_read == levels_read when all valid
                }

                // Normal path: decode RLE data into the bitmap
                let start = nulls.len();
                let levels_read = decoder.read(nulls, num_levels)?;

                let values_read = count_set_bits(nulls.as_slice(), start..start + levels_read);
                Ok((values_read, levels_read))
            }
            _ => unreachable!("inconsistent null mask"),
        }
    }

    fn skip_def_levels(&mut self, num_levels: usize) -> Result<(usize, usize)> {
        match &mut self.decoder {
            MaybePacked::Fallback(decoder) => decoder.skip_def_levels(num_levels),
            MaybePacked::Packed(decoder) => decoder.skip(num_levels),
        }
    }
}

/// An optimized decoder for decoding [RLE] and [BIT_PACKED] data with a bit width of 1
/// directly into a bitmask
///
/// This is significantly faster than decoding the data into `[i16]` and then computing
/// a bitmask from this, as not only can it skip this buffer allocation and construction,
/// but it can exploit properties of the encoded data to reduce work further
///
/// In particular:
///
/// * Packed runs are already bitmask encoded and can simply be appended
/// * Runs of 1 or 0 bits can be efficiently appended with byte (or larger) operations
///
/// [RLE]: https://github.com/apache/parquet-format/blob/master/Encodings.md#run-length-encoding--bit-packing-hybrid-rle--3
/// [BIT_PACKED]: https://github.com/apache/parquet-format/blob/master/Encodings.md#bit-packed-deprecated-bit_packed--4
struct PackedDecoder {
    data: Bytes,
    data_offset: usize,
    rle_left: usize,
    rle_value: bool,
    packed_count: usize,
    packed_offset: usize,
}

impl PackedDecoder {
    fn next_rle_block(&mut self) -> Result<()> {
        let indicator_value = self.decode_header()?;
        if indicator_value & 1 == 1 {
            let len = (indicator_value >> 1) as usize;
            self.packed_count = len * 8;
            self.packed_offset = 0;
        } else {
            self.rle_left = (indicator_value >> 1) as usize;
            let byte = *self.data.as_ref().get(self.data_offset).ok_or_else(|| {
                ParquetError::EOF(
                    "unexpected end of file whilst decoding definition levels rle value".into(),
                )
            })?;

            self.data_offset += 1;
            self.rle_value = byte != 0;
        }
        Ok(())
    }

    /// Decodes a VLQ encoded little endian integer and returns it
    fn decode_header(&mut self) -> Result<i64> {
        let mut offset = 0;
        let mut v: i64 = 0;
        while offset < 10 {
            let byte = *self
                .data
                .as_ref()
                .get(self.data_offset + offset)
                .ok_or_else(|| {
                    ParquetError::EOF(
                        "unexpected end of file whilst decoding definition levels rle header"
                            .into(),
                    )
                })?;

            v |= ((byte & 0x7F) as i64) << (offset * 7);
            offset += 1;
            if byte & 0x80 == 0 {
                self.data_offset += offset;
                return Ok(v);
            }
        }
        Err(general_err!("too many bytes for VLQ"))
    }
}

impl PackedDecoder {
    fn new() -> Self {
        Self {
            data: Bytes::from(vec![]),
            data_offset: 0,
            rle_left: 0,
            rle_value: false,
            packed_count: 0,
            packed_offset: 0,
        }
    }

    fn set_data(&mut self, encoding: Encoding, data: Bytes) {
        self.rle_left = 0;
        self.rle_value = false;
        self.packed_offset = 0;
        self.packed_count = match encoding {
            Encoding::RLE => 0,
            #[allow(deprecated)]
            Encoding::BIT_PACKED => data.len() * 8,
            _ => unreachable!("invalid level encoding: {}", encoding),
        };
        self.data = data;
        self.data_offset = 0;
    }

    /// Try to consume `len` levels if all are valid (max definition level).
    ///
    /// Returns `Ok(Some(count))` if successfully consumed `count` all-valid levels.
    /// Returns `Ok(None)` if there are any nulls or packed data that prevents fast path.
    ///
    /// Note: On `None`, the decoder state may have advanced to the next RLE block,
    /// but only if `rle_left` was zero (i.e., the block would have been loaded
    /// on the next read anyway).
    fn try_consume_all_valid(&mut self, len: usize) -> Result<Option<usize>> {
        // If no active run and no packed data pending, try to parse the next RLE block
        if self.rle_left == 0 && self.packed_count == self.packed_offset {
            if self.data_offset < self.data.len() {
                self.next_rle_block()?;
            } else {
                // No more data available
                return Ok(None);
            }
        }

        // Fast path only works when we have an active RLE run of true values
        // that covers the entire requested length.
        if self.rle_left >= len && self.rle_value {
            self.rle_left -= len;
            return Ok(Some(len));
        }

        // Any other case (null run, packed data, or insufficient length)
        // falls back to normal path
        Ok(None)
    }

    fn read(&mut self, buffer: &mut BooleanBufferBuilder, len: usize) -> Result<usize> {
        let mut read = 0;
        while read != len {
            if self.rle_left != 0 {
                let to_read = self.rle_left.min(len - read);
                buffer.append_n(to_read, self.rle_value);
                self.rle_left -= to_read;
                read += to_read;
            } else if self.packed_count != self.packed_offset {
                let to_read = (self.packed_count - self.packed_offset).min(len - read);
                let offset = self.data_offset * 8 + self.packed_offset;
                buffer.append_packed_range(offset..offset + to_read, self.data.as_ref());
                self.packed_offset += to_read;
                read += to_read;

                if self.packed_offset == self.packed_count {
                    self.data_offset += self.packed_count / 8;
                }
            } else if self.data_offset == self.data.len() {
                break;
            } else {
                self.next_rle_block()?
            }
        }
        Ok(read)
    }

    /// Skips `level_num` definition levels
    ///
    /// Returns the number of values skipped and the number of levels skipped
    fn skip(&mut self, level_num: usize) -> Result<(usize, usize)> {
        let mut skipped_value = 0;
        let mut skipped_level = 0;
        while skipped_level != level_num {
            if self.rle_left != 0 {
                let to_skip = self.rle_left.min(level_num - skipped_level);
                self.rle_left -= to_skip;
                skipped_level += to_skip;
                if self.rle_value {
                    skipped_value += to_skip;
                }
            } else if self.packed_count != self.packed_offset {
                let to_skip =
                    (self.packed_count - self.packed_offset).min(level_num - skipped_level);
                let offset = self.data_offset * 8 + self.packed_offset;
                let bit_chunk = UnalignedBitChunk::new(self.data.as_ref(), offset, to_skip);
                skipped_value += bit_chunk.count_ones();
                self.packed_offset += to_skip;
                skipped_level += to_skip;
                if self.packed_offset == self.packed_count {
                    self.data_offset += self.packed_count / 8;
                }
            } else if self.data_offset == self.data.len() {
                break;
            } else {
                self.next_rle_block()?
            }
        }
        Ok((skipped_value, skipped_level))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

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

    #[test]
    fn test_build_validity_bitmap_unfiltered_word_chunk() {
        // 65 levels forces the unfiltered path to process one full u64 word
        // with append_word, plus a remainder bit.
        let def_levels = (0..65)
            .map(|i| if i % 3 == 0 { 2 } else { 1 })
            .collect::<Vec<_>>();
        let mut bitmap = BooleanBufferBuilder::new(0);

        assert_eq!(
            build_filtered_validity_bitmap(&def_levels, None, None, 2, &mut bitmap),
            def_levels.len()
        );

        let bitmap = bitmap.finish();
        for (idx, def) in def_levels.iter().enumerate() {
            assert_eq!(bitmap.value(idx), *def >= 2);
        }
    }

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

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

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

        // 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.random_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_packed_decoder_skip() {
        let mut rng = rng();
        let len: usize = rng.random_range(512..1024);

        let mut expected = BooleanBufferBuilder::new(len);
        let mut encoder = RleEncoder::new(1, 1024);

        let mut total_value = 0;
        for _ in 0..len {
            let bool = rng.random_bool(0.8);
            encoder.put(bool as u64);
            expected.append(bool);
            if bool {
                total_value += 1;
            }
        }
        assert_eq!(expected.len(), len);

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

        let mut skip_value = 0;
        let mut read_value = 0;
        let mut skip_level = 0;
        let mut read_level = 0;

        loop {
            let offset = skip_level + read_level;
            let remaining_levels = len - offset;
            if remaining_levels == 0 {
                break;
            }
            let to_read_or_skip_level = rng.random_range(1..=remaining_levels);
            if rng.random_bool(0.5) {
                let (skip_val_num, skip_level_num) = decoder.skip(to_read_or_skip_level).unwrap();
                skip_value += skip_val_num;
                skip_level += skip_level_num
            } else {
                let mut decoded = BooleanBufferBuilder::new(to_read_or_skip_level);
                let read_level_num = decoder.read(&mut decoded, to_read_or_skip_level).unwrap();
                read_level += read_level_num;
                for i in 0..read_level_num {
                    assert!(!decoded.is_empty());
                    //check each read bit
                    let read_bit = decoded.get_bit(i);
                    if read_bit {
                        read_value += 1;
                    }
                    let expect_bit = expected.get_bit(i + offset);
                    assert_eq!(read_bit, expect_bit);
                }
            }
        }
        assert_eq!(read_level + skip_level, len);
        assert_eq!(read_value + skip_value, total_value);
    }

    #[test]
    fn test_try_consume_all_valid() {
        // Test with all-valid data (all 1s) - single RLE run
        let len = 100;
        let mut encoder = RleEncoder::new(1, 1024);
        for _ in 0..len {
            encoder.put(1); // all valid
        }
        let encoded = encoder.consume();
        let mut decoder = PackedDecoder::new();
        decoder.set_data(Encoding::RLE, encoded.into());

        // try_consume_all_valid now parses the RLE block itself, no need to read first
        let result = decoder.try_consume_all_valid(len).unwrap();
        assert_eq!(result, Some(len));

        // Test with all-null data (all 0s)
        let mut encoder = RleEncoder::new(1, 1024);
        for _ in 0..len {
            encoder.put(0); // all null
        }
        let encoded = encoder.consume();
        let mut decoder = PackedDecoder::new();
        decoder.set_data(Encoding::RLE, encoded.into());

        // Should return None because rle_value is false (all nulls)
        let result = decoder.try_consume_all_valid(len).unwrap();
        assert_eq!(result, None);

        // Test when requesting more than available in current RLE run
        let mut encoder = RleEncoder::new(1, 1024);
        for _ in 0..10 {
            encoder.put(1); // small run of valid
        }
        for _ in 0..10 {
            encoder.put(0); // followed by nulls
        }
        let encoded = encoder.consume();
        let mut decoder = PackedDecoder::new();
        decoder.set_data(Encoding::RLE, encoded.into());

        // Request more than the valid run - should return None
        // (because we don't look ahead to next block)
        let result = decoder.try_consume_all_valid(20).unwrap();
        assert_eq!(result, None);

        // Reset decoder and try requesting within the run
        decoder.set_data(Encoding::RLE, {
            let mut encoder = RleEncoder::new(1, 1024);
            for _ in 0..10 {
                encoder.put(1);
            }
            for _ in 0..10 {
                encoder.put(0);
            }
            encoder.consume().into()
        });

        let result = decoder.try_consume_all_valid(5).unwrap();
        assert_eq!(result, Some(5));

        // After skipping 5, we should have 5 left in the valid RLE run
        let result = decoder.try_consume_all_valid(5).unwrap();
        assert_eq!(result, Some(5));

        // Now the valid run is exhausted, next call should parse the null run and return None
        let result = decoder.try_consume_all_valid(5).unwrap();
        assert_eq!(result, None);
    }
}