parquet/arrow/arrow_reader/selection/
selector.rs

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

parquet/arrow/arrow_reader/selection/
selector.rs
