1use super::{RowSelection, RowSelector};
26use crate::errors::ParquetError;
27use arrow_array::BooleanArray;
28use arrow_buffer::BooleanBuffer;
29use arrow_buffer::bit_chunk_iterator::BitChunks;
30use std::collections::VecDeque;
31use std::ops::Range;
32use std::sync::Arc;
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq)]
36pub enum RowSelectionPolicy {
37 Selectors,
39 Mask,
41 Auto {
43 threshold: usize,
45 },
46}
47
48impl Default for RowSelectionPolicy {
49 fn default() -> Self {
50 Self::Auto { threshold: 32 }
51 }
52}
53
54impl RowSelectionPolicy {
55 pub(crate) fn resolve(self, selection: Option<&RowSelection>) -> RowSelectionStrategy {
57 match self {
58 Self::Selectors => RowSelectionStrategy::Selectors,
59 Self::Mask => RowSelectionStrategy::Mask,
60 Self::Auto { threshold } => selection
61 .map(|selection| selection.auto_selection_strategy(threshold))
62 .unwrap_or(RowSelectionStrategy::Selectors),
63 }
64 }
65
66 pub(crate) fn apply(self, selection: RowSelection) -> RowSelection {
69 match (
70 self.resolve(Some(&selection)),
71 selection.as_mask().is_some(),
72 ) {
73 (RowSelectionStrategy::Mask, true) | (RowSelectionStrategy::Selectors, false) => {
74 selection
75 }
76 (RowSelectionStrategy::Mask, false) => selection.force_mask(),
77 (RowSelectionStrategy::Selectors, true) => {
78 RowSelection::from(Vec::<RowSelector>::from(selection))
79 }
80 }
81 }
82}
83
84#[derive(Clone, Copy, Debug, Eq, PartialEq)]
89pub(crate) enum RowSelectionStrategy {
90 Selectors,
92 Mask,
94}
95
96impl RowSelectionStrategy {
97 pub(crate) fn build_cursor(
100 self,
101 selection: RowSelection,
102 loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
103 ) -> RowSelectionCursor {
104 match self {
105 Self::Mask => RowSelectionCursor::new_mask_from_buffer(
106 selection.into_boolean_buffer(),
107 loaded_row_ranges,
108 ),
109 Self::Selectors => RowSelectionCursor::new_selectors(selection.into()),
110 }
111 }
112}
113
114#[derive(Debug)]
120pub enum RowSelectionCursor {
121 All,
123 Mask(MaskCursor),
125 Selectors(SelectorsCursor),
127}
128
129impl RowSelectionCursor {
130 pub(crate) fn new_mask_from_buffer(
132 mask: BooleanBuffer,
133 loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
134 ) -> Self {
135 debug_assert!(
136 mask.is_empty() || mask.value(mask.len() - 1),
137 "Mask selections must not end with a skip"
138 );
139 Self::Mask(MaskCursor {
140 mask,
141 position: 0,
142 loaded_row_ranges,
143 })
144 }
145
146 pub(crate) fn new_selectors(selectors: Vec<RowSelector>) -> Self {
148 Self::Selectors(SelectorsCursor {
149 selectors: selectors.into(),
150 position: 0,
151 })
152 }
153
154 pub(crate) fn new_all() -> Self {
156 Self::All
157 }
158}
159
160#[derive(Debug)]
165pub struct SelectorsCursor {
166 selectors: VecDeque<RowSelector>,
167 position: usize,
169}
170
171impl SelectorsCursor {
172 pub fn is_empty(&self) -> bool {
174 self.selectors.is_empty()
175 }
176
177 pub(crate) fn next_selector(&mut self) -> RowSelector {
179 let selector = self.selectors.pop_front().unwrap();
180 self.position += selector.row_count;
181 selector
182 }
183
184 pub(crate) fn return_selector(&mut self, selector: RowSelector) {
186 self.position = self.position.saturating_sub(selector.row_count);
187 self.selectors.push_front(selector);
188 }
189}
190
191#[derive(Debug)]
216pub struct MaskCursor {
217 mask: BooleanBuffer,
218 position: usize,
220 loaded_row_ranges: Option<Arc<LoadedRowRanges>>,
222}
223
224impl MaskCursor {
225 pub fn is_empty(&self) -> bool {
227 self.position >= self.mask.len()
228 }
229
230 pub fn next_mask_chunk(&mut self, batch_size: usize) -> Option<MaskChunk> {
232 if self.is_empty() {
233 return None;
234 }
235
236 Some(self.next_mask_chunk_non_empty(batch_size))
237 }
238
239 fn next_mask_chunk_non_empty(&mut self, batch_size: usize) -> MaskChunk {
241 debug_assert!(!self.is_empty());
242 let mask_start = self.next_selected_row().unwrap_or(self.mask.len());
243 debug_assert!(
244 mask_start < self.mask.len(),
245 "ReadPlan must remove trailing skips from Mask selections"
246 );
247 self.scan_chunk(mask_start, self.mask.len(), batch_size)
248 }
249
250 fn next_selected_row(&self) -> Option<usize> {
252 let chunks = BitChunks::new(
253 self.mask.values(),
254 self.mask.offset() + self.position,
255 self.mask.len() - self.position,
256 );
257 chunks.iter_padded().enumerate().find_map(|(idx, word)| {
258 (word != 0).then(|| self.position + idx * 64 + word.trailing_zeros() as usize)
259 })
260 }
261
262 #[inline(never)]
267 fn scan_chunk(&mut self, mask_start: usize, end: usize, batch_size: usize) -> MaskChunk {
268 let chunks = BitChunks::new(
269 self.mask.values(),
270 self.mask.offset() + mask_start,
271 end - mask_start,
272 );
273 let mut selected_rows = 0;
274 let mut chunk_end = mask_start;
275
276 if batch_size > 0 {
277 for (idx, mut word) in chunks.iter_padded().enumerate() {
278 let count = word.count_ones() as usize;
279 if count == 0 {
280 continue;
281 }
282
283 let remaining = batch_size - selected_rows;
284 let last_bit = if count <= remaining {
285 selected_rows += count;
286 63 - word.leading_zeros() as usize
287 } else {
288 for _ in 1..remaining {
290 word &= word - 1;
291 }
292 selected_rows = batch_size;
293 word.trailing_zeros() as usize
294 };
295 chunk_end = mask_start + idx * 64 + last_bit + 1;
296 if selected_rows == batch_size {
297 break;
298 }
299 }
300 }
301
302 let initial_skip = mask_start - self.position;
303 self.position = chunk_end;
306 MaskChunk {
307 initial_skip,
308 chunk_rows: chunk_end - mask_start,
309 selected_rows,
310 mask_start,
311 }
312 }
313
314 pub(crate) fn next_chunk(&mut self, batch_size: usize) -> Result<MaskChunk, ParquetError> {
323 debug_assert!(batch_size > 0);
324 debug_assert!(!self.is_empty());
325
326 if self.loaded_row_ranges.is_none() {
327 return Ok(self.next_mask_chunk_non_empty(batch_size));
328 }
329
330 let mask_start = self.next_selected_row().ok_or_else(|| {
331 ParquetError::General(
332 "Internal Error: Mask cursor reached the end without finding a selected row; \
333 ReadPlan must remove trailing skips"
334 .to_string(),
335 )
336 })?;
337
338 let loaded_range_end = self
339 .loaded_row_ranges
340 .as_ref()
341 .and_then(|ranges| ranges.end_containing(mask_start))
342 .ok_or_else(|| {
343 ParquetError::General(format!(
344 "Internal Error: selected row {mask_start} has no loaded page range"
345 ))
346 })?;
347
348 Ok(self.scan_chunk(
349 mask_start,
350 loaded_range_end.min(self.mask.len()),
351 batch_size,
352 ))
353 }
354
355 pub fn mask_values_for(&self, chunk: &MaskChunk) -> Result<BooleanArray, ParquetError> {
357 if chunk.mask_start.saturating_add(chunk.chunk_rows) > self.mask.len() {
358 return Err(ParquetError::General(
359 "Internal Error: MaskChunk exceeds mask length".to_string(),
360 ));
361 }
362 Ok(BooleanArray::from(
363 self.mask.slice(chunk.mask_start, chunk.chunk_rows),
364 ))
365 }
366}
367
368#[derive(Debug)]
370pub struct MaskChunk {
371 pub initial_skip: usize,
373 pub chunk_rows: usize,
375 pub selected_rows: usize,
377 pub mask_start: usize,
379}
380
381#[derive(Clone, Debug)]
383pub(crate) struct LoadedRowRanges(Vec<Range<usize>>);
384
385impl LoadedRowRanges {
386 pub(crate) fn from_selection(selection: RowSelection) -> Self {
387 let selectors: Vec<RowSelector> = selection.into();
388 let mut position = 0;
389 let ranges = selectors
390 .into_iter()
391 .filter_map(|selector| {
392 let start = position;
393 position += selector.row_count;
394 (!selector.skip).then_some(start..position)
395 })
396 .collect();
397 Self(ranges)
398 }
399
400 fn end_containing(&self, row: usize) -> Option<usize> {
401 let idx = self.0.partition_point(|range| range.end <= row);
402 self.0
403 .get(idx)
404 .filter(|range| range.start <= row)
405 .map(|range| range.end)
406 }
407
408 #[cfg(test)]
409 pub(crate) fn ranges(&self) -> &[Range<usize>] {
410 &self.0
411 }
412}
413
414#[cfg(test)]
415mod tests {
416 use super::*;
417
418 fn assert_mask_chunks(mask: BooleanBuffer, loaded: Option<LoadedRowRanges>, batch_size: usize) {
422 let selected: Vec<_> = (0..mask.len()).filter(|&row| mask.value(row)).collect();
423 let loaded = loaded.map(Arc::new);
424 let selection = RowSelection::from_boolean_buffer(mask).trim();
425 let RowSelectionCursor::Mask(mut cursor) =
426 RowSelectionStrategy::Mask.build_cursor(selection, loaded.clone())
427 else {
428 unreachable!()
429 };
430 let mut next = 0;
431 let mut position = 0;
432 while next < selected.len() {
433 let first = selected[next];
434 let end = loaded
435 .as_ref()
436 .map_or(usize::MAX, |ranges| ranges.end_containing(first).unwrap());
437 let count = selected[next..]
438 .iter()
439 .take(batch_size)
440 .take_while(|&&row| row < end)
441 .count();
442 let last = selected[next + count - 1];
443 let chunk = cursor.next_chunk(batch_size).unwrap();
444 assert_eq!(chunk.initial_skip, first - position);
445 assert_eq!(chunk.mask_start, first);
446 assert_eq!(chunk.chunk_rows, last + 1 - first);
447 assert_eq!(chunk.selected_rows, count);
448 assert_eq!(cursor.position, last + 1);
449 let values = cursor.mask_values_for(&chunk).unwrap();
450 assert_eq!(values.true_count(), count);
451 assert!(values.value(0));
452 assert!(values.value(values.len() - 1));
453 next += count;
454 position = last + 1;
455 }
456 assert!(cursor.is_empty());
457 assert!(cursor.next_mask_chunk(batch_size).is_none());
458 }
459
460 #[test]
461 fn test_mask_chunks_sparse_aligned_words() {
462 let mask = BooleanBuffer::from_iter((0..256).map(|row| row == 0 || row == 255));
465 assert_mask_chunks(mask, None, 1);
466 }
467
468 #[test]
469 fn test_mask_chunks_word_boundaries_and_offsets() {
470 use rand::{RngExt, SeedableRng, rngs::StdRng};
471
472 for len in [0, 1, 7, 8, 63, 64, 65, 127, 128, 129, 511, 2049] {
473 for offset in [0, 1, 3, 7, 8, 31, 63, 64, 65] {
474 for pattern in 0..7 {
475 let mut rng = StdRng::seed_from_u64(42);
476 let mask = BooleanBuffer::from_iter((0..offset + len + 64).map(|row| {
477 if !(offset..offset + len).contains(&row) {
478 return true;
481 }
482 let row = row - offset;
483 match pattern {
484 0 => false,
485 1 => true,
486 2 => row % 2 == 0,
487 3 => row % 129 == 128,
488 4 => rng.random_bool(0.1),
489 5 => rng.random_bool(0.5),
490 _ => rng.random_bool(0.9),
491 }
492 }))
493 .slice(offset, len);
494 for batch_size in [1, 2, 3, 63, 64, 65, 1024] {
495 assert_mask_chunks(mask.clone(), None, batch_size);
496 }
497 }
498 }
499 }
500 }
501
502 #[test]
503 fn test_mask_chunks_loaded_ranges_word_boundaries() {
504 for page_rows in [1, 31, 63, 64, 65, 127, 257] {
505 let len = page_rows * 7 + 13;
506 let ranges: Vec<_> = (0..len)
507 .step_by(page_rows)
508 .enumerate()
509 .filter(|(page, _)| page % 3 != 1)
510 .map(|(_, start)| start..(start + page_rows).min(len))
511 .collect();
512 let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
513 ranges.iter().cloned(),
514 len,
515 ));
516 for offset in [0, 1, 3, 7, 31, 63, 64, 65] {
517 for pattern in 0..4 {
518 let mask = BooleanBuffer::from_iter((0..offset + len + 64).map(|row| {
519 if !(offset..offset + len).contains(&row) {
520 return true;
521 }
522 let row = row - offset;
523 ranges.iter().any(|range| range.contains(&row))
524 && match pattern {
525 0 => true,
526 1 => row % 2 == 0,
527 2 => row % page_rows == 0,
529 _ => {
532 row % page_rows == page_rows - 1
533 || row % 64 == 0
534 || row % 64 == 63
535 }
536 }
537 }))
538 .slice(offset, len);
539 for batch_size in [1, 2, 3, 63, 64, 65, 1024] {
540 assert_mask_chunks(mask.clone(), Some(loaded.clone()), batch_size);
541 }
542 }
543 }
544 }
545 }
546
547 #[test]
548 fn test_mask_chunk_does_not_read_selected_rows_outside_loaded_range() {
549 let mut cursor = MaskCursor {
550 mask: BooleanBuffer::new_set(130),
551 position: 0,
552 loaded_row_ranges: Some(Arc::new(LoadedRowRanges(vec![0..65, 66..130]))),
553 };
554 let chunk = cursor.next_chunk(1024).unwrap();
555 assert_eq!(chunk.chunk_rows, 65);
556 assert_eq!(chunk.selected_rows, 65);
557 assert_eq!(cursor.position, 65);
558 let err = cursor.next_chunk(1024).unwrap_err();
559 assert!(
560 err.to_string()
561 .contains("selected row 65 has no loaded page range")
562 );
563 assert_eq!(cursor.position, 65);
564 }
565
566 #[test]
567 fn selection_respects_policy() {
568 let mask = BooleanBuffer::from(vec![true, true, false, false, true, true, false, false]);
569 for (policy, expect_mask) in [
570 (RowSelectionPolicy::Mask, true),
571 (RowSelectionPolicy::Selectors, false),
572 (RowSelectionPolicy::Auto { threshold: 2 }, false),
573 (RowSelectionPolicy::Auto { threshold: 3 }, true),
574 ] {
575 for selection in [
576 RowSelection::from_boolean_buffer(mask.clone()),
577 RowSelection::from(vec![
578 RowSelector::select(2),
579 RowSelector::skip(2),
580 RowSelector::select(2),
581 RowSelector::skip(2),
582 ]),
583 ] {
584 let selection = policy.apply(selection);
585 assert_eq!(selection.as_mask().is_some(), expect_mask, "{policy:?}");
586 assert_eq!(selection.into_boolean_buffer(), mask);
587 }
588 }
589 }
590
591 #[test]
592 fn adaptive_selection_uses_selectors_for_long_runs() {
593 let mask = BooleanBuffer::from_iter((0..1_024).map(|idx| (256..768).contains(&idx)));
594 let selection =
595 RowSelectionPolicy::default().apply(RowSelection::from_boolean_buffer(mask));
596 assert!(selection.as_mask().is_none());
597 }
598
599 #[test]
600 fn adaptive_selection_keeps_fragmented_masks() {
601 let mask = BooleanBuffer::from_iter((0..1_024).map(|idx| idx % 2 == 0));
602 let selection =
603 RowSelectionPolicy::default().apply(RowSelection::from_boolean_buffer(mask));
604 assert!(selection.as_mask().is_some());
605 }
606
607 #[test]
608 fn test_loaded_mask_chunk_stops_at_trimmed_mask_end() {
609 let loaded = LoadedRowRanges::from_selection(RowSelection::from_consecutive_ranges(
610 std::iter::once(0..5),
611 10,
612 ));
613 let RowSelectionCursor::Mask(mut cursor) = RowSelectionStrategy::Mask.build_cursor(
614 RowSelection::from(vec![RowSelector::select(1)]),
615 Some(loaded.into()),
616 ) else {
617 unreachable!()
618 };
619
620 let chunk = cursor.next_chunk(10).unwrap();
621 assert_eq!(chunk.chunk_rows, 1);
622 assert!(cursor.is_empty());
623 }
624
625 #[test]
626 fn test_next_mask_chunk_until_cursor_is_empty() {
627 let RowSelectionCursor::Mask(mut cursor) = RowSelectionStrategy::Mask.build_cursor(
628 RowSelection::from(vec![
629 RowSelector::skip(2),
630 RowSelector::select(2),
631 RowSelector::skip(1),
632 RowSelector::select(1),
633 ]),
634 None,
635 ) else {
636 unreachable!()
637 };
638
639 let first = cursor.next_mask_chunk(2).unwrap();
640 assert_eq!(first.initial_skip, 2);
641 assert_eq!(first.chunk_rows, 2);
642 assert_eq!(first.selected_rows, 2);
643
644 let second = cursor.next_mask_chunk(2).unwrap();
645 assert_eq!(second.initial_skip, 1);
646 assert_eq!(second.chunk_rows, 1);
647 assert_eq!(second.selected_rows, 1);
648
649 assert!(cursor.next_mask_chunk(2).is_none());
650 }
651}