parquet/arrow/push_decoder/reader_builder/
data.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//! [`DataRequest`] tracks and holds data needed to construct InMemoryRowGroups
19
20use crate::arrow::ProjectionMask;
21use crate::arrow::arrow_reader::RowSelection;
22use crate::arrow::in_memory_row_group::{ColumnChunkData, FetchRanges, InMemoryRowGroup};
23use crate::errors::ParquetError;
24use crate::file::metadata::ParquetMetaData;
25use crate::file::reader::ChunkReader;
26use crate::util::push_buffers::PushBuffers;
27use bytes::Bytes;
28use std::ops::Range;
29use std::sync::Arc;
30
31/// Contains in-progress state to construct InMemoryRowGroups
32///
33/// See [`DataRequestBuilder`] for creating new requests
34#[derive(Debug)]
35pub(super) struct DataRequest {
36    /// Any previously read column chunk data
37    column_chunks: Vec<Option<Arc<ColumnChunkData>>>,
38    /// The ranges of data that are needed next
39    ranges: Vec<Range<u64>>,
40    /// Optional page start offsets for each requested range. This is used
41    /// to create the relevant InMemoryRowGroup
42    page_start_offsets: Option<Vec<Vec<u64>>>,
43}
44
45impl DataRequest {
46    /// return what ranges are still needed to satisfy this request. Returns an empty vec
47    /// if all ranges are satisfied
48    pub fn needed_ranges(&self, buffers: &PushBuffers) -> Vec<Range<u64>> {
49        self.ranges
50            .iter()
51            .filter(|&range| !buffers.has_range(range))
52            .cloned()
53            .collect()
54    }
55
56    /// Returns the chunks from the buffers that satisfy this request
57    fn get_chunks(&self, buffers: &PushBuffers) -> Result<Vec<Bytes>, ParquetError> {
58        self.ranges
59            .iter()
60            .map(|range| {
61                let length: usize = (range.end - range.start)
62                    .try_into()
63                    .expect("overflow for offset");
64                // should have all the data due to the check above
65                buffers.get_bytes(range.start, length).map_err(|e| {
66                    ParquetError::General(format!(
67                        "Internal Error missing data for range {range:?} in buffers: {e}",
68                    ))
69                })
70            })
71            .collect()
72    }
73
74    /// Create a new InMemoryRowGroup, and fill it with provided data
75    ///
76    /// Assumes that all needed data is present in the buffers
77    /// and clears any explicitly requested ranges
78    pub fn try_into_in_memory_row_group<'a>(
79        self,
80        row_group_idx: usize,
81        row_count: usize,
82        parquet_metadata: &'a ParquetMetaData,
83        projection: &ProjectionMask,
84        buffers: &mut PushBuffers,
85    ) -> Result<InMemoryRowGroup<'a>, ParquetError> {
86        let chunks = self.get_chunks(buffers)?;
87
88        let Self {
89            column_chunks,
90            ranges,
91            page_start_offsets,
92        } = self;
93
94        let page_index = if parquet_metadata
95            .page_index()
96            .is_some_and(|pi| pi.has_offset_indexes())
97        {
98            Some(parquet_metadata.page_index_for_row_group(row_group_idx))
99        } else {
100            None
101        };
102
103        // Create an InMemoryRowGroup to hold the column chunks, this is a
104        // temporary structure used to tell the ArrowReaders what pages are
105        // needed for decoding
106        let mut in_memory_row_group = InMemoryRowGroup {
107            row_count,
108            column_chunks,
109            page_index,
110            row_group_idx,
111            metadata: parquet_metadata,
112        };
113
114        in_memory_row_group.fill_column_chunks(projection, page_start_offsets, chunks);
115
116        // Clear the ranges that were explicitly requested
117        buffers.clear_ranges(&ranges);
118
119        Ok(in_memory_row_group)
120    }
121}
122
123/// Builder for [`DataRequest`]
124pub(super) struct DataRequestBuilder<'a> {
125    /// The row group index
126    row_group_idx: usize,
127    /// The number of rows in the row group
128    row_count: usize,
129    /// The batch size to read
130    batch_size: usize,
131    /// The parquet metadata
132    parquet_metadata: &'a ParquetMetaData,
133    /// The projection mask (which columns to read)
134    projection: &'a ProjectionMask,
135    /// Optional row selection to apply
136    selection: Option<&'a RowSelection>,
137    /// Optional projection mask if using
138    /// [`RowGroupCache`](crate::arrow::array_reader::RowGroupCache)
139    /// for caching decoded columns.
140    cache_projection: Option<&'a ProjectionMask>,
141    /// Any previously read column chunks
142    column_chunks: Option<Vec<Option<Arc<ColumnChunkData>>>>,
143}
144
145impl<'a> DataRequestBuilder<'a> {
146    pub(super) fn new(
147        row_group_idx: usize,
148        row_count: usize,
149        batch_size: usize,
150        parquet_metadata: &'a ParquetMetaData,
151        projection: &'a ProjectionMask,
152    ) -> Self {
153        Self {
154            row_group_idx,
155            row_count,
156            batch_size,
157            parquet_metadata,
158            projection,
159            selection: None,
160            cache_projection: None,
161            column_chunks: None,
162        }
163    }
164
165    /// Set an optional row selection to apply
166    pub(super) fn with_selection(mut self, selection: Option<&'a RowSelection>) -> Self {
167        self.selection = selection;
168        self
169    }
170
171    /// set columns to cache, if any
172    pub(super) fn with_cache_projection(
173        mut self,
174        cache_projection: Option<&'a ProjectionMask>,
175    ) -> Self {
176        self.cache_projection = cache_projection;
177        self
178    }
179
180    /// Provide any previously read column chunks
181    pub(super) fn with_column_chunks(
182        mut self,
183        column_chunks: Option<Vec<Option<Arc<ColumnChunkData>>>>,
184    ) -> Self {
185        self.column_chunks = column_chunks;
186        self
187    }
188
189    pub(crate) fn build(self) -> DataRequest {
190        let Self {
191            row_group_idx,
192            row_count,
193            batch_size,
194            parquet_metadata,
195            projection,
196            selection,
197            cache_projection,
198            column_chunks,
199        } = self;
200
201        let row_group_meta_data = parquet_metadata.row_group(row_group_idx);
202
203        // If no previously read column chunks are provided, create a new location to hold them
204        let column_chunks =
205            column_chunks.unwrap_or_else(|| vec![None; row_group_meta_data.columns().len()]);
206
207        let page_index = if parquet_metadata
208            .page_index()
209            .is_some_and(|pi| pi.has_offset_indexes())
210        {
211            Some(parquet_metadata.page_index_for_row_group(row_group_idx))
212        } else {
213            None
214        };
215
216        // Create an InMemoryRowGroup to hold the column chunks, this is a
217        // temporary structure used to tell the ArrowReaders what pages are
218        // needed for decoding
219        let row_group = InMemoryRowGroup {
220            row_count,
221            column_chunks,
222            page_index,
223            row_group_idx,
224            metadata: parquet_metadata,
225        };
226
227        let FetchRanges {
228            ranges,
229            page_start_offsets,
230        } = row_group.fetch_ranges(projection, selection, batch_size, cache_projection);
231
232        DataRequest {
233            // Save any previously read column chunks
234            column_chunks: row_group.column_chunks,
235            ranges,
236            page_start_offsets,
237        }
238    }
239}
parquet/arrow/push_decoder/reader_builder/data.rs

parquet/arrow/push_decoder/reader_builder/
data.rs
