parquet/arrow/arrow_reader/selection/
selector.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//! The run length backed representation of a [`RowSelection`]: [`RowSelector`]
19//! and the primitives operating on a sequence of them.
20//!
21//! This is the counterpart of the bitmap backing in the `boolean` module, and
22//! provides the same set of transforms (`split_off`, `trim`, `offset`,
23//! `limit`) over `Vec<RowSelector>` instead of a `BooleanBuffer`.
24//!
25//! [`RowSelection`]: crate::arrow::arrow_reader::RowSelection
26
27/// [`RowSelection`] is a collection of [`RowSelector`] used to skip rows when
28/// scanning a parquet file
29///
30/// [`RowSelection`]: crate::arrow::arrow_reader::RowSelection
31#[derive(Debug, Clone, Copy, Eq, PartialEq)]
32pub struct RowSelector {
33    /// The number of rows
34    pub row_count: usize,
35
36    /// If true, skip `row_count` rows
37    pub skip: bool,
38}
39
40impl RowSelector {
41    /// Select `row_count` rows
42    pub fn select(row_count: usize) -> Self {
43        Self {
44            row_count,
45            skip: false,
46        }
47    }
48
49    /// Skip `row_count` rows
50    pub fn skip(row_count: usize) -> Self {
51        Self {
52            row_count,
53            skip: true,
54        }
55    }
56}
57
58/// Splits `selectors` at the first `row_count` rows, returning `(head, tail)`.
59pub(super) fn split_off_selectors(
60    mut selectors: Vec<RowSelector>,
61    row_count: usize,
62) -> (Vec<RowSelector>, Vec<RowSelector>) {
63    let mut total_count = 0;
64
65    // Find the index where the selector exceeds the row count
66    let find = selectors.iter().position(|selector| {
67        total_count += selector.row_count;
68        total_count > row_count
69    });
70
71    let split_idx = match find {
72        Some(idx) => idx,
73        None => return (selectors, Vec::new()),
74    };
75
76    // `selectors` keeps the head, `tail` takes the rest. The selector straddling
77    // the boundary is split between the two.
78    let mut tail = selectors.split_off(split_idx);
79
80    // Always present as `split_idx < selectors.len`
81    let next = tail.first_mut().unwrap();
82    let overflow = total_count - row_count;
83
84    if next.row_count != overflow {
85        selectors.push(RowSelector {
86            row_count: next.row_count - overflow,
87            skip: next.skip,
88        })
89    }
90    next.row_count = overflow;
91
92    (selectors, tail)
93}
94
95/// Skips the first `offset` selected rows of `selectors`.
96pub(super) fn offset_selectors(mut selectors: Vec<RowSelector>, offset: usize) -> Vec<RowSelector> {
97    let mut selected_count = 0;
98    let mut skipped_count = 0;
99
100    // Find the index where the selector exceeds the row count
101    let find = selectors.iter().position(|selector| match selector.skip {
102        true => {
103            skipped_count += selector.row_count;
104            false
105        }
106        false => {
107            selected_count += selector.row_count;
108            selected_count > offset
109        }
110    });
111
112    let split_idx = match find {
113        Some(idx) => idx,
114        None => {
115            selectors.clear();
116            return selectors;
117        }
118    };
119
120    let mut new_selectors = Vec::with_capacity(selectors.len() - split_idx + 1);
121    new_selectors.push(RowSelector::skip(skipped_count + offset));
122    new_selectors.push(RowSelector::select(selected_count - offset));
123    new_selectors.extend_from_slice(&selectors[split_idx + 1..]);
124
125    new_selectors
126}
127
128/// Keeps only the first `limit` selected rows of `selectors`.
129pub(super) fn limit_selectors(
130    mut selectors: Vec<RowSelector>,
131    mut limit: usize,
132) -> Vec<RowSelector> {
133    if limit == 0 {
134        selectors.clear();
135    }
136
137    for (idx, selection) in selectors.iter_mut().enumerate() {
138        if !selection.skip {
139            if selection.row_count >= limit {
140                selection.row_count = limit;
141                selectors.truncate(idx + 1);
142                break;
143            } else {
144                limit -= selection.row_count;
145            }
146        }
147    }
148    selectors
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::arrow::arrow_reader::selection::RowSelection;
155
156    #[test]
157    fn test_from_selectors_skips_empty_selectors() {
158        let selection = RowSelection::from(vec![
159            RowSelector::select(0),
160            RowSelector::skip(0),
161            RowSelector::select(2),
162            RowSelector::select(0),
163            RowSelector::skip(1),
164        ]);
165        assert_eq!(
166            selection.selectors(),
167            vec![RowSelector::select(2), RowSelector::skip(1)]
168        );
169    }
170
171    #[test]
172    fn test_split_off() {
173        let mut selection = RowSelection::from(vec![
174            RowSelector::skip(34),
175            RowSelector::select(12),
176            RowSelector::skip(3),
177            RowSelector::select(35),
178        ]);
179
180        let split = selection.split_off(34);
181        assert_eq!(split.selectors(), vec![RowSelector::skip(34)]);
182        assert_eq!(
183            selection.selectors(),
184            vec![
185                RowSelector::select(12),
186                RowSelector::skip(3),
187                RowSelector::select(35)
188            ]
189        );
190
191        let split = selection.split_off(5);
192        assert_eq!(split.selectors(), vec![RowSelector::select(5)]);
193        assert_eq!(
194            selection.selectors(),
195            vec![
196                RowSelector::select(7),
197                RowSelector::skip(3),
198                RowSelector::select(35)
199            ]
200        );
201
202        let split = selection.split_off(8);
203        assert_eq!(
204            split.selectors(),
205            vec![RowSelector::select(7), RowSelector::skip(1)]
206        );
207        assert_eq!(
208            selection.selectors(),
209            vec![RowSelector::skip(2), RowSelector::select(35)]
210        );
211
212        let split = selection.split_off(200);
213        assert_eq!(
214            split.selectors(),
215            vec![RowSelector::skip(2), RowSelector::select(35)]
216        );
217        assert!(selection.selectors().is_empty());
218    }
219
220    #[test]
221    fn test_offset() {
222        let selection = RowSelection::from(vec![
223            RowSelector::select(5),
224            RowSelector::skip(23),
225            RowSelector::select(7),
226            RowSelector::skip(33),
227            RowSelector::select(6),
228        ]);
229
230        let selection = selection.offset(2);
231        assert_eq!(
232            selection.selectors(),
233            vec![
234                RowSelector::skip(2),
235                RowSelector::select(3),
236                RowSelector::skip(23),
237                RowSelector::select(7),
238                RowSelector::skip(33),
239                RowSelector::select(6),
240            ]
241        );
242
243        let selection = selection.offset(5);
244        assert_eq!(
245            selection.selectors(),
246            vec![
247                RowSelector::skip(30),
248                RowSelector::select(5),
249                RowSelector::skip(33),
250                RowSelector::select(6),
251            ]
252        );
253
254        let selection = selection.offset(3);
255        assert_eq!(
256            selection.selectors(),
257            vec![
258                RowSelector::skip(33),
259                RowSelector::select(2),
260                RowSelector::skip(33),
261                RowSelector::select(6),
262            ]
263        );
264
265        let selection = selection.offset(2);
266        assert_eq!(
267            selection.selectors(),
268            vec![RowSelector::skip(68), RowSelector::select(6),]
269        );
270
271        let selection = selection.offset(3);
272        assert_eq!(
273            selection.selectors(),
274            vec![RowSelector::skip(71), RowSelector::select(3),]
275        );
276    }
277
278    #[test]
279    fn test_combine() {
280        let a = vec![
281            RowSelector::skip(3),
282            RowSelector::skip(3),
283            RowSelector::select(10),
284            RowSelector::skip(4),
285        ];
286
287        let b = vec![
288            RowSelector::skip(3),
289            RowSelector::skip(3),
290            RowSelector::select(10),
291            RowSelector::skip(4),
292            RowSelector::skip(0),
293        ];
294
295        let c = vec![
296            RowSelector::skip(2),
297            RowSelector::skip(4),
298            RowSelector::select(3),
299            RowSelector::select(3),
300            RowSelector::select(4),
301            RowSelector::skip(3),
302            RowSelector::skip(1),
303            RowSelector::skip(0),
304        ];
305
306        let expected = RowSelection::from(vec![
307            RowSelector::skip(6),
308            RowSelector::select(10),
309            RowSelector::skip(4),
310        ]);
311
312        assert_eq!(RowSelection::from_iter(a), expected);
313        assert_eq!(RowSelection::from_iter(b), expected);
314        assert_eq!(RowSelection::from_iter(c), expected);
315    }
316
317    #[test]
318    fn test_combine_2elements() {
319        let a = vec![RowSelector::select(10), RowSelector::select(5)];
320        let a_expect = vec![RowSelector::select(15)];
321        assert_eq!(RowSelection::from_iter(a).selectors(), a_expect);
322
323        let b = vec![RowSelector::select(10), RowSelector::skip(5)];
324        let b_expect = vec![RowSelector::select(10), RowSelector::skip(5)];
325        assert_eq!(RowSelection::from_iter(b).selectors(), b_expect);
326
327        let c = vec![RowSelector::skip(10), RowSelector::select(5)];
328        let c_expect = vec![RowSelector::skip(10), RowSelector::select(5)];
329        assert_eq!(RowSelection::from_iter(c).selectors(), c_expect);
330
331        let d = vec![RowSelector::skip(10), RowSelector::skip(5)];
332        let d_expect = vec![RowSelector::skip(15)];
333        assert_eq!(RowSelection::from_iter(d).selectors(), d_expect);
334    }
335
336    #[test]
337    fn test_from_one_and_empty() {
338        let a = vec![RowSelector::select(10)];
339        let selection1 = RowSelection::from(a.clone());
340        assert_eq!(selection1.selectors(), a);
341
342        let b = vec![];
343        let selection1 = RowSelection::from(b.clone());
344        assert_eq!(selection1.selectors(), b)
345    }
346
347    #[test]
348    fn test_limit() {
349        // Limit to existing limit should no-op
350        let selection = RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(90)]);
351        let limited = selection.limit(10);
352        assert_eq!(RowSelection::from(vec![RowSelector::select(10)]), limited);
353
354        let selection = RowSelection::from(vec![
355            RowSelector::select(10),
356            RowSelector::skip(10),
357            RowSelector::select(10),
358            RowSelector::skip(10),
359            RowSelector::select(10),
360        ]);
361
362        let limited = selection.clone().limit(5);
363        let expected = vec![RowSelector::select(5)];
364        assert_eq!(limited.selectors(), expected);
365
366        let limited = selection.clone().limit(15);
367        let expected = vec![
368            RowSelector::select(10),
369            RowSelector::skip(10),
370            RowSelector::select(5),
371        ];
372        assert_eq!(limited.selectors(), expected);
373
374        let limited = selection.clone().limit(0);
375        let expected = vec![];
376        assert_eq!(limited.selectors(), expected);
377
378        let limited = selection.clone().limit(30);
379        let expected = vec![
380            RowSelector::select(10),
381            RowSelector::skip(10),
382            RowSelector::select(10),
383            RowSelector::skip(10),
384            RowSelector::select(10),
385        ];
386        assert_eq!(limited.selectors(), expected);
387
388        let limited = selection.limit(100);
389        let expected = vec![
390            RowSelector::select(10),
391            RowSelector::skip(10),
392            RowSelector::select(10),
393            RowSelector::skip(10),
394            RowSelector::select(10),
395        ];
396        assert_eq!(limited.selectors(), expected);
397    }
398
399    #[test]
400    fn test_from_ranges() {
401        let ranges = [1..3, 4..6, 6..6, 8..8, 9..10];
402        let selection = RowSelection::from_consecutive_ranges(ranges.into_iter(), 10);
403        assert_eq!(
404            selection.selectors(),
405            vec![
406                RowSelector::skip(1),
407                RowSelector::select(2),
408                RowSelector::skip(1),
409                RowSelector::select(2),
410                RowSelector::skip(3),
411                RowSelector::select(1)
412            ]
413        );
414
415        let out_of_order_ranges = [1..3, 8..10, 4..7];
416        let result = std::panic::catch_unwind(|| {
417            RowSelection::from_consecutive_ranges(out_of_order_ranges.into_iter(), 10)
418        });
419        assert!(result.is_err());
420    }
421
422    #[test]
423    fn test_empty_selector() {
424        let selection = RowSelection::from(vec![
425            RowSelector::skip(0),
426            RowSelector::select(2),
427            RowSelector::skip(0),
428            RowSelector::select(2),
429        ]);
430        assert_eq!(selection.selectors(), vec![RowSelector::select(4)]);
431
432        let selection = RowSelection::from(vec![
433            RowSelector::select(0),
434            RowSelector::skip(2),
435            RowSelector::select(0),
436            RowSelector::skip(2),
437        ]);
438        assert_eq!(selection.selectors(), vec![RowSelector::skip(4)]);
439    }
440
441    #[test]
442    fn test_trim() {
443        let selection = RowSelection::from(vec![
444            RowSelector::skip(34),
445            RowSelector::select(12),
446            RowSelector::skip(3),
447            RowSelector::select(35),
448        ]);
449
450        let expected = vec![
451            RowSelector::skip(34),
452            RowSelector::select(12),
453            RowSelector::skip(3),
454            RowSelector::select(35),
455        ];
456
457        assert_eq!(selection.trim().selectors(), expected);
458
459        let selection = RowSelection::from(vec![
460            RowSelector::skip(34),
461            RowSelector::select(12),
462            RowSelector::skip(3),
463        ]);
464
465        let expected = vec![RowSelector::skip(34), RowSelector::select(12)];
466
467        assert_eq!(selection.trim().selectors(), expected);
468    }
469}
parquet/arrow/arrow_reader/selection/selector.rs

parquet/arrow/arrow_reader/selection/
selector.rs
