parquet/arrow/decoder/
dictionary_index.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 bytes::Bytes;
19
20use crate::encodings::rle::{MAX_RLE_DICTIONARY_BIT_WIDTH, RleDecoder};
21use crate::errors::{ParquetError, Result};
22
23/// Decoder for `Encoding::RLE_DICTIONARY` indices
24pub struct DictIndexDecoder {
25    /// Decoder for the dictionary offsets array
26    decoder: RleDecoder,
27
28    /// We want to decode the offsets in chunks so we will maintain an internal buffer of decoded
29    /// offsets
30    index_buf: Box<[i32; 1024]>,
31    /// Current length of `index_buf`
32    index_buf_len: usize,
33    /// Current offset into `index_buf`. If `index_buf_offset` == `index_buf_len` then we've consumed
34    /// the entire buffer and need to decode another chunk of offsets.
35    index_offset: usize,
36
37    /// This is a maximum as the null count is not always known, e.g. value data from
38    /// a v1 data page
39    max_remaining_values: usize,
40}
41
42impl DictIndexDecoder {
43    /// Create a new [`DictIndexDecoder`] with the provided data page, the number of levels
44    /// associated with this data page, and the number of non-null values (if known)
45    pub fn new(data: Bytes, num_levels: usize, num_values: Option<usize>) -> Result<Self> {
46        let bit_width = *data
47            .first()
48            .ok_or_else(|| general_err!("dictionary index page is empty"))?;
49        if bit_width > MAX_RLE_DICTIONARY_BIT_WIDTH {
50            return Err(general_err!(
51                "Invalid or corrupted RLE bit width {bit_width}. Max allowed is {MAX_RLE_DICTIONARY_BIT_WIDTH}"
52            ));
53        }
54        let mut decoder = RleDecoder::new(bit_width);
55        decoder.set_data(data.slice(1..))?;
56
57        Ok(Self {
58            decoder,
59            index_buf: Box::new([0; 1024]),
60            index_buf_len: 0,
61            index_offset: 0,
62            max_remaining_values: num_values.unwrap_or(num_levels),
63        })
64    }
65
66    /// Read up to `len` values, returning the number of values read
67    /// and calling `f` with each decoded dictionary index
68    ///
69    /// Will short-circuit and return on error
70    #[inline(always)]
71    pub fn read<F: FnMut(&[i32]) -> Result<()>>(&mut self, len: usize, mut f: F) -> Result<usize> {
72        let total_to_read = len.min(self.max_remaining_values);
73
74        let mut values_read = 0;
75
76        let index_buf = self.index_buf.as_mut();
77        while values_read < total_to_read {
78            if self.index_offset == self.index_buf_len {
79                // We've consumed the entire index buffer so we need to reload it before proceeding
80                let read = self.decoder.get_batch(index_buf)?;
81                if read == 0 {
82                    break;
83                }
84                self.index_buf_len = read;
85                self.index_offset = 0;
86            }
87
88            let available = self.index_buf_len - self.index_offset;
89            let n = available.min(total_to_read - values_read);
90
91            f(&index_buf[self.index_offset..self.index_offset + n])?;
92
93            self.index_offset += n;
94            values_read += n;
95        }
96        self.max_remaining_values -= values_read;
97
98        Ok(values_read)
99    }
100
101    /// Skip up to `to_skip` values, returning the number of values skipped
102    pub fn skip(&mut self, to_skip: usize) -> Result<usize> {
103        let to_skip = to_skip.min(self.max_remaining_values);
104
105        let mut values_skip = 0;
106        while values_skip < to_skip {
107            if self.index_offset == self.index_buf_len {
108                // Instead of reloading the buffer, just skip in the decoder
109                let skip = self.decoder.skip(to_skip - values_skip)?;
110
111                if skip == 0 {
112                    break;
113                }
114
115                self.max_remaining_values -= skip;
116                values_skip += skip;
117            } else {
118                // We still have indices buffered, so skip within the buffer
119                let skip = (to_skip - values_skip).min(self.index_buf_len - self.index_offset);
120
121                self.index_offset += skip;
122                self.max_remaining_values -= skip;
123                values_skip += skip;
124            }
125        }
126        Ok(values_skip)
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn test_dict_index_decoder_empty_data() {
136        let Err(err) = DictIndexDecoder::new(Bytes::new(), 0, None) else {
137            panic!("expected error");
138        };
139        assert_eq!(
140            err.to_string(),
141            "Parquet error: dictionary index page is empty"
142        );
143    }
144
145    #[test]
146    fn test_dict_index_decoder_invalid_bit_width() {
147        let data = Bytes::from_static(&[33, 0, 0, 0]);
148        let Err(err) = DictIndexDecoder::new(data, 1, None) else {
149            panic!("expected error");
150        };
151        assert_eq!(
152            err.to_string(),
153            "Parquet error: Invalid or corrupted RLE bit width 33. Max allowed is 32"
154        );
155    }
156}
parquet/arrow/decoder/dictionary_index.rs

parquet/arrow/decoder/
dictionary_index.rs
