vortex_array/compute/arrays/
is_not_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::hash::Hasher;
5
6use vortex_dtype::DType;
7use vortex_dtype::Nullability::NonNullable;
8use vortex_error::VortexResult;
9use vortex_mask::Mask;
10use vortex_vector::VectorOps;
11use vortex_vector::bool::BoolVector;
12
13use crate::execution::{BatchKernelRef, BindCtx, kernel};
14use crate::stats::{ArrayStats, StatsSetRef};
15use crate::vtable::{ArrayVTable, NotSupported, OperatorVTable, VTable, VisitorVTable};
16use crate::{
17    ArrayBufferVisitor, ArrayChildVisitor, ArrayEq, ArrayHash, ArrayRef, EncodingId, EncodingRef,
18    Precision, vtable,
19};
20
21vtable!(IsNotNull);
22
23#[derive(Debug, Clone)]
24pub struct IsNotNullArray {
25    child: ArrayRef,
26    stats: ArrayStats,
27}
28
29impl IsNotNullArray {
30    /// Create a new is_not_null array.
31    pub fn new(child: ArrayRef) -> Self {
32        Self {
33            child,
34            stats: ArrayStats::default(),
35        }
36    }
37}
38
39#[derive(Debug, Clone)]
40pub struct IsNotNullEncoding;
41
42impl VTable for IsNotNullVTable {
43    type Array = IsNotNullArray;
44    type Encoding = IsNotNullEncoding;
45    type ArrayVTable = Self;
46    type CanonicalVTable = NotSupported;
47    type OperationsVTable = NotSupported;
48    type ValidityVTable = NotSupported;
49    type VisitorVTable = Self;
50    type ComputeVTable = NotSupported;
51    type EncodeVTable = NotSupported;
52    type SerdeVTable = NotSupported;
53    type OperatorVTable = Self;
54
55    fn id(_encoding: &Self::Encoding) -> EncodingId {
56        EncodingId::from("vortex.is_not_null")
57    }
58
59    fn encoding(_array: &Self::Array) -> EncodingRef {
60        EncodingRef::from(IsNotNullEncoding.as_ref())
61    }
62}
63
64impl ArrayVTable<IsNotNullVTable> for IsNotNullVTable {
65    fn len(array: &IsNotNullArray) -> usize {
66        array.child.len()
67    }
68
69    fn dtype(_array: &IsNotNullArray) -> &DType {
70        &DType::Bool(NonNullable)
71    }
72
73    fn stats(array: &IsNotNullArray) -> StatsSetRef<'_> {
74        array.stats.to_ref(array.as_ref())
75    }
76
77    fn array_hash<H: Hasher>(array: &IsNotNullArray, state: &mut H, precision: Precision) {
78        array.child.array_hash(state, precision);
79    }
80
81    fn array_eq(array: &IsNotNullArray, other: &IsNotNullArray, precision: Precision) -> bool {
82        array.child.array_eq(&other.child, precision)
83    }
84}
85
86impl VisitorVTable<IsNotNullVTable> for IsNotNullVTable {
87    fn visit_buffers(_array: &IsNotNullArray, _visitor: &mut dyn ArrayBufferVisitor) {
88        // No buffers
89    }
90
91    fn visit_children(array: &IsNotNullArray, visitor: &mut dyn ArrayChildVisitor) {
92        visitor.visit_child("child", array.child.as_ref());
93    }
94}
95
96impl OperatorVTable<IsNotNullVTable> for IsNotNullVTable {
97    fn bind(
98        array: &IsNotNullArray,
99        selection: Option<&ArrayRef>,
100        ctx: &mut dyn BindCtx,
101    ) -> VortexResult<BatchKernelRef> {
102        let child = ctx.bind(&array.child, selection)?;
103        Ok(kernel(move || {
104            let child = child.execute()?;
105            let is_null = child.validity().to_bit_buffer();
106            Ok(BoolVector::new(is_null, Mask::AllTrue(child.len())).into())
107        }))
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use vortex_buffer::{bitbuffer, buffer};
114    use vortex_error::VortexResult;
115    use vortex_vector::VectorOps;
116
117    use super::IsNotNullArray;
118    use crate::IntoArray;
119    use crate::arrays::PrimitiveArray;
120    use crate::validity::Validity;
121
122    #[test]
123    fn test_is_null() -> VortexResult<()> {
124        let validity = bitbuffer![1 0 1];
125        let array = PrimitiveArray::new(
126            buffer![0, 1, 2],
127            Validity::Array(validity.clone().into_array()),
128        )
129        .into_array();
130
131        let result = IsNotNullArray::new(array).execute()?.into_bool();
132        assert!(result.validity().all_true());
133        assert_eq!(result.bits(), &validity);
134
135        Ok(())
136    }
137}