parquet/arrow/record_reader/
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::errors::Result;
20
21/// A buffer that supports padding with nulls
22pub trait ValuesBuffer {
23    /// Create a new buffer with capacity for at least `capacity` logical values
24    fn with_capacity(capacity: usize) -> Self;
25
26    /// Reserve capacity for `additional` more logical values.
27    ///
28    /// Callers should use this once the level decoders have determined how
29    /// many output slots are actually needed for the next read.
30    fn reserve_exact(&mut self, additional: usize);
31
32    /// If a column contains nulls, more level data may be read than value data, as null
33    /// values are not encoded. Therefore, first the levels data is read, the null count
34    /// determined, and then the corresponding number of values read to a [`ValuesBuffer`].
35    ///
36    /// It is then necessary to move this values data into positions that correspond to
37    /// the non-null level positions. This is what this method does.
38    ///
39    /// It is provided with:
40    ///
41    /// - `read_offset` - the offset in [`ValuesBuffer`] to start null padding from
42    /// - `values_read` - the number of values read
43    /// - `levels_read` - the number of levels read
44    /// - `valid_mask` - a packed mask of valid levels
45    ///
46    /// Returns an error if the inputs are inconsistent, for example because the
47    /// decoded data was corrupt. This must not panic on such input.
48    fn pad_nulls(
49        &mut self,
50        read_offset: usize,
51        values_read: usize,
52        levels_read: usize,
53        valid_mask: &[u8],
54    ) -> Result<()>;
55}
56
57impl<T: Copy + Default> ValuesBuffer for Vec<T> {
58    fn with_capacity(capacity: usize) -> Self {
59        Vec::with_capacity(capacity)
60    }
61
62    fn reserve_exact(&mut self, additional: usize) {
63        Vec::reserve_exact(self, additional)
64    }
65
66    fn pad_nulls(
67        &mut self,
68        read_offset: usize,
69        values_read: usize,
70        levels_read: usize,
71        valid_mask: &[u8],
72    ) -> Result<()> {
73        self.resize(read_offset + levels_read, T::default());
74
75        let values_range = read_offset..read_offset + values_read;
76        for (value_pos, level_pos) in values_range.rev().zip(iter_set_bits_rev(valid_mask)) {
77            debug_assert!(level_pos >= value_pos);
78            if level_pos <= value_pos {
79                break;
80            }
81            self[level_pos] = self[value_pos];
82        }
83        Ok(())
84    }
85}
parquet/arrow/record_reader/buffer.rs

parquet/arrow/record_reader/
buffer.rs
