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

parquet/arrow/buffer/
offset_buffer.rs
