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::page_index::RowGroupPageIndex;
24use crate::file::metadata::{ParquetMetaData, RowGroupMetaData};
25use crate::file::page_index::offset_index::PageLocation;
26use crate::file::reader::{ChunkReader, Length, SerializedPageReader};
27use bytes::{Buf, Bytes};
28use std::ops::Range;
29use std::sync::Arc;
30
31/// An in-memory collection of column chunks
32#[derive(Debug)]
33pub(crate) struct InMemoryRowGroup<'a> {
34    pub(crate) page_index: Option<RowGroupPageIndex>,
35    /// Column chunks for this row group
36    pub(crate) column_chunks: Vec<Option<Arc<ColumnChunkData>>>,
37    pub(crate) row_count: usize,
38    pub(crate) row_group_idx: usize,
39    pub(crate) metadata: &'a ParquetMetaData,
40}
41
42/// What ranges to fetch for the columns in this row group
43#[derive(Debug)]
44pub(crate) struct FetchRanges {
45    /// The byte ranges to fetch
46    pub(crate) ranges: Vec<Range<u64>>,
47    /// If `Some`, the start offsets of each page for each column chunk, or
48    /// `None` for a column chunk without an offset index (fetched in full)
49    pub(crate) page_start_offsets: Option<Vec<Option<Vec<u64>>>>,
50}
51
52impl InMemoryRowGroup<'_> {
53    /// Returns the byte ranges to fetch for the columns specified in
54    /// `projection` and `selection`.
55    ///
56    /// `cache_mask` indicates which columns, if any, are being cached by
57    /// [`RowGroupCache`](crate::arrow::array_reader::RowGroupCache).
58    /// The `selection` for Cached columns is expanded to batch boundaries to simplify
59    /// accounting for what data is cached.
60    ///
61    /// The ranges of each column chunk come from [`ColumnFetch`].
62    pub(crate) fn fetch_ranges(
63        &self,
64        projection: &ProjectionMask,
65        selection: Option<&RowSelection>,
66        batch_size: usize,
67        cache_mask: Option<&ProjectionMask>,
68    ) -> FetchRanges {
69        let metadata = self.metadata.row_group(self.row_group_idx);
70        // With a `RowSelection` and an `OffsetIndex`, only fetch the pages
71        // required for the `RowSelection`
72        let page_index = self.page_index.as_ref().filter(|_| selection.is_some());
73        let expanded_selection = selection
74            .filter(|_| page_index.is_some() && cache_mask.is_some())
75            .map(|selection| selection.expand_to_batch_boundaries(batch_size, self.row_count));
76
77        // Consider preallocating outer vec: https://github.com/apache/arrow-rs/issues/8667
78        let mut page_start_offsets: Option<Vec<Option<Vec<u64>>>> = page_index.map(|_| vec![]);
79        let mut ranges = vec![];
80        let columns = columns_to_fetch(projection, self.column_chunks.len(), |idx| {
81            self.column_chunks[idx].is_some()
82        });
83        for idx in columns {
84            let locations = page_index
85                .and_then(|page_index| page_index.offset_index(idx))
86                .map(|offset_index| offset_index.page_locations().as_slice());
87            let column_selection =
88                column_selection(selection, expanded_selection.as_ref(), cache_mask, idx);
89            let (start, len) = metadata.column(idx).byte_range();
90            let fetch = ColumnFetch::new(start..start + len, locations, column_selection);
91            let first = ranges.len();
92            ranges.extend(fetch.ranges());
93            if let Some(page_start_offsets) = page_start_offsets.as_mut() {
94                page_start_offsets.push(match fetch {
95                    // No offset index for this column, fetch the entire column
96                    ColumnFetch::Chunk { .. } => None,
97                    ColumnFetch::Pages { .. } => {
98                        Some(ranges[first..].iter().map(|range| range.start).collect())
99                    }
100                });
101            }
102        }
103        FetchRanges {
104            ranges,
105            page_start_offsets,
106        }
107    }
108
109    /// Fills in `self.column_chunks` with the data fetched from `chunk_data`.
110    ///
111    /// This function **must** be called with the data from the ranges returned by
112    /// `fetch_ranges` and the corresponding page_start_offsets, with the exact same and `selection`.
113    pub(crate) fn fill_column_chunks<I>(
114        &mut self,
115        projection: &ProjectionMask,
116        page_start_offsets: Option<Vec<Option<Vec<u64>>>>,
117        chunk_data: I,
118    ) where
119        I: IntoIterator<Item = Bytes>,
120    {
121        let mut chunk_data = chunk_data.into_iter();
122        let metadata = self.metadata.row_group(self.row_group_idx);
123        if let Some(page_start_offsets) = page_start_offsets {
124            // If we have a `RowSelection` and an `OffsetIndex` then only fetch pages required for the
125            // `RowSelection`
126            let mut page_start_offsets = page_start_offsets.into_iter();
127
128            for (idx, chunk) in self.column_chunks.iter_mut().enumerate() {
129                if chunk.is_some() || !projection.leaf_included(idx) {
130                    continue;
131                }
132
133                match page_start_offsets.next() {
134                    // No offset index: `fetch_ranges` requested the whole chunk
135                    Some(None) => {
136                        if let Some(data) = chunk_data.next() {
137                            *chunk = Some(Arc::new(ColumnChunkData::Dense {
138                                offset: metadata.column(idx).byte_range().0 as usize,
139                                data,
140                            }));
141                        }
142                    }
143                    Some(Some(offsets)) => {
144                        let mut chunks = Vec::with_capacity(offsets.len());
145                        for _ in 0..offsets.len() {
146                            chunks.push(chunk_data.next().unwrap());
147                        }
148
149                        *chunk = Some(Arc::new(ColumnChunkData::Sparse {
150                            length: metadata.column(idx).byte_range().1 as usize,
151                            data: offsets
152                                .into_iter()
153                                .map(|x| x as usize)
154                                .zip(chunks)
155                                .collect(),
156                        }))
157                    }
158                    None => {}
159                }
160            }
161        } else {
162            for (idx, chunk) in self.column_chunks.iter_mut().enumerate() {
163                if chunk.is_some() || !projection.leaf_included(idx) {
164                    continue;
165                }
166
167                if let Some(data) = chunk_data.next() {
168                    *chunk = Some(Arc::new(ColumnChunkData::Dense {
169                        offset: metadata.column(idx).byte_range().0 as usize,
170                        data,
171                    }));
172                }
173            }
174        }
175    }
176}
177
178/// The leaf columns in `projection` that are not yet read, in column order.
179///
180/// A decoding stage of a row group does not fetch a column that an earlier
181/// stage of the same row group read.
182#[inline]
183pub(crate) fn columns_to_fetch<'a>(
184    projection: &'a ProjectionMask,
185    num_columns: usize,
186    is_read: impl Fn(usize) -> bool + 'a,
187) -> impl Iterator<Item = usize> + 'a {
188    (0..num_columns).filter(move |&idx| projection.leaf_included(idx) && !is_read(idx))
189}
190
191/// The selection that [`InMemoryRowGroup::fetch_ranges`] uses to choose the
192/// pages of column `idx`: `expanded_selection` (the selection expanded to
193/// batch boundaries) if `cache_mask` includes the column, else `selection`.
194#[inline]
195pub(crate) fn column_selection<'a>(
196    selection: Option<&'a RowSelection>,
197    expanded_selection: Option<&'a RowSelection>,
198    cache_mask: Option<&ProjectionMask>,
199    idx: usize,
200) -> Option<&'a RowSelection> {
201    match expanded_selection {
202        Some(expanded) if cache_mask.is_some_and(|mask| mask.leaf_included(idx)) => Some(expanded),
203        _ => selection,
204    }
205}
206
207/// What [`InMemoryRowGroup::fetch_ranges`] fetches for one column chunk.
208///
209/// This is the one place that decides which bytes of a column chunk are
210/// requested, so other users of this rule cannot differ from the decoder.
211#[derive(Debug, Clone, PartialEq, Eq)]
212pub(crate) enum ColumnFetch<'a> {
213    /// The whole column chunk, as one range. This is the case if there is
214    /// no selection or the column has no offset index.
215    Chunk {
216        /// Byte range of the column chunk.
217        range: Range<u64>,
218    },
219    /// The dictionary page, if any, and the data pages that contain a
220    /// selected row.
221    Pages {
222        /// Byte range of the dictionary page, if the column chunk has one.
223        dictionary: Option<Range<u64>>,
224        /// The page locations of the column chunk.
225        locations: &'a [PageLocation],
226        /// Indexes in `locations` of the data pages to fetch, in page order.
227        pages: Vec<usize>,
228    },
229}
230
231impl<'a> ColumnFetch<'a> {
232    /// The fetch of the column chunk at byte range `chunk`, with page
233    /// locations `locations` from its offset index, if any, and row
234    /// selection `selection`, if any.
235    #[inline]
236    pub(crate) fn new(
237        chunk: Range<u64>,
238        locations: Option<&'a [PageLocation]>,
239        selection: Option<&RowSelection>,
240    ) -> Self {
241        let (Some(selection), Some(locations)) = (selection, locations) else {
242            return Self::Chunk { range: chunk };
243        };
244        // `scan_ranges` returns the ranges of the selected pages in page
245        // order. Map them back to page indexes.
246        let fetched = selection.scan_ranges(locations);
247        let mut fetched = fetched.iter().peekable();
248        let pages = locations
249            .iter()
250            .enumerate()
251            .filter(|(_, location)| {
252                fetched
253                    .next_if(|range| range.start == location.offset as u64)
254                    .is_some()
255            })
256            .map(|(idx, _)| idx)
257            .collect();
258        Self::Pages {
259            dictionary: dictionary_range(chunk.start, locations),
260            locations,
261            pages,
262        }
263    }
264
265    /// The byte ranges to fetch, in file order.
266    pub(crate) fn ranges(&self) -> impl Iterator<Item = Range<u64>> + '_ {
267        let (chunk, dictionary, pages) = match self {
268            Self::Chunk { range, .. } => (Some(range.clone()), None, None),
269            Self::Pages {
270                dictionary,
271                locations,
272                pages,
273            } => (None, dictionary.clone(), Some((*locations, pages))),
274        };
275        let data_pages = pages
276            .into_iter()
277            .flat_map(|(locations, pages)| pages.iter().map(|&page| page_range(&locations[page])));
278        chunk.into_iter().chain(dictionary).chain(data_pages)
279    }
280}
281
282/// Byte range of the dictionary page of a column chunk that starts at
283/// `chunk_start`: the bytes before the first data page, if any.
284#[inline]
285pub(crate) fn dictionary_range(chunk_start: u64, locations: &[PageLocation]) -> Option<Range<u64>> {
286    match locations.first() {
287        Some(first) if first.offset as u64 != chunk_start => Some(chunk_start..first.offset as u64),
288        _ => None,
289    }
290}
291
292/// Byte range of a data page.
293#[inline]
294pub(crate) fn page_range(location: &PageLocation) -> Range<u64> {
295    let start = location.offset as u64;
296    start..start + location.compressed_page_size as u64
297}
298
299impl RowGroups for InMemoryRowGroup<'_> {
300    fn num_rows(&self) -> usize {
301        self.row_count
302    }
303
304    /// Return chunks for column i
305    fn column_chunks(&self, i: usize) -> crate::errors::Result<Box<dyn PageIterator>> {
306        match &self.column_chunks[i] {
307            None => Err(ParquetError::General(format!(
308                "Invalid column index {i}, column was not fetched"
309            ))),
310            Some(data) => {
311                let page_locations = self
312                    .page_index
313                    .as_ref()
314                    .and_then(|pi| pi.page_locations(i).cloned());
315                let column_chunk_metadata = self.metadata.row_group(self.row_group_idx).column(i);
316                let page_reader = SerializedPageReader::new(
317                    data.clone(),
318                    column_chunk_metadata,
319                    self.row_count,
320                    page_locations,
321                )?;
322                let page_reader = page_reader.add_crypto_context(
323                    self.row_group_idx,
324                    i,
325                    self.metadata,
326                    column_chunk_metadata,
327                )?;
328
329                let page_reader: Box<dyn PageReader> = Box::new(page_reader);
330
331                Ok(Box::new(ColumnChunkIterator {
332                    reader: Some(Ok(page_reader)),
333                }))
334            }
335        }
336    }
337
338    fn row_groups(&self) -> Box<dyn Iterator<Item = &RowGroupMetaData> + '_> {
339        Box::new(std::iter::once(self.metadata.row_group(self.row_group_idx)))
340    }
341
342    fn metadata(&self) -> &ParquetMetaData {
343        self.metadata
344    }
345}
346
347/// An in-memory column chunk.
348/// This allows us to hold either dense column chunks or sparse column chunks and easily
349/// access them by offset.
350#[derive(Clone, Debug)]
351pub(crate) enum ColumnChunkData {
352    /// Column chunk data representing only a subset of data pages.
353    /// For example if a row selection (possibly caused by a filter in a query) causes us to read only
354    /// a subset of the rows in the column.
355    Sparse {
356        /// Length of the full column chunk
357        length: usize,
358        /// Subset of data pages included in this sparse chunk.
359        ///
360        /// Each element is a tuple of (page offset within file, page data).
361        /// Each entry is a complete page and the list is ordered by offset.
362        data: Vec<(usize, Bytes)>,
363    },
364    /// Full column chunk and the offset within the original file
365    Dense { offset: usize, data: Bytes },
366}
367
368impl ColumnChunkData {
369    /// Return the data for this column chunk at the given offset
370    fn get(&self, start: u64) -> crate::errors::Result<Bytes> {
371        match &self {
372            ColumnChunkData::Sparse { data, .. } => data
373                .binary_search_by_key(&start, |(offset, _)| *offset as u64)
374                .map(|idx| data[idx].1.clone())
375                .map_err(|_| {
376                    ParquetError::General(format!(
377                        "Invalid offset in sparse column chunk data: {start}, no matching page found.\
378                         If you are using a `SelectionStrategyPolicy::Mask`, ensure that the OffsetIndex is provided when \
379                         creating the InMemoryRowGroup."
380                    ))
381                }),
382            ColumnChunkData::Dense { offset, data } => {
383                let start = start as usize - *offset;
384                Ok(data.slice(start..))
385            }
386        }
387    }
388}
389
390impl Length for ColumnChunkData {
391    /// Return the total length of the full column chunk
392    fn len(&self) -> u64 {
393        match &self {
394            ColumnChunkData::Sparse { length, .. } => *length as u64,
395            ColumnChunkData::Dense { data, .. } => data.len() as u64,
396        }
397    }
398}
399
400impl ChunkReader for ColumnChunkData {
401    type T = bytes::buf::Reader<Bytes>;
402
403    fn get_read(&self, start: u64) -> crate::errors::Result<Self::T> {
404        Ok(self.get(start)?.reader())
405    }
406
407    fn get_bytes(&self, start: u64, length: usize) -> crate::errors::Result<Bytes> {
408        let data = self.get(start)?;
409        if data.len() < length {
410            return Err(general_err!(
411                "column chunk data at offset {start} has {} bytes, expected {length}",
412                data.len()
413            ));
414        }
415        Ok(data.slice(..length))
416    }
417}
418
419/// Implements [`PageIterator`] for a single column chunk, yielding a single [`PageReader`]
420struct ColumnChunkIterator {
421    reader: Option<crate::errors::Result<Box<dyn PageReader>>>,
422}
423
424impl Iterator for ColumnChunkIterator {
425    type Item = crate::errors::Result<Box<dyn PageReader>>;
426
427    fn next(&mut self) -> Option<Self::Item> {
428        self.reader.take()
429    }
430}
431
432impl PageIterator for ColumnChunkIterator {}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437
438    #[test]
439    fn get_bytes_errors_on_short_data() {
440        let dense = ColumnChunkData::Dense {
441            offset: 100,
442            data: Bytes::from_static(b"0123456789"),
443        };
444        assert_eq!(
445            dense.get_bytes(105, 5).unwrap(),
446            Bytes::from_static(b"56789")
447        );
448        let err = dense.get_bytes(105, 6).unwrap_err().to_string();
449        assert!(err.contains("has 5 bytes, expected 6"), "{err}");
450    }
451}
parquet/arrow/in_memory_row_group.rs

parquet/arrow/
in_memory_row_group.rs
