parquet/encodings/
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 std::mem;
19
20use super::rle::RleEncoder;
21
22use crate::basic::Encoding;
23use crate::data_type::AsBytes;
24use crate::util::bit_util::{ceil, num_required_bits, BitWriter};
25
26/// Computes max buffer size for level encoder/decoder based on encoding, max
27/// repetition/definition level and number of total buffered values (includes null
28/// values).
29#[inline]
30pub fn max_buffer_size(encoding: Encoding, max_level: i16, num_buffered_values: usize) -> usize {
31    let bit_width = num_required_bits(max_level as u64);
32    match encoding {
33        Encoding::RLE => RleEncoder::max_buffer_size(bit_width, num_buffered_values),
34        #[allow(deprecated)]
35        Encoding::BIT_PACKED => ceil(num_buffered_values * bit_width as usize, 8),
36        _ => panic!("Unsupported encoding type {encoding}"),
37    }
38}
39
40/// Encoder for definition/repetition levels.
41/// Currently only supports Rle and BitPacked (dev/null) encoding, including v2.
42pub enum LevelEncoder {
43    Rle(RleEncoder),
44    RleV2(RleEncoder),
45    BitPacked(u8, BitWriter),
46}
47
48impl LevelEncoder {
49    /// Creates new level encoder based on encoding, max level and underlying byte buffer.
50    /// For bit packed encoding it is assumed that buffer is already allocated with
51    /// `levels::max_buffer_size` method.
52    ///
53    /// Used to encode levels for Data Page v1.
54    ///
55    /// Panics, if encoding is not supported.
56    pub fn v1(encoding: Encoding, max_level: i16, capacity: usize) -> Self {
57        let capacity_bytes = max_buffer_size(encoding, max_level, capacity);
58        let mut buffer = Vec::with_capacity(capacity_bytes);
59        let bit_width = num_required_bits(max_level as u64);
60        match encoding {
61            Encoding::RLE => {
62                // Reserve space for length header
63                buffer.extend_from_slice(&[0; 4]);
64                LevelEncoder::Rle(RleEncoder::new_from_buf(bit_width, buffer))
65            }
66            #[allow(deprecated)]
67            Encoding::BIT_PACKED => {
68                // Here we set full byte buffer without adjusting for num_buffered_values,
69                // because byte buffer will already be allocated with size from
70                // `max_buffer_size()` method.
71                LevelEncoder::BitPacked(bit_width, BitWriter::new_from_buf(buffer))
72            }
73            _ => panic!("Unsupported encoding type {encoding}"),
74        }
75    }
76
77    /// Creates new level encoder based on RLE encoding. Used to encode Data Page v2
78    /// repetition and definition levels.
79    pub fn v2(max_level: i16, capacity: usize) -> Self {
80        let capacity_bytes = max_buffer_size(Encoding::RLE, max_level, capacity);
81        let buffer = Vec::with_capacity(capacity_bytes);
82        let bit_width = num_required_bits(max_level as u64);
83        LevelEncoder::RleV2(RleEncoder::new_from_buf(bit_width, buffer))
84    }
85
86    /// Put/encode levels vector into this level encoder.
87    /// Returns number of encoded values that are less than or equal to length of the
88    /// input buffer.
89    #[inline]
90    pub fn put(&mut self, buffer: &[i16]) -> usize {
91        let mut num_encoded = 0;
92        match *self {
93            LevelEncoder::Rle(ref mut encoder) | LevelEncoder::RleV2(ref mut encoder) => {
94                for value in buffer {
95                    encoder.put(*value as u64);
96                    num_encoded += 1;
97                }
98                encoder.flush();
99            }
100            LevelEncoder::BitPacked(bit_width, ref mut encoder) => {
101                for value in buffer {
102                    encoder.put_value(*value as u64, bit_width as usize);
103                    num_encoded += 1;
104                }
105                encoder.flush();
106            }
107        }
108        num_encoded
109    }
110
111    /// Finalizes level encoder, flush all intermediate buffers and return resulting
112    /// encoded buffer. Returned buffer is already truncated to encoded bytes only.
113    #[inline]
114    pub fn consume(self) -> Vec<u8> {
115        match self {
116            LevelEncoder::Rle(encoder) => {
117                let mut encoded_data = encoder.consume();
118                // Account for the buffer offset
119                let encoded_len = encoded_data.len() - mem::size_of::<i32>();
120                let len = (encoded_len as i32).to_le();
121                let len_bytes = len.as_bytes();
122                encoded_data[0..len_bytes.len()].copy_from_slice(len_bytes);
123                encoded_data
124            }
125            LevelEncoder::RleV2(encoder) => encoder.consume(),
126            LevelEncoder::BitPacked(_, encoder) => encoder.consume(),
127        }
128    }
129}
parquet/encodings/levels.rs

parquet/encodings/
levels.rs
