parquet/arrow/
in_memory_row_group.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::ProjectionMask;
19use crate::arrow::array_reader::RowGroups;
20use crate::arrow::arrow_reader::RowSelection;
21use crate::column::page::{PageIterator, PageReader};
22use crate::errors::ParquetError;
23use crate::file::metadata::{ParquetMetaData, RowGroupMetaData};
24use crate::file::page_index::offset_index::OffsetIndexMetaData;
25use crate::file::reader::{ChunkReader, Length, SerializedPageReader};
26use bytes::{Buf, Bytes};
27use std::ops::Range;
28use std::sync::Arc;
29
30/// An in-memory collection of column chunks
31#[derive(Debug)]
32pub(crate) struct InMemoryRowGroup<'a> {
33    pub(crate) offset_index: Option<&'a [Option<OffsetIndexMetaData>]>,
34    /// Column chunks for this row group
35    pub(crate) column_chunks: Vec<Option<Arc<ColumnChunkData>>>,
36    pub(crate) row_count: usize,
37    pub(crate) row_group_idx: usize,
38    pub(crate) metadata: &'a ParquetMetaData,
39}
40
41/// What ranges to fetch for the columns in this row group
42#[derive(Debug)]
43pub(crate) struct FetchRanges {
44    /// The byte ranges to fetch
45    pub(crate) ranges: Vec<Range<u64>>,
46    /// If `Some`, the start offsets of each page for each column chunk
47    pub(crate) page_start_offsets: Option<Vec<Vec<u64>>>,
48}
49
50impl InMemoryRowGroup<'_> {
51    /// Returns the byte ranges to fetch for the columns specified in
52    /// `projection` and `selection`.
53    ///
54    /// `cache_mask` indicates which columns, if any, are being cached by
55    /// [`RowGroupCache`](crate::arrow::array_reader::RowGroupCache).
56    /// The `selection` for Cached columns is expanded to batch boundaries to simplify
57    /// accounting for what data is cached.
58    pub(crate) fn fetch_ranges(
59        &self,
60        projection: &ProjectionMask,
61        selection: Option<&RowSelection>,
62        batch_size: usize,
63        cache_mask: Option<&ProjectionMask>,
64    ) -> FetchRanges {
65        let metadata = self.metadata.row_group(self.row_group_idx);
66        if let Some((selection, offset_index)) = selection.zip(self.offset_index) {
67            let expanded_selection =
68                selection.expand_to_batch_boundaries(batch_size, self.row_count);
69
70            // If we have a `RowSelection` and an `OffsetIndex` then only fetch
71            // pages required for the `RowSelection`
72            // Consider preallocating outer vec: https://github.com/apache/arrow-rs/issues/8667
73            let mut page_start_offsets: Vec<Vec<u64>> = vec![];
74
75            let ranges = self
76                .column_chunks
77                .iter()
78                .zip(metadata.columns())
79                .enumerate()
80                .filter(|&(idx, (chunk, _chunk_meta))| {
81                    chunk.is_none() && projection.leaf_included(idx)
82                })
83                .flat_map(|(idx, (_chunk, chunk_meta))| {
84                    // If the first page does not start at the beginning of the column,
85                    // then we need to also fetch a dictionary page.
86                    let mut ranges: Vec<Range<u64>> = vec![];
87                    let (start, len) = chunk_meta.byte_range();
88                    let Some(offset_idx) = offset_index[idx].as_ref() else {
89                        // No offset index for this column, fetch the entire column
90                        ranges.push(start..start + len);
91                        return ranges;
92                    };
93
94                    match offset_idx.page_locations.first() {
95                        Some(first) if first.offset as u64 != start => {
96                            ranges.push(start..first.offset as u64);
97                        }
98                        _ => (),
99                    }
100
101                    // Expand selection to batch boundaries if needed for caching
102                    // (see doc comment for this function for details on `cache_mask`)
103                    let use_expanded = cache_mask.map(|m| m.leaf_included(idx)).unwrap_or(false);
104                    if use_expanded {
105                        ranges.extend(expanded_selection.scan_ranges(&offset_idx.page_locations));
106                    } else {
107                        ranges.extend(selection.scan_ranges(&offset_idx.page_locations));
108                    }
109                    page_start_offsets.push(ranges.iter().map(|range| range.start).collect());
110
111                    ranges
112                })
113                .collect();
114            FetchRanges {
115                ranges,
116                page_start_offsets: Some(page_start_offsets),
117            }
118        } else {
119            let ranges = self
120                .column_chunks
121                .iter()
122                .enumerate()
123                .filter(|&(idx, chunk)| chunk.is_none() && projection.leaf_included(idx))
124                .map(|(idx, _chunk)| {
125                    let column = metadata.column(idx);
126                    let (start, length) = column.byte_range();
127                    start..(start + length)
128                })
129                .collect();
130            FetchRanges {
131                ranges,
132                page_start_offsets: None,
133            }
134        }
135    }
136
137    /// Fills in `self.column_chunks` with the data fetched from `chunk_data`.
138    ///
139    /// This function **must** be called with the data from the ranges returned by
140    /// `fetch_ranges` and the corresponding page_start_offsets, with the exact same and `selection`.
141    pub(crate) fn fill_column_chunks<I>(
142        &mut self,
143        projection: &ProjectionMask,
144        page_start_offsets: Option<Vec<Vec<u64>>>,
145        chunk_data: I,
146    ) where
147        I: IntoIterator<Item = Bytes>,
148    {
149        let mut chunk_data = chunk_data.into_iter();
150        let metadata = self.metadata.row_group(self.row_group_idx);
151        if let Some(page_start_offsets) = page_start_offsets {
152            // If we have a `RowSelection` and an `OffsetIndex` then only fetch pages required for the
153            // `RowSelection`
154            let mut page_start_offsets = page_start_offsets.into_iter();
155
156            for (idx, chunk) in self.column_chunks.iter_mut().enumerate() {
157                if chunk.is_some() || !projection.leaf_included(idx) {
158                    continue;
159                }
160
161                if let Some(offsets) = page_start_offsets.next() {
162                    let mut chunks = Vec::with_capacity(offsets.len());
163                    for _ in 0..offsets.len() {
164                        chunks.push(chunk_data.next().unwrap());
165                    }
166
167                    *chunk = Some(Arc::new(ColumnChunkData::Sparse {
168                        length: metadata.column(idx).byte_range().1 as usize,
169                        data: offsets
170                            .into_iter()
171                            .map(|x| x as usize)
172                            .zip(chunks)
173                            .collect(),
174                    }))
175                }
176            }
177        } else {
178            for (idx, chunk) in self.column_chunks.iter_mut().enumerate() {
179                if chunk.is_some() || !projection.leaf_included(idx) {
180                    continue;
181                }
182
183                if let Some(data) = chunk_data.next() {
184                    *chunk = Some(Arc::new(ColumnChunkData::Dense {
185                        offset: metadata.column(idx).byte_range().0 as usize,
186                        data,
187                    }));
188                }
189            }
190        }
191    }
192}
193
194impl RowGroups for InMemoryRowGroup<'_> {
195    fn num_rows(&self) -> usize {
196        self.row_count
197    }
198
199    /// Return chunks for column i
200    fn column_chunks(&self, i: usize) -> crate::errors::Result<Box<dyn PageIterator>> {
201        match &self.column_chunks[i] {
202            None => Err(ParquetError::General(format!(
203                "Invalid column index {i}, column was not fetched"
204            ))),
205            Some(data) => {
206                let page_locations = self
207                    .offset_index
208                    // filter out empty offset indexes (old versions specified Some(vec![]) when no present)
209                    .filter(|index| !index.is_empty())
210                    .and_then(|index| index[i].as_ref().map(|idx| idx.page_locations.clone()));
211                let column_chunk_metadata = self.metadata.row_group(self.row_group_idx).column(i);
212                let page_reader = SerializedPageReader::new(
213                    data.clone(),
214                    column_chunk_metadata,
215                    self.row_count,
216                    page_locations,
217                )?;
218                let page_reader = page_reader.add_crypto_context(
219                    self.row_group_idx,
220                    i,
221                    self.metadata,
222                    column_chunk_metadata,
223                )?;
224
225                let page_reader: Box<dyn PageReader> = Box::new(page_reader);
226
227                Ok(Box::new(ColumnChunkIterator {
228                    reader: Some(Ok(page_reader)),
229                }))
230            }
231        }
232    }
233
234    fn row_groups(&self) -> Box<dyn Iterator<Item = &RowGroupMetaData> + '_> {
235        Box::new(std::iter::once(self.metadata.row_group(self.row_group_idx)))
236    }
237
238    fn metadata(&self) -> &ParquetMetaData {
239        self.metadata
240    }
241}
242
243/// An in-memory column chunk.
244/// This allows us to hold either dense column chunks or sparse column chunks and easily
245/// access them by offset.
246#[derive(Clone, Debug)]
247pub(crate) enum ColumnChunkData {
248    /// Column chunk data representing only a subset of data pages.
249    /// For example if a row selection (possibly caused by a filter in a query) causes us to read only
250    /// a subset of the rows in the column.
251    Sparse {
252        /// Length of the full column chunk
253        length: usize,
254        /// Subset of data pages included in this sparse chunk.
255        ///
256        /// Each element is a tuple of (page offset within file, page data).
257        /// Each entry is a complete page and the list is ordered by offset.
258        data: Vec<(usize, Bytes)>,
259    },
260    /// Full column chunk and the offset within the original file
261    Dense { offset: usize, data: Bytes },
262}
263
264impl ColumnChunkData {
265    /// Return the data for this column chunk at the given offset
266    fn get(&self, start: u64) -> crate::errors::Result<Bytes> {
267        match &self {
268            ColumnChunkData::Sparse { data, .. } => data
269                .binary_search_by_key(&start, |(offset, _)| *offset as u64)
270                .map(|idx| data[idx].1.clone())
271                .map_err(|_| {
272                    ParquetError::General(format!(
273                        "Invalid offset in sparse column chunk data: {start}, no matching page found.\
274                         If you are using a `SelectionStrategyPolicy::Mask`, ensure that the OffsetIndex is provided when \
275                         creating the InMemoryRowGroup."
276                    ))
277                }),
278            ColumnChunkData::Dense { offset, data } => {
279                let start = start as usize - *offset;
280                Ok(data.slice(start..))
281            }
282        }
283    }
284}
285
286impl Length for ColumnChunkData {
287    /// Return the total length of the full column chunk
288    fn len(&self) -> u64 {
289        match &self {
290            ColumnChunkData::Sparse { length, .. } => *length as u64,
291            ColumnChunkData::Dense { data, .. } => data.len() as u64,
292        }
293    }
294}
295
296impl ChunkReader for ColumnChunkData {
297    type T = bytes::buf::Reader<Bytes>;
298
299    fn get_read(&self, start: u64) -> crate::errors::Result<Self::T> {
300        Ok(self.get(start)?.reader())
301    }
302
303    fn get_bytes(&self, start: u64, length: usize) -> crate::errors::Result<Bytes> {
304        Ok(self.get(start)?.slice(..length))
305    }
306}
307
308/// Implements [`PageIterator`] for a single column chunk, yielding a single [`PageReader`]
309struct ColumnChunkIterator {
310    reader: Option<crate::errors::Result<Box<dyn PageReader>>>,
311}
312
313impl Iterator for ColumnChunkIterator {
314    type Item = crate::errors::Result<Box<dyn PageReader>>;
315
316    fn next(&mut self) -> Option<Self::Item> {
317        self.reader.take()
318    }
319}
320
321impl PageIterator for ColumnChunkIterator {}
parquet/arrow/in_memory_row_group.rs

parquet/arrow/
in_memory_row_group.rs
