datatypes/vectors/
validity.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use arrow::array::ArrayData;
16use arrow::buffer::NullBuffer;
17
18#[derive(Debug, PartialEq)]
19enum ValidityKind {
20    /// Whether the array slot is valid or not (null).
21    Slots {
22        bitmap: NullBuffer,
23        len: usize,
24        null_count: usize,
25    },
26    /// All slots are valid.
27    AllValid { len: usize },
28    /// All slots are null.
29    AllNull { len: usize },
30}
31
32/// Validity of a vector.
33#[derive(Debug, PartialEq)]
34pub struct Validity {
35    kind: ValidityKind,
36}
37
38impl Validity {
39    /// Creates a `Validity` from [`ArrayData`].
40    pub fn from_array_data(data: ArrayData) -> Validity {
41        match data.nulls() {
42            Some(null_buf) => Validity {
43                kind: ValidityKind::Slots {
44                    bitmap: null_buf.clone(),
45                    len: data.len(),
46                    null_count: data.null_count(),
47                },
48            },
49            None => Validity::all_valid(data.len()),
50        }
51    }
52
53    /// Creates a `Validity` from a logical null buffer.
54    pub(crate) fn from_null_buffer(bitmap: NullBuffer) -> Validity {
55        let len = bitmap.len();
56        let null_count = bitmap.null_count();
57        Validity {
58            kind: ValidityKind::Slots {
59                bitmap,
60                len,
61                null_count,
62            },
63        }
64    }
65
66    /// Returns `Validity` that all elements are valid.
67    pub fn all_valid(len: usize) -> Validity {
68        Validity {
69            kind: ValidityKind::AllValid { len },
70        }
71    }
72
73    /// Returns `Validity` that all elements are null.
74    pub fn all_null(len: usize) -> Validity {
75        Validity {
76            kind: ValidityKind::AllNull { len },
77        }
78    }
79
80    /// Returns whether `i-th` bit is set.
81    pub fn is_set(&self, i: usize) -> bool {
82        match &self.kind {
83            ValidityKind::Slots { bitmap, .. } => bitmap.is_valid(i),
84            ValidityKind::AllValid { len } => i < *len,
85            ValidityKind::AllNull { .. } => false,
86        }
87    }
88
89    /// Returns true if all bits are null.
90    pub fn is_all_null(&self) -> bool {
91        match self.kind {
92            ValidityKind::Slots {
93                len, null_count, ..
94            } => len == null_count,
95            ValidityKind::AllValid { .. } => false,
96            ValidityKind::AllNull { .. } => true,
97        }
98    }
99
100    /// Returns true if all bits are valid.
101    pub fn is_all_valid(&self) -> bool {
102        match self.kind {
103            ValidityKind::Slots { null_count, .. } => null_count == 0,
104            ValidityKind::AllValid { .. } => true,
105            ValidityKind::AllNull { .. } => false,
106        }
107    }
108
109    /// The number of null slots.
110    pub fn null_count(&self) -> usize {
111        match self.kind {
112            ValidityKind::Slots { null_count, .. } => null_count,
113            ValidityKind::AllValid { .. } => 0,
114            ValidityKind::AllNull { len } => len,
115        }
116    }
117}
118
119#[cfg(test)]
120mod tests {
121    use arrow::array::{Array, Int32Array};
122
123    use super::*;
124
125    #[test]
126    fn test_all_valid() {
127        let validity = Validity::all_valid(5);
128        assert!(validity.is_all_valid());
129        assert!(!validity.is_all_null());
130        assert_eq!(0, validity.null_count());
131        for i in 0..5 {
132            assert!(validity.is_set(i));
133        }
134        assert!(!validity.is_set(5));
135    }
136
137    #[test]
138    fn test_all_null() {
139        let validity = Validity::all_null(5);
140        assert!(validity.is_all_null());
141        assert!(!validity.is_all_valid());
142        assert_eq!(5, validity.null_count());
143        for i in 0..5 {
144            assert!(!validity.is_set(i));
145        }
146        assert!(!validity.is_set(5));
147    }
148
149    #[test]
150    fn test_from_array_data() {
151        let array = Int32Array::from_iter([None, Some(1), None]);
152        let validity = Validity::from_array_data(array.to_data());
153        assert_eq!(2, validity.null_count());
154        assert!(!validity.is_set(0));
155        assert!(validity.is_set(1));
156        assert!(!validity.is_set(2));
157        assert!(!validity.is_all_null());
158        assert!(!validity.is_all_valid());
159
160        let array = Int32Array::from_iter([None, None]);
161        let validity = Validity::from_array_data(array.to_data());
162        assert!(validity.is_all_null());
163        assert!(!validity.is_all_valid());
164        assert_eq!(2, validity.null_count());
165
166        let array = Int32Array::from_iter_values([1, 2]);
167        let validity = Validity::from_array_data(array.to_data());
168        assert!(!validity.is_all_null());
169        assert!(validity.is_all_valid());
170        assert_eq!(0, validity.null_count());
171    }
172}
datatypes/vectors/validity.rs

datatypes/vectors/
validity.rs
