parquet/arrow/arrow_reader/selection/
ranges.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//! Mapping the [`RowSelector`] runs of a selection onto ranges: the byte ranges
19//! of the data pages that must be fetched ([`RowSelection::scan_ranges`]) and
20//! the expansion of a selection to batch boundaries.
21//!
22//! Both are shared by the selector and mask backings, which stream their runs
23//! from a slice and a [`MaskRunIter`] respectively.
24//!
25//! [`MaskRunIter`]: crate::arrow::arrow_reader::MaskRunIter
26
27use super::{RowSelection, RowSelector};
28use crate::file::page_index::offset_index::PageLocation;
29use std::ops::Range;
30
31/// Byte ranges of the data pages containing at least one selected row.
32#[inline]
33pub(super) fn scan_ranges_from_selectors<I>(
34    selectors: I,
35    page_locations: &[PageLocation],
36) -> Vec<Range<u64>>
37where
38    I: IntoIterator<Item = RowSelector>,
39{
40    let mut ranges: Vec<Range<u64>> = vec![];
41    let mut row_offset = 0;
42
43    let mut pages = page_locations.iter().peekable();
44    let mut selectors = selectors.into_iter();
45    let mut current_selector = selectors.next();
46    let mut current_page = pages.next();
47
48    let mut current_page_included = false;
49
50    while let Some((selector, page)) = current_selector.as_mut().zip(current_page) {
51        if !(selector.skip || current_page_included) {
52            let start = page.offset as u64;
53            let end = start + page.compressed_page_size as u64;
54            ranges.push(start..end);
55            current_page_included = true;
56        }
57
58        if let Some(next_page) = pages.peek() {
59            if row_offset + selector.row_count > next_page.first_row_index as usize {
60                let remaining_in_page = next_page.first_row_index as usize - row_offset;
61                selector.row_count -= remaining_in_page;
62                row_offset += remaining_in_page;
63                current_page = pages.next();
64                current_page_included = false;
65
66                continue;
67            } else {
68                if row_offset + selector.row_count == next_page.first_row_index as usize {
69                    current_page = pages.next();
70                    current_page_included = false;
71                }
72                row_offset += selector.row_count;
73                current_selector = selectors.next();
74            }
75        } else {
76            if !(selector.skip || current_page_included) {
77                let start = page.offset as u64;
78                let end = start + page.compressed_page_size as u64;
79                ranges.push(start..end);
80            }
81            current_selector = selectors.next()
82        }
83    }
84
85    ranges
86}
87
88/// Grows each selected run to the batch boundaries containing it, merging the
89/// runs that overlap as a result.
90#[inline]
91pub(super) fn expand_to_batch_boundaries_from_selectors<I>(
92    selectors: I,
93    batch_size: usize,
94    total_rows: usize,
95) -> RowSelection
96where
97    I: IntoIterator<Item = RowSelector>,
98{
99    let mut expanded_ranges = Vec::new();
100    let mut row_offset = 0;
101
102    for selector in selectors {
103        if selector.skip {
104            row_offset += selector.row_count;
105        } else {
106            let start = row_offset;
107            let end = row_offset + selector.row_count;
108
109            // Expand start to batch boundary
110            let expanded_start = (start / batch_size) * batch_size;
111            // Expand end to batch boundary
112            let expanded_end = end.div_ceil(batch_size) * batch_size;
113            let expanded_end = expanded_end.min(total_rows);
114
115            expanded_ranges.push(expanded_start..expanded_end);
116            row_offset += selector.row_count;
117        }
118    }
119
120    // Sort ranges by start position
121    expanded_ranges.sort_by_key(|range| range.start);
122
123    // Merge overlapping or consecutive ranges
124    let mut merged_ranges: Vec<Range<usize>> = Vec::new();
125    for range in expanded_ranges {
126        if let Some(last) = merged_ranges.last_mut() {
127            if range.start <= last.end {
128                // Overlapping or consecutive - merge them
129                last.end = last.end.max(range.end);
130            } else {
131                // No overlap - add new range
132                merged_ranges.push(range);
133            }
134        } else {
135            // First range
136            merged_ranges.push(range);
137        }
138    }
139
140    RowSelection::from_consecutive_ranges(merged_ranges.into_iter(), total_rows)
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn test_scan_ranges() {
149        let index = vec![
150            PageLocation {
151                offset: 0,
152                compressed_page_size: 10,
153                first_row_index: 0,
154            },
155            PageLocation {
156                offset: 10,
157                compressed_page_size: 10,
158                first_row_index: 10,
159            },
160            PageLocation {
161                offset: 20,
162                compressed_page_size: 10,
163                first_row_index: 20,
164            },
165            PageLocation {
166                offset: 30,
167                compressed_page_size: 10,
168                first_row_index: 30,
169            },
170            PageLocation {
171                offset: 40,
172                compressed_page_size: 10,
173                first_row_index: 40,
174            },
175            PageLocation {
176                offset: 50,
177                compressed_page_size: 10,
178                first_row_index: 50,
179            },
180            PageLocation {
181                offset: 60,
182                compressed_page_size: 10,
183                first_row_index: 60,
184            },
185        ];
186
187        let selection = RowSelection::from(vec![
188            // Skip first page
189            RowSelector::skip(10),
190            // Multiple selects in same page
191            RowSelector::select(3),
192            RowSelector::skip(3),
193            RowSelector::select(4),
194            // Select to page boundary
195            RowSelector::skip(5),
196            RowSelector::select(5),
197            // Skip full page past page boundary
198            RowSelector::skip(12),
199            // Select across page boundaries
200            RowSelector::select(12),
201            // Skip final page
202            RowSelector::skip(12),
203        ]);
204
205        let ranges = selection.scan_ranges(&index);
206
207        // assert_eq!(mask, vec![false, true, true, false, true, true, false]);
208        assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60]);
209        assert_eq!(
210            selection.row_ranges_for_selected_pages(&index, 70),
211            vec![10..20, 20..30, 40..50, 50..60]
212        );
213
214        let selection = RowSelection::from(vec![
215            // Skip first page
216            RowSelector::skip(10),
217            // Multiple selects in same page
218            RowSelector::select(3),
219            RowSelector::skip(3),
220            RowSelector::select(4),
221            // Select to page boundary
222            RowSelector::skip(5),
223            RowSelector::select(5),
224            // Skip full page past page boundary
225            RowSelector::skip(12),
226            // Select across page boundaries
227            RowSelector::select(12),
228            RowSelector::skip(1),
229            // Select across page boundaries including final page
230            RowSelector::select(8),
231        ]);
232
233        let ranges = selection.scan_ranges(&index);
234
235        // assert_eq!(mask, vec![false, true, true, false, true, true, true]);
236        assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]);
237
238        let selection = RowSelection::from(vec![
239            // Skip first page
240            RowSelector::skip(10),
241            // Multiple selects in same page
242            RowSelector::select(3),
243            RowSelector::skip(3),
244            RowSelector::select(4),
245            // Select to page boundary
246            RowSelector::skip(5),
247            RowSelector::select(5),
248            // Skip full page past page boundary
249            RowSelector::skip(12),
250            // Select to final page boundary
251            RowSelector::select(12),
252            RowSelector::skip(1),
253            // Skip across final page boundary
254            RowSelector::skip(8),
255            // Select from final page
256            RowSelector::select(4),
257        ]);
258
259        let ranges = selection.scan_ranges(&index);
260
261        // assert_eq!(mask, vec![false, true, true, false, true, true, true]);
262        assert_eq!(ranges, vec![10..20, 20..30, 40..50, 50..60, 60..70]);
263
264        let selection = RowSelection::from(vec![
265            // Skip first page
266            RowSelector::skip(10),
267            // Multiple selects in same page
268            RowSelector::select(3),
269            RowSelector::skip(3),
270            RowSelector::select(4),
271            // Select to remaining in page and first row of next page
272            RowSelector::skip(5),
273            RowSelector::select(6),
274            // Skip remaining
275            RowSelector::skip(50),
276        ]);
277
278        let ranges = selection.scan_ranges(&index);
279
280        // assert_eq!(mask, vec![false, true, true, false, true, true, true]);
281        assert_eq!(ranges, vec![10..20, 20..30, 30..40]);
282    }
283}
parquet/arrow/arrow_reader/selection/ranges.rs

parquet/arrow/arrow_reader/selection/
ranges.rs
