feat: add vortex-mask crate (#2019)

I decided to stick with the vortex-blah naming scheme.. instead of my
much better name "inthesea"

Author: gatesn

Cargo.lock

@@ -17,6 +17,7 @@ members = [
     "vortex-io",
     "vortex-ipc",
     "vortex-layout",
+    "vortex-mask",
     "vortex-proto",
     "vortex-sampling-compressor",
     "vortex-scalar",

Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "vortex-mask"
+description = "Vortex Mask - sorted, unique, positive integers"
+version.workspace = true
+homepage.workspace = true
+repository.workspace = true
+authors.workspace = true
+license.workspace = true
+keywords.workspace = true
+include.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+readme.workspace = true
+categories.workspace = true
+
+[dependencies]
+arrow-buffer = { workspace = true }
+vortex-error = { workspace = true }
+
+[lints]
+workspace = true

vortex-mask/Cargo.toml

@@ -0,0 +1,40 @@
+use std::ops::BitAnd;
+
+use vortex_error::vortex_panic;
+
+use crate::Mask;
+
+impl BitAnd for &Mask {
+    type Output = Mask;
+
+    fn bitand(self, rhs: Self) -> Self::Output {
+        if self.len() != rhs.len() {
+            vortex_panic!("FilterMasks must have the same length");
+        }
+        if self.true_count() == 0 || rhs.true_count() == 0 {
+            return Mask::new_false(self.len());
+        }
+        if self.true_count() == self.len() {
+            return rhs.clone();
+        }
+        if rhs.true_count() == self.len() {
+            return self.clone();
+        }
+
+        if let (Some(lhs), Some(rhs)) = (self.0.buffer.get(), rhs.0.buffer.get()) {
+            return Mask::from_buffer(lhs & rhs);
+        }
+
+        if let (Some(lhs), Some(rhs)) = (self.0.indices.get(), rhs.0.indices.get()) {
+            // TODO(ngates): this may only make sense for sparse indices.
+            return Mask::from_intersection_indices(
+                self.len(),
+                lhs.iter().copied(),
+                rhs.iter().copied(),
+            );
+        }
+
+        // TODO(ngates): we could perform a more efficient bitandion for slices.
+        Mask::from_buffer(self.boolean_buffer() & rhs.boolean_buffer())
+    }
+}

vortex-mask/src/bitand.rs

@@ -0,0 +1,66 @@
+use crate::Mask;
+
+impl PartialEq for Mask {
+    fn eq(&self, other: &Self) -> bool {
+        if self.len() != other.len() {
+            return false;
+        }
+        if self.true_count() != other.true_count() {
+            return false;
+        }
+
+        // Since the true counts are the same, a full or empty mask is equal to the other mask.
+        if self.true_count() == 0 || self.true_count() == self.len() {
+            return true;
+        }
+
+        // Compare the buffer if both masks are non-empty.
+        if let (Some(buffer), Some(other)) = (self.0.buffer.get(), other.0.buffer.get()) {
+            return buffer == other;
+        }
+
+        // Compare the indices if both masks are non-empty.
+        if let (Some(indices), Some(other)) = (self.0.indices.get(), other.0.indices.get()) {
+            return indices == other;
+        }
+
+        // Compare the slices if both masks are non-empty.
+        if let (Some(slices), Some(other)) = (self.0.slices.get(), other.0.slices.get()) {
+            return slices == other;
+        }
+
+        // Otherwise, we fall back to comparison based on sparsity.
+        // We could go further an exhaustively check whose OnceLocks are initialized, but that's
+        // probably not worth the effort.
+        self.boolean_buffer() == other.boolean_buffer()
+    }
+}
+
+impl Eq for Mask {}
+
+#[cfg(test)]
+mod test {
+    use arrow_buffer::BooleanBuffer;
+
+    use crate::Mask;
+
+    #[test]
+    fn filter_mask_eq() {
+        assert_eq!(
+            Mask::new_true(5),
+            Mask::from_buffer(BooleanBuffer::new_set(5))
+        );
+        assert_eq!(
+            Mask::new_false(5),
+            Mask::from_buffer(BooleanBuffer::new_unset(5))
+        );
+        assert_eq!(
+            Mask::from_indices(5, vec![0, 2, 3]),
+            Mask::from_slices(5, vec![(0, 1), (2, 4)])
+        );
+        assert_eq!(
+            Mask::from_indices(5, vec![0, 2, 3]),
+            Mask::from_buffer(BooleanBuffer::from_iter([true, false, true, true, false]))
+        );
+    }
+}

vortex-mask/src/eq.rs

@@ -0,0 +1,86 @@
+use crate::Mask;
+
+impl Mask {
+    /// take the intersection of the `mask` with the set of true values in `self`.
+    ///
+    /// We are more interested in low selectivity `self` (as indices) with a boolean buffer mask,
+    /// so we don't optimize for other cases, yet.
+    pub fn intersect_by_rank(&self, mask: &Mask) -> Mask {
+        assert_eq!(self.true_count(), mask.len());
+
+        if mask.true_count() == mask.len() {
+            return self.clone();
+        }
+
+        if mask.true_count() == 0 {
+            return Self::new_false(self.len());
+        }
+
+        // TODO(joe): support other fast paths, not converting self & mask into indices,
+        // however indices are better for sparse masks, so this is the common case for now.
+        let indices = self.0.indices();
+        Self::from_indices(
+            self.len(),
+            mask.indices()
+                .iter()
+                .map(|idx|
+                    // This is verified as safe because we know that the indices are less than the
+                    // mask.len() and we known mask.len() <= self.len(),
+                    // implied by `self.true_count() == mask.len()`.
+                    unsafe{*indices.get_unchecked(*idx)})
+                .collect(),
+        )
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use arrow_buffer::BooleanBuffer;
+
+    use crate::Mask;
+
+    #[test]
+    fn mask_bitand_all_as_bit_and() {
+        let this = Mask::from_buffer(BooleanBuffer::from_iter(vec![true, true, true, true, true]));
+        let mask = Mask::from_buffer(BooleanBuffer::from_iter(vec![
+            false, true, false, true, true,
+        ]));
+        assert_eq!(
+            this.intersect_by_rank(&mask),
+            Mask::from_indices(5, vec![1, 3, 4])
+        );
+    }
+
+    #[test]
+    fn mask_bitand_all_true() {
+        let this = Mask::from_buffer(BooleanBuffer::from_iter(vec![
+            false, false, true, true, true,
+        ]));
+        let mask = Mask::from_buffer(BooleanBuffer::from_iter(vec![true, true, true]));
+        assert_eq!(
+            this.intersect_by_rank(&mask),
+            Mask::from_indices(5, vec![2, 3, 4])
+        );
+    }
+
+    #[test]
+    fn mask_bitand_true() {
+        let this = Mask::from_buffer(BooleanBuffer::from_iter(vec![
+            true, false, false, true, true,
+        ]));
+        let mask = Mask::from_buffer(BooleanBuffer::from_iter(vec![true, false, true]));
+        assert_eq!(
+            this.intersect_by_rank(&mask),
+            Mask::from_indices(5, vec![0, 4])
+        );
+    }
+
+    #[test]
+    fn mask_bitand_false() {
+        let this = Mask::from_buffer(BooleanBuffer::from_iter(vec![
+            true, false, false, true, true,
+        ]));
+        let mask = Mask::from_buffer(BooleanBuffer::from_iter(vec![false, false, false]));
+        assert_eq!(this.intersect_by_rank(&mask), Mask::from_indices(5, vec![]));
+    }
+}