parquet/arrow/buffer/
view_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::record_reader::buffer::ValuesBuffer;
19use crate::errors::Result;
20use arrow_array::{ArrayRef, BinaryViewArray, StringViewArray};
21use arrow_buffer::{Buffer, NullBuffer, ScalarBuffer};
22use arrow_schema::DataType as ArrowType;
23use std::sync::Arc;
24
25/// A buffer of view type byte arrays that can be converted into
26/// `GenericByteViewArray`
27///
28/// Note this does not reuse `GenericByteViewBuilder` due to the need to call `pad_nulls`
29/// and reuse the existing logic for Vec in the parquet crate
30#[derive(Debug, Default)]
31pub struct ViewBuffer {
32    pub views: Vec<u128>,
33    pub buffers: Vec<Buffer>,
34}
35
36impl ViewBuffer {
37    /// Create a new ViewBuffer with capacity for the specified number of views
38    pub fn with_capacity(capacity: usize) -> Self {
39        Self {
40            views: Vec::with_capacity(capacity),
41            buffers: Vec::new(),
42        }
43    }
44
45    pub fn is_empty(&self) -> bool {
46        self.views.is_empty()
47    }
48
49    pub fn append_block(&mut self, block: Buffer) -> u32 {
50        let block_id = self.buffers.len() as u32;
51        self.buffers.push(block);
52        block_id
53    }
54
55    /// Converts this into an [`ArrayRef`] with the provided `data_type` and `null_buffer`
56    pub fn into_array(self, null_buffer: Option<Buffer>, data_type: &ArrowType) -> ArrayRef {
57        let len = self.views.len();
58        let views = ScalarBuffer::from(self.views);
59        let nulls = null_buffer.and_then(|b| NullBuffer::from_unsliced_buffer(b, len));
60        match data_type {
61            ArrowType::Utf8View => {
62                // Safety: views were created correctly, and checked that the data is utf8 when building the buffer
63                unsafe { Arc::new(StringViewArray::new_unchecked(views, self.buffers, nulls)) }
64            }
65            ArrowType::BinaryView => {
66                // Safety: views were created correctly
67                unsafe { Arc::new(BinaryViewArray::new_unchecked(views, self.buffers, nulls)) }
68            }
69            _ => panic!("Unsupported data type: {data_type}"),
70        }
71    }
72}
73
74impl ValuesBuffer for ViewBuffer {
75    fn with_capacity(capacity: usize) -> Self {
76        Self::with_capacity(capacity)
77    }
78
79    fn reserve_exact(&mut self, additional: usize) {
80        self.views.reserve_exact(additional);
81    }
82
83    fn pad_nulls(
84        &mut self,
85        read_offset: usize,
86        values_read: usize,
87        levels_read: usize,
88        valid_mask: &[u8],
89    ) -> Result<()> {
90        self.views
91            .pad_nulls(read_offset, values_read, levels_read, valid_mask)
92    }
93}
94
95#[cfg(test)]
96mod tests {
97
98    use arrow::array::make_view;
99    use arrow_array::Array;
100
101    use super::*;
102
103    #[test]
104    fn test_view_buffer_empty() {
105        let buffer = ViewBuffer::with_capacity(0);
106        let array = buffer.into_array(None, &ArrowType::Utf8View);
107        let strings = array
108            .as_any()
109            .downcast_ref::<arrow::array::StringViewArray>()
110            .unwrap();
111        assert_eq!(strings.len(), 0);
112    }
113
114    #[test]
115    fn test_view_buffer_append_view() {
116        let mut buffer = ViewBuffer::with_capacity(0);
117        let data = b"0123456789long string to test string view";
118        let string_buffer = Buffer::from(data);
119        let block_id = buffer.append_block(string_buffer);
120
121        buffer.views.push(make_view(&data[0..1], block_id, 0));
122        buffer.views.push(make_view(&data[1..10], block_id, 1));
123        buffer.views.push(make_view(&data[10..41], block_id, 10));
124
125        let array = buffer.into_array(None, &ArrowType::Utf8View);
126        let string_array = array
127            .as_any()
128            .downcast_ref::<arrow::array::StringViewArray>()
129            .unwrap();
130        assert_eq!(
131            string_array.iter().collect::<Vec<_>>(),
132            vec![
133                Some("0"),
134                Some("123456789"),
135                Some("long string to test string view"),
136            ]
137        );
138    }
139
140    #[test]
141    fn test_view_buffer_pad_null() {
142        let mut buffer = ViewBuffer::with_capacity(0);
143        let data = b"0123456789long string to test string view";
144        let string_buffer = Buffer::from(data);
145        let block_id = buffer.append_block(string_buffer);
146
147        buffer.views.push(make_view(&data[0..1], block_id, 0));
148        buffer.views.push(make_view(&data[1..10], block_id, 1));
149        buffer.views.push(make_view(&data[10..41], block_id, 10));
150
151        let valid = [true, false, false, true, false, false, true];
152        let valid_mask = Buffer::from_iter(valid.iter().copied());
153
154        buffer
155            .pad_nulls(1, 2, valid.len() - 1, valid_mask.as_slice())
156            .unwrap();
157
158        let array = buffer.into_array(Some(valid_mask), &ArrowType::Utf8View);
159        let strings = array
160            .as_any()
161            .downcast_ref::<arrow::array::StringViewArray>()
162            .unwrap();
163
164        assert_eq!(
165            strings.iter().collect::<Vec<_>>(),
166            vec![
167                Some("0"),
168                None,
169                None,
170                Some("123456789"),
171                None,
172                None,
173                Some("long string to test string view"),
174            ]
175        );
176    }
177}
parquet/arrow/buffer/view_buffer.rs

parquet/arrow/buffer/
view_buffer.rs
