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

parquet/arrow/arrow_reader/selection/
selector.rs
