polars_mem_engine/scan_predicate/
skip_files_mask.rs

1use arrow::bitmap::Bitmap;
2
3#[derive(Debug, Clone)]
4pub enum SkipFilesMask {
5    Exclusion(Bitmap),
6    Inclusion(Bitmap),
7}
8
9impl SkipFilesMask {
10    pub fn is_empty(&self) -> bool {
11        match self {
12            Self::Exclusion(mask) => mask.is_empty(),
13            Self::Inclusion(mask) => mask.is_empty(),
14        }
15    }
16
17    pub fn len(&self) -> usize {
18        match self {
19            Self::Exclusion(mask) => mask.len(),
20            Self::Inclusion(mask) => mask.len(),
21        }
22    }
23
24    pub fn is_skipped_file(&self, index: usize) -> bool {
25        match self {
26            Self::Exclusion(mask) => mask.get_bit(index),
27            Self::Inclusion(mask) => !mask.get_bit(index),
28        }
29    }
30
31    pub fn num_skipped_files(&self) -> usize {
32        match self {
33            Self::Exclusion(mask) => mask.set_bits(),
34            Self::Inclusion(mask) => mask.unset_bits(),
35        }
36    }
37
38    pub fn leading_skipped_files(&self) -> usize {
39        match self {
40            Self::Exclusion(mask) => mask.leading_ones(),
41            Self::Inclusion(mask) => mask.leading_zeros(),
42        }
43    }
44
45    pub fn trailing_skipped_files(&self) -> usize {
46        match self {
47            Self::Exclusion(mask) => mask.trailing_ones(),
48            Self::Inclusion(mask) => mask.trailing_zeros(),
49        }
50    }
51
52    pub fn sliced(self, offset: usize, len: usize) -> Self {
53        match self {
54            Self::Exclusion(mask) => Self::Exclusion(mask.sliced(offset, len)),
55            Self::Inclusion(mask) => Self::Inclusion(mask.sliced(offset, len)),
56        }
57    }
58
59    pub fn non_skipped_files_idx_iter(&self) -> impl Iterator<Item = usize> + Clone {
60        let range_end = self.len() - self.trailing_skipped_files();
61        let range_start = if range_end == 0 {
62            0
63        } else {
64            self.leading_skipped_files()
65        };
66
67        (range_start..range_end).filter(|i| !self.is_skipped_file(*i))
68    }
69}