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        let buffers = self.buffers.into();
61        match data_type {
62            ArrowType::Utf8View => {
63                // Safety: views were created correctly, and checked that the data is utf8 when building the buffer
64                unsafe { Arc::new(StringViewArray::new_unchecked(views, buffers, nulls)) }
65            }
66            ArrowType::BinaryView => {
67                // Safety: views were created correctly
68                unsafe { Arc::new(BinaryViewArray::new_unchecked(views, buffers, nulls)) }
69            }
70            _ => panic!("Unsupported data type: {data_type}"),
71        }
72    }
73}
74
75impl ValuesBuffer for ViewBuffer {
76    fn with_capacity(capacity: usize) -> Self {
77        Self::with_capacity(capacity)
78    }
79
80    fn reserve_exact(&mut self, additional: usize) {
81        self.views.reserve_exact(additional);
82    }
83
84    fn pad_nulls(
85        &mut self,
86        read_offset: usize,
87        values_read: usize,
88        levels_read: usize,
89        valid_mask: &[u8],
90    ) -> Result<()> {
91        self.views
92            .pad_nulls(read_offset, values_read, levels_read, valid_mask)
93    }
94}
95
96#[cfg(test)]
97mod tests {
98
99    use arrow::array::make_view;
100    use arrow_array::Array;
101
102    use super::*;
103
104    #[test]
105    fn test_view_buffer_empty() {
106        let buffer = ViewBuffer::with_capacity(0);
107        let array = buffer.into_array(None, &ArrowType::Utf8View);
108        let strings = array
109            .as_any()
110            .downcast_ref::<arrow::array::StringViewArray>()
111            .unwrap();
112        assert_eq!(strings.len(), 0);
113    }
114
115    #[test]
116    fn test_view_buffer_append_view() {
117        let mut buffer = ViewBuffer::with_capacity(0);
118        let data = b"0123456789long string to test string view";
119        let string_buffer = Buffer::from(data);
120        let block_id = buffer.append_block(string_buffer);
121
122        buffer.views.push(make_view(&data[0..1], block_id, 0));
123        buffer.views.push(make_view(&data[1..10], block_id, 1));
124        buffer.views.push(make_view(&data[10..41], block_id, 10));
125
126        let array = buffer.into_array(None, &ArrowType::Utf8View);
127        let string_array = array
128            .as_any()
129            .downcast_ref::<arrow::array::StringViewArray>()
130            .unwrap();
131        assert_eq!(
132            string_array.iter().collect::<Vec<_>>(),
133            vec![
134                Some("0"),
135                Some("123456789"),
136                Some("long string to test string view"),
137            ]
138        );
139    }
140
141    #[test]
142    fn test_view_buffer_pad_null() {
143        let mut buffer = ViewBuffer::with_capacity(0);
144        let data = b"0123456789long string to test string view";
145        let string_buffer = Buffer::from(data);
146        let block_id = buffer.append_block(string_buffer);
147
148        buffer.views.push(make_view(&data[0..1], block_id, 0));
149        buffer.views.push(make_view(&data[1..10], block_id, 1));
150        buffer.views.push(make_view(&data[10..41], block_id, 10));
151
152        let valid = [true, false, false, true, false, false, true];
153        let valid_mask = Buffer::from_iter(valid.iter().copied());
154
155        buffer
156            .pad_nulls(1, 2, valid.len() - 1, valid_mask.as_slice())
157            .unwrap();
158
159        let array = buffer.into_array(Some(valid_mask), &ArrowType::Utf8View);
160        let strings = array
161            .as_any()
162            .downcast_ref::<arrow::array::StringViewArray>()
163            .unwrap();
164
165        assert_eq!(
166            strings.iter().collect::<Vec<_>>(),
167            vec![
168                Some("0"),
169                None,
170                None,
171                Some("123456789"),
172                None,
173                None,
174                Some("long string to test string view"),
175            ]
176        );
177    }
178}
parquet/arrow/buffer/view_buffer.rs

parquet/arrow/buffer/
view_buffer.rs
