MaskedArray + masked compute + is_null/is_not_null

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/arrays/masked/vtable/mod.rs

@@ -4,6 +4,7 @@
 mod array;
 mod canonical;
 mod operations;
+mod operator;
 mod serde;
 mod validity;
 
@@ -28,7 +29,7 @@ impl VTable for MaskedVTable {
     type ComputeVTable = NotSupported;
     type EncodeVTable = NotSupported;
     type SerdeVTable = Self;
-    type OperatorVTable = NotSupported;
+    type OperatorVTable = Self;
 
     fn id(_encoding: &Self::Encoding) -> EncodingId {
         EncodingId::new_ref("vortex.masked")

vortex-array/src/arrays/masked/vtable/operator.rs

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_compute::mask::MaskValidity;
+use vortex_error::VortexResult;
+
+use crate::ArrayRef;
+use crate::arrays::{MaskedArray, MaskedVTable};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::vtable::OperatorVTable;
+
+impl OperatorVTable<MaskedVTable> for MaskedVTable {
+    fn bind(
+        array: &MaskedArray,
+        selection: Option<&ArrayRef>,
+        ctx: &mut dyn BindCtx,
+    ) -> VortexResult<BatchKernelRef> {
+        // A masked array performs the intersection if the mask validity with the child validity.
+        let mask = ctx.bind_validity(&array.validity, array.len(), selection)?;
+        let child = ctx.bind(&array.child, selection)?;
+
+        Ok(kernel(move || {
+            let mask = mask.execute()?;
+            let child = child.execute()?;
+            Ok(MaskValidity::mask_validity(child, &mask))
+        }))
+    }
+}

vortex-array/src/compute/arrays/is_not_null.rs

@@ -0,0 +1,136 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::hash::Hasher;
+
+use vortex_dtype::DType;
+use vortex_dtype::Nullability::NonNullable;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::{BoolVector, VectorOps};
+
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::stats::{ArrayStats, StatsSetRef};
+use crate::vtable::{ArrayVTable, NotSupported, OperatorVTable, VTable, VisitorVTable};
+use crate::{
+    ArrayBufferVisitor, ArrayChildVisitor, ArrayEq, ArrayHash, ArrayRef, EncodingId, EncodingRef,
+    Precision, vtable,
+};
+
+vtable!(IsNotNull);
+
+#[derive(Debug, Clone)]
+pub struct IsNotNullArray {
+    child: ArrayRef,
+    stats: ArrayStats,
+}
+
+impl IsNotNullArray {
+    /// Create a new is_not_null array.
+    pub fn new(child: ArrayRef) -> Self {
+        Self {
+            child,
+            stats: ArrayStats::default(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct IsNotNullEncoding;
+
+impl VTable for IsNotNullVTable {
+    type Array = IsNotNullArray;
+    type Encoding = IsNotNullEncoding;
+    type ArrayVTable = Self;
+    type CanonicalVTable = NotSupported;
+    type OperationsVTable = NotSupported;
+    type ValidityVTable = NotSupported;
+    type VisitorVTable = Self;
+    type ComputeVTable = NotSupported;
+    type EncodeVTable = NotSupported;
+    type SerdeVTable = NotSupported;
+    type OperatorVTable = Self;
+
+    fn id(_encoding: &Self::Encoding) -> EncodingId {
+        EncodingId::from("vortex.is_null")
+    }
+
+    fn encoding(_array: &Self::Array) -> EncodingRef {
+        EncodingRef::from(IsNotNullEncoding.as_ref())
+    }
+}
+
+impl ArrayVTable<IsNotNullVTable> for IsNotNullVTable {
+    fn len(array: &IsNotNullArray) -> usize {
+        array.len()
+    }
+
+    fn dtype(_array: &IsNotNullArray) -> &DType {
+        &DType::Bool(NonNullable)
+    }
+
+    fn stats(array: &IsNotNullArray) -> StatsSetRef<'_> {
+        array.stats.to_ref(array.as_ref())
+    }
+
+    fn array_hash<H: Hasher>(array: &IsNotNullArray, state: &mut H, precision: Precision) {
+        array.child.array_hash(state, precision);
+    }
+
+    fn array_eq(array: &IsNotNullArray, other: &IsNotNullArray, precision: Precision) -> bool {
+        array.child.array_eq(&other.child, precision)
+    }
+}
+
+impl VisitorVTable<IsNotNullVTable> for IsNotNullVTable {
+    fn visit_buffers(_array: &IsNotNullArray, _visitor: &mut dyn ArrayBufferVisitor) {
+        // No buffers
+    }
+
+    fn visit_children(array: &IsNotNullArray, visitor: &mut dyn ArrayChildVisitor) {
+        visitor.visit_child("child", array.child.as_ref());
+    }
+}
+
+impl OperatorVTable<IsNotNullVTable> for IsNotNullVTable {
+    fn bind(
+        array: &IsNotNullArray,
+        selection: Option<&ArrayRef>,
+        ctx: &mut dyn BindCtx,
+    ) -> VortexResult<BatchKernelRef> {
+        let child = ctx.bind(&array.child, selection)?;
+        Ok(kernel(move || {
+            let child = child.execute()?;
+            let is_null = child.validity().to_bit_buffer();
+            Ok(BoolVector::new(is_null, Mask::AllTrue(child.len())).into())
+        }))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::{bitbuffer, buffer};
+    use vortex_error::VortexResult;
+    use vortex_vector::VectorOps;
+
+    use super::IsNotNullArray;
+    use crate::IntoArray;
+    use crate::arrays::PrimitiveArray;
+    use crate::validity::Validity;
+
+    #[test]
+    fn test_is_null() -> VortexResult<()> {
+        let validity = bitbuffer![1 0 1];
+        let array = PrimitiveArray::new(
+            buffer![0, 1, 2],
+            Validity::Array(validity.clone().into_array()),
+        )
+        .into_array();
+
+        let result = IsNotNullArray::new(array).execute()?.into_bool();
+        assert!(result.validity().all_true());
+        assert_eq!(result.bits(), &validity);
+
+        Ok(())
+    }
+}

vortex-array/src/compute/arrays/is_null.rs

@@ -0,0 +1,139 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::hash::Hasher;
+use std::ops::Not;
+
+use vortex_dtype::DType;
+use vortex_dtype::Nullability::NonNullable;
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::{BoolVector, VectorOps};
+
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::stats::{ArrayStats, StatsSetRef};
+use crate::vtable::{ArrayVTable, NotSupported, OperatorVTable, VTable, VisitorVTable};
+use crate::{
+    ArrayBufferVisitor, ArrayChildVisitor, ArrayEq, ArrayHash, ArrayRef, EncodingId, EncodingRef,
+    Precision, vtable,
+};
+
+vtable!(IsNull);
+
+#[derive(Debug, Clone)]
+pub struct IsNullArray {
+    child: ArrayRef,
+    stats: ArrayStats,
+}
+
+impl IsNullArray {
+    /// Create a new is_null array.
+    pub fn new(child: ArrayRef) -> Self {
+        Self {
+            child,
+            stats: ArrayStats::default(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct IsNullEncoding;
+
+impl VTable for IsNullVTable {
+    type Array = IsNullArray;
+    type Encoding = IsNullEncoding;
+    type ArrayVTable = Self;
+    type CanonicalVTable = NotSupported;
+    type OperationsVTable = NotSupported;
+    type ValidityVTable = NotSupported;
+    type VisitorVTable = Self;
+    type ComputeVTable = NotSupported;
+    type EncodeVTable = NotSupported;
+    type SerdeVTable = NotSupported;
+    type OperatorVTable = Self;
+
+    fn id(_encoding: &Self::Encoding) -> EncodingId {
+        EncodingId::from("vortex.is_null")
+    }
+
+    fn encoding(_array: &Self::Array) -> EncodingRef {
+        EncodingRef::from(IsNullEncoding.as_ref())
+    }
+}
+
+impl ArrayVTable<IsNullVTable> for IsNullVTable {
+    fn len(array: &IsNullArray) -> usize {
+        array.len()
+    }
+
+    fn dtype(_array: &IsNullArray) -> &DType {
+        &DType::Bool(NonNullable)
+    }
+
+    fn stats(array: &IsNullArray) -> StatsSetRef<'_> {
+        array.stats.to_ref(array.as_ref())
+    }
+
+    fn array_hash<H: Hasher>(array: &IsNullArray, state: &mut H, precision: Precision) {
+        array.child.array_hash(state, precision);
+    }
+
+    fn array_eq(array: &IsNullArray, other: &IsNullArray, precision: Precision) -> bool {
+        array.child.array_eq(&other.child, precision)
+    }
+}
+
+impl VisitorVTable<IsNullVTable> for IsNullVTable {
+    fn visit_buffers(_array: &IsNullArray, _visitor: &mut dyn ArrayBufferVisitor) {
+        // No buffers
+    }
+
+    fn visit_children(array: &IsNullArray, visitor: &mut dyn ArrayChildVisitor) {
+        visitor.visit_child("child", array.child.as_ref());
+    }
+}
+
+impl OperatorVTable<IsNullVTable> for IsNullVTable {
+    fn bind(
+        array: &IsNullArray,
+        selection: Option<&ArrayRef>,
+        ctx: &mut dyn BindCtx,
+    ) -> VortexResult<BatchKernelRef> {
+        let child = ctx.bind(&array.child, selection)?;
+        Ok(kernel(move || {
+            let child = child.execute()?;
+            let is_null = child.validity().not().to_bit_buffer();
+            Ok(BoolVector::new(is_null, Mask::AllTrue(child.len())).into())
+        }))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::ops::Not;
+
+    use vortex_buffer::{bitbuffer, buffer};
+    use vortex_error::VortexResult;
+    use vortex_vector::VectorOps;
+
+    use crate::IntoArray;
+    use crate::arrays::PrimitiveArray;
+    use crate::compute::arrays::is_null::IsNullArray;
+    use crate::validity::Validity;
+
+    #[test]
+    fn test_is_null() -> VortexResult<()> {
+        let validity = bitbuffer![1 0 1];
+        let array = PrimitiveArray::new(
+            buffer![0, 1, 2],
+            Validity::Array(validity.clone().into_array()),
+        )
+        .into_array();
+
+        let result = IsNullArray::new(array).execute()?.into_bool();
+        assert!(result.validity().all_true());
+        assert_eq!(result.bits(), &validity.not());
+
+        Ok(())
+    }
+}

vortex-array/src/compute/arrays/mod.rs

@@ -2,4 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 pub mod arithmetic;
+pub mod is_not_null;
+pub mod is_null;
 pub mod logical;

vortex-compute/Cargo.toml

@@ -28,9 +28,10 @@ vortex-vector = { workspace = true }
 num-traits = { workspace = true }
 
 [features]
-default = ["arithmetic", "comparison", "filter", "logical"]
+default = ["arithmetic", "comparison", "filter", "logical", "mask"]
 
 arithmetic = []
 comparison = []
 filter = []
 logical = []
+mask = []

vortex-compute/src/lib.rs

@@ -15,3 +15,5 @@ pub mod comparison;
 pub mod filter;
 #[cfg(feature = "logical")]
 pub mod logical;
+#[cfg(feature = "mask")]
+pub mod mask;

vortex-compute/src/logical/and.rs

@@ -3,7 +3,6 @@
 
 use std::ops::BitAnd;
 
-use vortex_mask::Mask;
 use vortex_vector::{BoolVector, VectorOps};
 
 use crate::logical::LogicalAnd;
@@ -28,22 +27,6 @@ impl LogicalAnd<&BoolVector> for BoolVector {
     }
 }
 
-impl LogicalAnd for &Mask {
-    type Output = Mask;
-
-    fn and(self, other: Self) -> Self::Output {
-        self.bitand(other)
-    }
-}
-
-impl LogicalAnd<&Mask> for Mask {
-    type Output = Mask;
-
-    fn and(self, other: &Mask) -> Self::Output {
-        self.bitand(other)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use vortex_buffer::bitbuffer;