parquet/arrow/buffer/
offset_buffer.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 crate::arrow::buffer::bit_util::iter_set_bits_rev;
19use crate::arrow::record_reader::buffer::ValuesBuffer;
20use crate::errors::{ParquetError, Result};
21use crate::util::utf8::check_valid_utf8;
22use arrow_array::{ArrayRef, OffsetSizeTrait, make_array};
23use arrow_buffer::{ArrowNativeType, Buffer};
24use arrow_data::ArrayDataBuilder;
25use arrow_schema::DataType as ArrowType;
26
27/// A buffer of variable-sized byte arrays that can be converted into
28/// a corresponding [`ArrayRef`]
29#[derive(Debug)]
30pub struct OffsetBuffer<I: OffsetSizeTrait> {
31    pub offsets: Vec<I>,
32    pub values: Vec<u8>,
33}
34
35impl<I: OffsetSizeTrait> OffsetBuffer<I> {
36    /// Create a new `OffsetBuffer` with capacity for at least `capacity` elements
37    ///
38    /// Pre-allocates the offsets vector to avoid reallocations during reading.
39    /// The values vector is not pre-allocated as its size is unpredictable.
40    pub fn with_capacity(capacity: usize) -> Self {
41        let mut offsets = Vec::with_capacity(capacity + 1);
42        offsets.push(I::default());
43        Self {
44            offsets,
45            values: Vec::new(),
46        }
47    }
48
49    /// Returns the number of byte arrays in this buffer
50    pub fn len(&self) -> usize {
51        self.offsets.len() - 1
52    }
53
54    pub fn is_empty(&self) -> bool {
55        self.len() == 0
56    }
57
58    /// If `validate_utf8` this verifies that the first character of `data` is
59    /// the start of a UTF-8 codepoint
60    ///
61    /// Note: This does not verify that the entirety of `data` is valid
62    /// UTF-8. This should be done by calling [`Self::check_valid_utf8`] after
63    /// all data has been written
64    pub fn try_push(&mut self, data: &[u8], validate_utf8: bool) -> Result<()> {
65        if validate_utf8 && let Some(&b) = data.first() {
66            // A valid code-point iff it does not start with 0b10xxxxxx
67            // Bit-magic taken from `std::str::is_char_boundary`
68            if (b as i8) < -0x40 {
69                return Err(ParquetError::General(
70                    "encountered non UTF-8 data".to_string(),
71                ));
72            }
73        }
74
75        self.values.extend_from_slice(data);
76
77        let index_offset = I::from_usize(self.values.len())
78            .ok_or_else(|| general_err!("index overflow decoding byte array"))?;
79
80        self.offsets.push(index_offset);
81        Ok(())
82    }
83
84    /// Extends this buffer with a list of keys
85    ///
86    /// For each value `key` in `keys` this will insert
87    /// `&dict_values[dict_offsets[key]..dict_offsets[key+1]]`
88    ///
89    /// Note: This will validate offsets are valid
90    pub fn extend_from_dictionary<K: ArrowNativeType, V: ArrowNativeType>(
91        &mut self,
92        keys: &[K],
93        dict_offsets: &[V],
94        dict_values: &[u8],
95    ) -> Result<()> {
96        self.offsets.reserve(keys.len());
97
98        for key in keys {
99            let index = key.as_usize();
100            if index + 1 >= dict_offsets.len() {
101                return Err(general_err!(
102                    "dictionary key beyond bounds of dictionary: 0..{}",
103                    dict_offsets.len().saturating_sub(1)
104                ));
105            }
106            let start_offset = dict_offsets[index].as_usize();
107            let end_offset = dict_offsets[index + 1].as_usize();
108
109            // Dictionary values are verified when decoding dictionary page
110            self.values
111                .extend_from_slice(&dict_values[start_offset..end_offset]);
112            let index_offset = I::from_usize(self.values.len())
113                .ok_or_else(|| general_err!("index overflow decoding byte array"))?;
114            self.offsets.push(index_offset);
115        }
116        Ok(())
117    }
118
119    /// Validates that `&self.values[start_offset..]` is a valid UTF-8 sequence
120    ///
121    /// This MUST be combined with validating that the offsets start on a character
122    /// boundary, otherwise it would be possible for the values array to be a valid UTF-8
123    /// sequence, but not the individual string slices it contains
124    ///
125    /// [`Self::try_push`] can perform this validation check on insertion
126    pub fn check_valid_utf8(&self, start_offset: usize) -> Result<()> {
127        check_valid_utf8(&self.values.as_slice()[start_offset..])
128    }
129
130    /// Converts this into an [`ArrayRef`] with the provided `data_type` and `null_buffer`
131    pub fn into_array(self, null_buffer: Option<Buffer>, data_type: ArrowType) -> ArrayRef {
132        let array_data_builder = ArrayDataBuilder::new(data_type)
133            .len(self.len())
134            .add_buffer(Buffer::from_vec(self.offsets))
135            .add_buffer(Buffer::from_vec(self.values))
136            .null_bit_buffer(null_buffer);
137
138        let data = match cfg!(debug_assertions) {
139            true => array_data_builder.build().unwrap(),
140            false => unsafe { array_data_builder.build_unchecked() },
141        };
142
143        make_array(data)
144    }
145}
146
147impl<I: OffsetSizeTrait> ValuesBuffer for OffsetBuffer<I> {
148    fn with_capacity(capacity: usize) -> Self {
149        Self::with_capacity(capacity)
150    }
151
152    fn reserve_exact(&mut self, additional: usize) {
153        self.offsets.reserve_exact(additional);
154    }
155
156    fn pad_nulls(
157        &mut self,
158        read_offset: usize,
159        values_read: usize,
160        levels_read: usize,
161        valid_mask: &[u8],
162    ) -> Result<()> {
163        if self.offsets.len() != read_offset + values_read + 1 {
164            return Err(general_err!(
165                "found inconsistent offsets while padding nulls: expected {} offsets, got {}",
166                read_offset + values_read + 1,
167                self.offsets.len()
168            ));
169        }
170        self.offsets
171            .resize(read_offset + levels_read + 1, I::default());
172
173        let offsets = &mut self.offsets;
174
175        let mut last_pos = read_offset + levels_read + 1;
176        let mut last_start_offset = I::from_usize(self.values.len()).unwrap();
177
178        let values_range = read_offset..read_offset + values_read;
179        for (value_pos, level_pos) in values_range
180            .clone()
181            .rev()
182            .zip(iter_set_bits_rev(valid_mask))
183        {
184            if level_pos < value_pos || level_pos >= last_pos {
185                return Err(general_err!("found corrupt level data while padding nulls"));
186            }
187
188            let end_offset = offsets[value_pos + 1];
189            let start_offset = offsets[value_pos];
190
191            // Fill in any nulls
192            for x in &mut offsets[level_pos + 1..last_pos] {
193                *x = end_offset;
194            }
195
196            if level_pos == value_pos {
197                return Ok(());
198            }
199
200            offsets[level_pos] = start_offset;
201            last_pos = level_pos;
202            last_start_offset = start_offset;
203        }
204
205        // Pad leading nulls up to `last_offset`
206        for x in &mut offsets[values_range.start + 1..last_pos] {
207            *x = last_start_offset
208        }
209        Ok(())
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use arrow_array::{Array, LargeStringArray, StringArray};
217
218    #[test]
219    fn test_offset_buffer_empty() {
220        let buffer = OffsetBuffer::<i32>::with_capacity(0);
221        let array = buffer.into_array(None, ArrowType::Utf8);
222        let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
223        assert_eq!(strings.len(), 0);
224    }
225
226    #[test]
227    fn test_offset_buffer_append() {
228        let mut buffer = OffsetBuffer::<i64>::with_capacity(0);
229        buffer.try_push("hello".as_bytes(), true).unwrap();
230        buffer.try_push("bar".as_bytes(), true).unwrap();
231        buffer
232            .extend_from_dictionary(&[1, 3, 0, 2], &[0, 2, 4, 5, 6], "abcdef".as_bytes())
233            .unwrap();
234
235        let array = buffer.into_array(None, ArrowType::LargeUtf8);
236        let strings = array.as_any().downcast_ref::<LargeStringArray>().unwrap();
237        assert_eq!(
238            strings.iter().map(|x| x.unwrap()).collect::<Vec<_>>(),
239            vec!["hello", "bar", "cd", "f", "ab", "e"]
240        )
241    }
242
243    #[test]
244    fn test_offset_buffer() {
245        let mut buffer = OffsetBuffer::<i32>::with_capacity(0);
246        for v in ["hello", "world", "cupcakes", "a", "b", "c"] {
247            buffer.try_push(v.as_bytes(), false).unwrap()
248        }
249        let split = std::mem::replace(&mut buffer, OffsetBuffer::with_capacity(0));
250
251        let array = split.into_array(None, ArrowType::Utf8);
252        let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
253        assert_eq!(
254            strings.iter().map(|x| x.unwrap()).collect::<Vec<_>>(),
255            vec!["hello", "world", "cupcakes", "a", "b", "c"]
256        );
257
258        buffer.try_push("test".as_bytes(), false).unwrap();
259        let array = buffer.into_array(None, ArrowType::Utf8);
260        let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
261        assert_eq!(
262            strings.iter().map(|x| x.unwrap()).collect::<Vec<_>>(),
263            vec!["test"]
264        );
265    }
266
267    #[test]
268    fn test_offset_buffer_pad_nulls() {
269        let mut buffer = OffsetBuffer::<i32>::with_capacity(0);
270        let values = ["a", "b", "c", "def", "gh"];
271        for v in &values {
272            buffer.try_push(v.as_bytes(), false).unwrap()
273        }
274
275        let valid = [
276            true, false, false, true, false, true, false, true, true, false, false,
277        ];
278        let valid_mask = Buffer::from_iter(valid.iter().copied());
279
280        // Both trailing and leading nulls
281        buffer
282            .pad_nulls(1, values.len() - 1, valid.len() - 1, valid_mask.as_slice())
283            .unwrap();
284
285        let array = buffer.into_array(Some(valid_mask), ArrowType::Utf8);
286        let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
287        assert_eq!(
288            strings.iter().collect::<Vec<_>>(),
289            vec![
290                Some("a"),
291                None,
292                None,
293                Some("b"),
294                None,
295                Some("c"),
296                None,
297                Some("def"),
298                Some("gh"),
299                None,
300                None
301            ]
302        );
303    }
304
305    #[test]
306    fn test_utf8_validation() {
307        let valid_2_byte_utf8 = &[0b11001000, 0b10001000];
308        std::str::from_utf8(valid_2_byte_utf8).unwrap();
309        let valid_3_byte_utf8 = &[0b11101000, 0b10001000, 0b10001000];
310        std::str::from_utf8(valid_3_byte_utf8).unwrap();
311        let valid_4_byte_utf8 = &[0b11110010, 0b10101000, 0b10101001, 0b10100101];
312        std::str::from_utf8(valid_4_byte_utf8).unwrap();
313
314        let mut buffer = OffsetBuffer::<i32>::with_capacity(0);
315        buffer.try_push(valid_2_byte_utf8, true).unwrap();
316        buffer.try_push(valid_3_byte_utf8, true).unwrap();
317        buffer.try_push(valid_4_byte_utf8, true).unwrap();
318
319        // Cannot append string starting with incomplete codepoint
320        buffer.try_push(&valid_2_byte_utf8[1..], true).unwrap_err();
321        buffer.try_push(&valid_3_byte_utf8[1..], true).unwrap_err();
322        buffer.try_push(&valid_3_byte_utf8[2..], true).unwrap_err();
323        buffer.try_push(&valid_4_byte_utf8[1..], true).unwrap_err();
324        buffer.try_push(&valid_4_byte_utf8[2..], true).unwrap_err();
325        buffer.try_push(&valid_4_byte_utf8[3..], true).unwrap_err();
326
327        // Can append data containing an incomplete codepoint
328        buffer.try_push(&[0b01111111, 0b10111111], true).unwrap();
329
330        assert_eq!(buffer.len(), 4);
331        assert_eq!(buffer.values.len(), 11);
332
333        buffer.try_push(valid_3_byte_utf8, true).unwrap();
334
335        // Should fail due to incomplete codepoint
336        buffer.check_valid_utf8(0).unwrap_err();
337
338        // After broken codepoint -> success
339        buffer.check_valid_utf8(11).unwrap();
340
341        // Fails if run from middle of codepoint
342        buffer.check_valid_utf8(12).unwrap_err();
343    }
344
345    #[test]
346    fn test_pad_nulls_corrupt_input_returns_err() {
347        // Corrupt input must produce a decode error rather than panicking.
348
349        // Offsets inconsistent with `values_read`: only one value was pushed,
350        // but three are claimed to have been read.
351        let mut buffer = OffsetBuffer::<i32>::with_capacity(0);
352        buffer.try_push("a".as_bytes(), false).unwrap();
353        let valid_mask = Buffer::from_iter([true, false, false]);
354        let err = buffer
355            .pad_nulls(0, 3, 3, valid_mask.as_slice())
356            .unwrap_err();
357        assert!(
358            err.to_string().contains("inconsistent offsets"),
359            "unexpected error: {err}"
360        );
361
362        // Valid mask has fewer set bits than `values_read`, which previously
363        // tripped an assertion in the null-padding loop.
364        let mut buffer = OffsetBuffer::<i32>::with_capacity(0);
365        for v in ["a", "b", "c"] {
366            buffer.try_push(v.as_bytes(), false).unwrap();
367        }
368        let valid_mask = Buffer::from_iter([true, false, false]);
369        let err = buffer
370            .pad_nulls(0, 3, 3, valid_mask.as_slice())
371            .unwrap_err();
372        assert!(
373            err.to_string().contains("corrupt level data"),
374            "unexpected error: {err}"
375        );
376    }
377
378    #[test]
379    fn test_pad_nulls_empty() {
380        let mut buffer = OffsetBuffer::<i32>::with_capacity(0);
381        let valid_mask = Buffer::from_iter(std::iter::repeat_n(false, 9));
382        buffer.pad_nulls(0, 0, 9, valid_mask.as_slice()).unwrap();
383
384        let array = buffer.into_array(Some(valid_mask), ArrowType::Utf8);
385        let strings = array.as_any().downcast_ref::<StringArray>().unwrap();
386
387        assert_eq!(strings.len(), 9);
388        assert!(strings.iter().all(|x| x.is_none()))
389    }
390}
parquet/arrow/buffer/offset_buffer.rs

parquet/arrow/buffer/
offset_buffer.rs
