parquet/arrow/array_reader/
mod.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
18//! Logic for reading into arrow arrays
19
20use crate::errors::Result;
21use arrow_array::ArrayRef;
22use arrow_schema::DataType as ArrowType;
23use std::any::Any;
24use std::sync::Arc;
25
26use crate::arrow::record_reader::buffer::ValuesBuffer;
27use crate::arrow::record_reader::GenericRecordReader;
28use crate::column::page::PageIterator;
29use crate::column::reader::decoder::ColumnValueDecoder;
30use crate::file::reader::{FilePageIterator, FileReader};
31
32mod builder;
33mod byte_array;
34mod byte_array_dictionary;
35mod byte_view_array;
36mod empty_array;
37mod fixed_len_byte_array;
38mod fixed_size_list_array;
39mod list_array;
40mod map_array;
41mod null_array;
42mod primitive_array;
43mod struct_array;
44
45#[cfg(test)]
46mod test_util;
47
48pub(crate) use builder::ArrayReaderBuilder;
49pub use byte_array::make_byte_array_reader;
50pub use byte_array_dictionary::make_byte_array_dictionary_reader;
51#[allow(unused_imports)] // Only used for benchmarks
52pub use byte_view_array::make_byte_view_array_reader;
53#[allow(unused_imports)] // Only used for benchmarks
54pub use fixed_len_byte_array::make_fixed_len_byte_array_reader;
55pub use fixed_size_list_array::FixedSizeListArrayReader;
56pub use list_array::ListArrayReader;
57pub use map_array::MapArrayReader;
58pub use null_array::NullArrayReader;
59pub use primitive_array::PrimitiveArrayReader;
60pub use struct_array::StructArrayReader;
61
62/// Array reader reads parquet data into arrow array.
63pub trait ArrayReader: Send {
64    // TODO: this function is never used, and the trait is not public. Perhaps this should be
65    // removed.
66    #[allow(dead_code)]
67    fn as_any(&self) -> &dyn Any;
68
69    /// Returns the arrow type of this array reader.
70    fn get_data_type(&self) -> &ArrowType;
71
72    /// Reads at most `batch_size` records into an arrow array and return it.
73    #[cfg(any(feature = "experimental", test))]
74    fn next_batch(&mut self, batch_size: usize) -> Result<ArrayRef> {
75        self.read_records(batch_size)?;
76        self.consume_batch()
77    }
78
79    /// Reads at most `batch_size` records' bytes into buffer
80    ///
81    /// Returns the number of records read, which can be less than `batch_size` if
82    /// pages is exhausted.
83    fn read_records(&mut self, batch_size: usize) -> Result<usize>;
84
85    /// Consume all currently stored buffer data
86    /// into an arrow array and return it.
87    fn consume_batch(&mut self) -> Result<ArrayRef>;
88
89    /// Skips over `num_records` records, returning the number of rows skipped
90    fn skip_records(&mut self, num_records: usize) -> Result<usize>;
91
92    /// If this array has a non-zero definition level, i.e. has a nullable parent
93    /// array, returns the definition levels of data from the last call of `next_batch`
94    ///
95    /// Otherwise returns None
96    ///
97    /// This is used by parent [`ArrayReader`] to compute their null bitmaps
98    fn get_def_levels(&self) -> Option<&[i16]>;
99
100    /// If this array has a non-zero repetition level, i.e. has a repeated parent
101    /// array, returns the repetition levels of data from the last call of `next_batch`
102    ///
103    /// Otherwise returns None
104    ///
105    /// This is used by parent [`ArrayReader`] to compute their array offsets
106    fn get_rep_levels(&self) -> Option<&[i16]>;
107}
108
109/// A collection of row groups
110pub trait RowGroups {
111    /// Get the number of rows in this collection
112    fn num_rows(&self) -> usize;
113
114    /// Returns a [`PageIterator`] for all pages in the specified column chunk
115    /// across all row groups in this collection.
116    fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>>;
117}
118
119impl RowGroups for Arc<dyn FileReader> {
120    fn num_rows(&self) -> usize {
121        self.metadata().file_metadata().num_rows() as usize
122    }
123
124    fn column_chunks(&self, column_index: usize) -> Result<Box<dyn PageIterator>> {
125        let iterator = FilePageIterator::new(column_index, Arc::clone(self))?;
126        Ok(Box::new(iterator))
127    }
128}
129
130/// Uses `record_reader` to read up to `batch_size` records from `pages`
131///
132/// Returns the number of records read, which can be less than `batch_size` if
133/// pages is exhausted.
134fn read_records<V, CV>(
135    record_reader: &mut GenericRecordReader<V, CV>,
136    pages: &mut dyn PageIterator,
137    batch_size: usize,
138) -> Result<usize>
139where
140    V: ValuesBuffer,
141    CV: ColumnValueDecoder<Buffer = V>,
142{
143    let mut records_read = 0usize;
144    while records_read < batch_size {
145        let records_to_read = batch_size - records_read;
146
147        let records_read_once = record_reader.read_records(records_to_read)?;
148        records_read += records_read_once;
149
150        // Record reader exhausted
151        if records_read_once < records_to_read {
152            if let Some(page_reader) = pages.next() {
153                // Read from new page reader (i.e. column chunk)
154                record_reader.set_page_reader(page_reader?)?;
155            } else {
156                // Page reader also exhausted
157                break;
158            }
159        }
160    }
161    Ok(records_read)
162}
163
164/// Uses `record_reader` to skip up to `batch_size` records from `pages`
165///
166/// Returns the number of records skipped, which can be less than `batch_size` if
167/// pages is exhausted
168fn skip_records<V, CV>(
169    record_reader: &mut GenericRecordReader<V, CV>,
170    pages: &mut dyn PageIterator,
171    batch_size: usize,
172) -> Result<usize>
173where
174    V: ValuesBuffer,
175    CV: ColumnValueDecoder<Buffer = V>,
176{
177    let mut records_skipped = 0usize;
178    while records_skipped < batch_size {
179        let records_to_read = batch_size - records_skipped;
180
181        let records_skipped_once = record_reader.skip_records(records_to_read)?;
182        records_skipped += records_skipped_once;
183
184        // Record reader exhausted
185        if records_skipped_once < records_to_read {
186            if let Some(page_reader) = pages.next() {
187                // Read from new page reader (i.e. column chunk)
188                record_reader.set_page_reader(page_reader?)?;
189            } else {
190                // Page reader also exhausted
191                break;
192            }
193        }
194    }
195    Ok(records_skipped)
196}
parquet/arrow/array_reader/mod.rs

parquet/arrow/array_reader/
mod.rs
