vortex_scan/
row_mask.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
use std::cmp::{max, min};
use std::fmt::{Display, Formatter};
use std::ops::RangeBounds;

use vortex_array::array::BooleanBuffer;
use vortex_array::compute::{filter, slice, try_cast};
use vortex_array::{Array, IntoArrayVariant};
use vortex_dtype::Nullability::NonNullable;
use vortex_dtype::{DType, PType};
use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult};
use vortex_mask::Mask;

/// A RowMask captures a set of selected rows within a range.
///
/// The range itself can be [`u64`], but the length of the range must fit into a [`usize`], this
/// allows us to use a `usize` filter mask within a much larger file.
#[derive(Debug, Clone)]
pub struct RowMask {
    mask: Mask,
    begin: u64,
    end: u64,
}

// We don't want to implement full logical equality, this naive equality is sufficient for tests.
#[cfg(test)]
impl PartialEq for RowMask {
    fn eq(&self, other: &Self) -> bool {
        self.begin == other.begin && self.end == other.end && self.mask == other.mask
    }
}

impl Display for RowMask {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "RowSelector [{}..{}]", self.begin, self.end)
    }
}

impl RowMask {
    /// Define a new [`RowMask`] with the given mask and offset into the file.
    pub fn new(mask: Mask, begin: u64) -> Self {
        let end = begin + (mask.len() as u64);
        Self { mask, begin, end }
    }

    /// Construct a RowMask which is valid in the given range.
    ///
    /// ## Panics
    ///
    /// If the size of the range is too large to fit into a usize.
    pub fn new_valid_between(begin: u64, end: u64) -> Self {
        let length =
            usize::try_from(end - begin).vortex_expect("Range length does not fit into a usize");
        RowMask::new(Mask::from(BooleanBuffer::new_set(length)), begin)
    }

    /// Construct a RowMask which is invalid everywhere in the given range.
    pub fn new_invalid_between(begin: u64, end: u64) -> Self {
        let length =
            usize::try_from(end - begin).vortex_expect("Range length does not fit into a usize");
        RowMask::new(Mask::from(BooleanBuffer::new_unset(length)), begin)
    }

    /// Creates a RowMask from an array, only supported boolean and integer types.
    pub fn from_array(array: &Array, begin: u64, end: u64) -> VortexResult<Self> {
        if array.dtype().is_int() {
            Self::from_index_array(array, begin, end)
        } else if array.dtype().is_boolean() {
            Self::from_mask_array(array, begin)
        } else {
            vortex_bail!(
                "RowMask can only be created from integer or boolean arrays, got {} instead.",
                array.dtype()
            );
        }
    }

    /// Construct a RowMask from a Boolean typed array.
    ///
    /// True-valued positions are kept by the returned mask.
    fn from_mask_array(array: &Array, begin: u64) -> VortexResult<Self> {
        Ok(Self::new(array.logical_validity()?, begin))
    }

    /// Construct a RowMask from an integral array.
    ///
    /// The array values are interpreted as indices and those indices are kept by the returned mask.
    #[allow(clippy::cast_possible_truncation)]
    fn from_index_array(array: &Array, begin: u64, end: u64) -> VortexResult<Self> {
        let length = usize::try_from(end - begin)
            .map_err(|_| vortex_err!("Range length does not fit into a usize"))?;

        let indices =
            try_cast(array, &DType::Primitive(PType::U64, NonNullable))?.into_primitive()?;

        let mask = Mask::from_indices(
            length,
            indices
                .as_slice::<u64>()
                .iter()
                .map(|i| *i as usize)
                .collect(),
        );

        Ok(RowMask::new(mask, begin))
    }

    /// Whether the mask is disjoint with the given range.
    ///
    /// This function may return false negatives, but never false positives.
    ///
    /// TODO(ngates): improve this function to take into account the [`Mask`].
    pub fn is_disjoint(&self, range: impl RangeBounds<u64>) -> bool {
        use std::ops::Bound;

        // Get the start bound of the input range
        let start = match range.start_bound() {
            Bound::Included(&n) => n,
            Bound::Excluded(&n) => n + 1,
            Bound::Unbounded => 0,
        };

        // Get the end bound of the input range
        let end = match range.end_bound() {
            Bound::Included(&n) => n + 1,
            Bound::Excluded(&n) => n,
            Bound::Unbounded => u64::MAX,
        };

        // Two ranges are disjoint if one ends before the other begins
        self.end <= start || end <= self.begin
    }

    /// The beginning of the masked range.
    #[inline]
    pub fn begin(&self) -> u64 {
        self.begin
    }

    /// The end of the masked range.
    #[inline]
    pub fn end(&self) -> u64 {
        self.end
    }

    /// The length of the mask is the number of possible rows between the `begin` and `end`,
    /// regardless of how many appear in the mask. For the number of masked rows, see `true_count`.
    #[inline]
    // There is good definition of is_empty, does it mean len == 0 or true_count == 0?
    #[allow(clippy::len_without_is_empty)]
    pub fn len(&self) -> usize {
        self.mask.len()
    }

    /// Returns the [`Mask`] whose true values are relative to the range of this `RowMask`.
    pub fn filter_mask(&self) -> &Mask {
        &self.mask
    }

    /// Limit mask to `[begin..end)` range
    pub fn slice(&self, begin: u64, end: u64) -> VortexResult<Self> {
        let range_begin = max(self.begin, begin);
        let range_end = min(self.end, end);
        Ok(RowMask::new(
            if range_begin == self.begin && range_end == self.end {
                self.mask.clone()
            } else {
                self.mask.slice(
                    usize::try_from(range_begin - self.begin)
                        .vortex_expect("we know this must fit into usize"),
                    usize::try_from(range_end - range_begin)
                        .vortex_expect("we know this must fit into usize"),
                )
            },
            range_begin,
        ))
    }

    /// Filter array with this `RowMask`.
    ///
    /// This function assumes that Array is no longer than the mask length and that the mask starts on same offset as the array,
    /// i.e. the beginning of the array corresponds to the beginning of the mask with begin = 0
    pub fn filter_array(&self, array: impl AsRef<Array>) -> VortexResult<Option<Array>> {
        let true_count = self.mask.true_count();
        if true_count == 0 {
            return Ok(None);
        }

        let array = array.as_ref();

        let sliced = if self.len() == array.len() {
            array
        } else {
            // TODO(ngates): I thought the point was the array only covers the valid row range of
            //  the mask?
            // FIXME(ngates): this is made more obvious by the unsafe u64 cast.
            &slice(
                array,
                usize::try_from(self.begin).vortex_expect("TODO(ngates): fix this bad cast"),
                usize::try_from(self.end).vortex_expect("TODO(ngates): fix this bad cast"),
            )?
        };

        if true_count == sliced.len() {
            return Ok(Some(sliced.clone()));
        }

        filter(sliced, &self.mask).map(Some)
    }

    /// Shift the [`RowMask`] down by the given offset.
    pub fn shift(self, offset: u64) -> VortexResult<RowMask> {
        let valid_shift = self.begin >= offset;
        if !valid_shift {
            vortex_bail!(
                "Can shift RowMask by at most {}, tried to shift by {offset}",
                self.begin
            )
        }
        Ok(RowMask::new(self.mask, self.begin - offset))
    }

    /// The number of masked rows within the range.
    pub fn true_count(&self) -> usize {
        self.mask.true_count()
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;
    use vortex_array::array::PrimitiveArray;
    use vortex_array::validity::Validity;
    use vortex_array::{IntoArray, IntoArrayVariant};
    use vortex_buffer::{buffer, Buffer};
    use vortex_error::VortexUnwrap;
    use vortex_mask::Mask;

    use super::*;

    #[rstest]
    #[case(
        RowMask::new(Mask::from_iter([true, true, true, false, false, false, false, false, true, true]), 0), (0, 1),
        RowMask::new(Mask::from_iter([true]), 0))]
    #[case(
        RowMask::new(Mask::from_iter([false, false, false, false, false, true, true, true, true, true]), 0), (2, 5),
        RowMask::new(Mask::from_iter([false, false, false]), 2)
    )]
    #[case(
        RowMask::new(Mask::from_iter([true, true, true, true, false, false, false, false, false, false]), 0), (2, 5),
        RowMask::new(Mask::from_iter([true, true, false]), 2)
    )]
    #[case(
        RowMask::new(Mask::from_iter([true, true, true, false, false, true, true, false, false, false]), 0), (2, 6),
        RowMask::new(Mask::from_iter([true, false, false, true]), 2))]
    #[case(
        RowMask::new(Mask::from_iter([false, false, false, false, false, true, true, true, true, true]), 0), (7, 11),
        RowMask::new(Mask::from_iter([true, true, true]), 7))]
    #[case(
        RowMask::new(Mask::from_iter([false, true, true, true, true, true]), 3), (0, 5),
        RowMask::new(Mask::from_iter([false, true]), 3))]
    #[cfg_attr(miri, ignore)]
    fn slice(#[case] first: RowMask, #[case] range: (u64, u64), #[case] expected: RowMask) {
        assert_eq!(first.slice(range.0, range.1).vortex_unwrap(), expected);
    }

    #[test]
    #[should_panic]
    #[cfg_attr(miri, ignore)]
    fn shift_invalid() {
        RowMask::new(Mask::from_iter([true, true, true, true, true]), 5)
            .shift(7)
            .unwrap();
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn shift() {
        assert_eq!(
            RowMask::new(Mask::from_iter([true, true, true, true, true]), 5)
                .shift(5)
                .unwrap(),
            RowMask::new(Mask::from_iter([true, true, true, true, true]), 0)
        );
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn filter_array() {
        let mask = RowMask::new(
            Mask::from_iter([
                false, false, false, false, false, true, true, true, true, true,
            ]),
            0,
        );
        let array = Buffer::from_iter(0..20).into_array();
        let filtered = mask.filter_array(array).unwrap().unwrap();
        assert_eq!(
            filtered.into_primitive().unwrap().as_slice::<i32>(),
            (5..10).collect::<Vec<_>>()
        );
    }

    #[test]
    #[should_panic]
    fn test_row_mask_type_validation() {
        let array = PrimitiveArray::new(buffer![1.0, 2.0], Validity::AllInvalid).into_array();
        RowMask::from_array(&array, 0, 2).unwrap();
    }
}