datatypes/vectors/
boolean.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 std::any::Any;
16use std::sync::Arc;
17
18use arrow::array::{Array, ArrayBuilder, ArrayIter, ArrayRef, BooleanArray, BooleanBuilder};
19use snafu::ResultExt;
20
21use crate::data_type::ConcreteDataType;
22use crate::error::Result;
23use crate::scalars::{ScalarVector, ScalarVectorBuilder};
24use crate::serialize::Serializable;
25use crate::value::{Value, ValueRef};
26use crate::vectors::{self, MutableVector, Validity, Vector, VectorRef};
27
28/// Vector of boolean.
29#[derive(Debug, PartialEq)]
30pub struct BooleanVector {
31    array: BooleanArray,
32}
33
34impl BooleanVector {
35    pub(crate) fn as_arrow(&self) -> &dyn Array {
36        &self.array
37    }
38
39    /// Get the inner boolean array.
40    pub fn as_boolean_array(&self) -> &BooleanArray {
41        &self.array
42    }
43
44    pub(crate) fn false_count(&self) -> usize {
45        self.array.false_count()
46    }
47}
48
49impl From<Vec<bool>> for BooleanVector {
50    fn from(data: Vec<bool>) -> Self {
51        BooleanVector {
52            array: BooleanArray::from(data),
53        }
54    }
55}
56
57impl From<BooleanArray> for BooleanVector {
58    fn from(array: BooleanArray) -> Self {
59        Self { array }
60    }
61}
62
63impl From<Vec<Option<bool>>> for BooleanVector {
64    fn from(data: Vec<Option<bool>>) -> Self {
65        BooleanVector {
66            array: BooleanArray::from(data),
67        }
68    }
69}
70
71impl FromIterator<Option<bool>> for BooleanVector {
72    fn from_iter<T: IntoIterator<Item = Option<bool>>>(iter: T) -> Self {
73        BooleanVector {
74            array: BooleanArray::from_iter(iter),
75        }
76    }
77}
78
79impl Vector for BooleanVector {
80    fn data_type(&self) -> ConcreteDataType {
81        ConcreteDataType::boolean_datatype()
82    }
83
84    fn vector_type_name(&self) -> String {
85        "BooleanVector".to_string()
86    }
87
88    fn as_any(&self) -> &dyn Any {
89        self
90    }
91
92    fn len(&self) -> usize {
93        self.array.len()
94    }
95
96    fn to_arrow_array(&self) -> ArrayRef {
97        Arc::new(self.array.clone())
98    }
99
100    fn to_boxed_arrow_array(&self) -> Box<dyn Array> {
101        Box::new(self.array.clone())
102    }
103
104    fn validity(&self) -> Validity {
105        vectors::impl_validity_for_vector!(self.array)
106    }
107
108    fn memory_size(&self) -> usize {
109        self.array.get_buffer_memory_size()
110    }
111
112    fn null_count(&self) -> usize {
113        self.array.null_count()
114    }
115
116    fn is_null(&self, row: usize) -> bool {
117        self.array.is_null(row)
118    }
119
120    fn slice(&self, offset: usize, length: usize) -> VectorRef {
121        Arc::new(Self::from(self.array.slice(offset, length)))
122    }
123
124    fn get(&self, index: usize) -> Value {
125        vectors::impl_get_for_vector!(self.array, index)
126    }
127
128    fn get_ref(&self, index: usize) -> ValueRef<'_> {
129        vectors::impl_get_ref_for_vector!(self.array, index)
130    }
131}
132
133impl ScalarVector for BooleanVector {
134    type OwnedItem = bool;
135    type RefItem<'a> = bool;
136    type Iter<'a> = ArrayIter<&'a BooleanArray>;
137    type Builder = BooleanVectorBuilder;
138
139    fn get_data(&self, idx: usize) -> Option<Self::RefItem<'_>> {
140        if self.array.is_valid(idx) {
141            Some(self.array.value(idx))
142        } else {
143            None
144        }
145    }
146
147    fn iter_data(&self) -> Self::Iter<'_> {
148        self.array.iter()
149    }
150}
151
152pub struct BooleanVectorBuilder {
153    mutable_array: BooleanBuilder,
154}
155
156impl MutableVector for BooleanVectorBuilder {
157    fn data_type(&self) -> ConcreteDataType {
158        ConcreteDataType::boolean_datatype()
159    }
160
161    fn len(&self) -> usize {
162        self.mutable_array.len()
163    }
164
165    fn as_any(&self) -> &dyn Any {
166        self
167    }
168
169    fn as_mut_any(&mut self) -> &mut dyn Any {
170        self
171    }
172
173    fn to_vector(&mut self) -> VectorRef {
174        Arc::new(self.finish())
175    }
176
177    fn to_vector_cloned(&self) -> VectorRef {
178        Arc::new(self.finish_cloned())
179    }
180
181    fn try_push_value_ref(&mut self, value: &ValueRef) -> Result<()> {
182        match value.try_into_boolean()? {
183            Some(v) => self.mutable_array.append_value(v),
184            None => self.mutable_array.append_null(),
185        }
186        Ok(())
187    }
188
189    fn extend_slice_of(&mut self, vector: &dyn Vector, offset: usize, length: usize) -> Result<()> {
190        vectors::impl_extend_for_builder!(self, vector, BooleanVector, offset, length)
191    }
192
193    fn push_null(&mut self) {
194        self.mutable_array.append_null()
195    }
196}
197
198impl ScalarVectorBuilder for BooleanVectorBuilder {
199    type VectorType = BooleanVector;
200
201    fn with_capacity(capacity: usize) -> Self {
202        Self {
203            mutable_array: BooleanBuilder::with_capacity(capacity),
204        }
205    }
206
207    fn push(&mut self, value: Option<<Self::VectorType as ScalarVector>::RefItem<'_>>) {
208        match value {
209            Some(v) => self.mutable_array.append_value(v),
210            None => self.mutable_array.append_null(),
211        }
212    }
213
214    fn finish(&mut self) -> Self::VectorType {
215        BooleanVector {
216            array: self.mutable_array.finish(),
217        }
218    }
219
220    fn finish_cloned(&self) -> Self::VectorType {
221        BooleanVector {
222            array: self.mutable_array.finish_cloned(),
223        }
224    }
225}
226
227impl Serializable for BooleanVector {
228    fn serialize_to_json(&self) -> Result<Vec<serde_json::Value>> {
229        self.iter_data()
230            .map(serde_json::to_value)
231            .collect::<serde_json::Result<_>>()
232            .context(crate::error::SerializeSnafu)
233    }
234}
235
236vectors::impl_try_from_arrow_array_for_vector!(BooleanArray, BooleanVector);
237
238#[cfg(test)]
239mod tests {
240    use arrow::datatypes::DataType as ArrowDataType;
241    use serde_json;
242
243    use super::*;
244    use crate::data_type::DataType;
245    use crate::serialize::Serializable;
246    use crate::types::BooleanType;
247
248    #[test]
249    fn test_boolean_vector_misc() {
250        let bools = vec![true, false, true, true, false, false, true, true, false];
251        let v = BooleanVector::from(bools.clone());
252        assert_eq!(9, v.len());
253        assert_eq!("BooleanVector", v.vector_type_name());
254        assert!(!v.is_const());
255        assert!(v.validity().is_all_valid());
256        assert!(!v.only_null());
257        assert_eq!(2, v.memory_size());
258
259        for (i, b) in bools.iter().enumerate() {
260            assert!(!v.is_null(i));
261            assert_eq!(Value::Boolean(*b), v.get(i));
262            assert_eq!(ValueRef::Boolean(*b), v.get_ref(i));
263        }
264
265        let arrow_arr = v.to_arrow_array();
266        assert_eq!(9, arrow_arr.len());
267        assert_eq!(&ArrowDataType::Boolean, arrow_arr.data_type());
268    }
269
270    #[test]
271    fn test_serialize_boolean_vector_to_json() {
272        let vector = BooleanVector::from(vec![true, false, true, true, false, false]);
273
274        let json_value = vector.serialize_to_json().unwrap();
275        assert_eq!(
276            "[true,false,true,true,false,false]",
277            serde_json::to_string(&json_value).unwrap(),
278        );
279    }
280
281    #[test]
282    fn test_serialize_boolean_vector_with_null_to_json() {
283        let vector = BooleanVector::from(vec![Some(true), None, Some(false)]);
284
285        let json_value = vector.serialize_to_json().unwrap();
286        assert_eq!(
287            "[true,null,false]",
288            serde_json::to_string(&json_value).unwrap(),
289        );
290    }
291
292    #[test]
293    fn test_boolean_vector_from_vec() {
294        let input = vec![false, true, false, true];
295        let vec = BooleanVector::from(input.clone());
296        assert_eq!(4, vec.len());
297        for (i, v) in input.into_iter().enumerate() {
298            assert_eq!(Some(v), vec.get_data(i), "Failed at {i}")
299        }
300    }
301
302    #[test]
303    fn test_boolean_vector_from_iter() {
304        let input = vec![Some(false), Some(true), Some(false), Some(true)];
305        let vec = input.iter().cloned().collect::<BooleanVector>();
306        assert_eq!(4, vec.len());
307        for (i, v) in input.into_iter().enumerate() {
308            assert_eq!(v, vec.get_data(i), "Failed at {i}")
309        }
310    }
311
312    #[test]
313    fn test_boolean_vector_from_vec_option() {
314        let input = vec![Some(false), Some(true), None, Some(true)];
315        let vec = BooleanVector::from(input.clone());
316        assert_eq!(4, vec.len());
317        for (i, v) in input.into_iter().enumerate() {
318            assert_eq!(v, vec.get_data(i), "failed at {i}")
319        }
320    }
321
322    #[test]
323    fn test_boolean_vector_build_get() {
324        let input = [Some(true), None, Some(false)];
325        let mut builder = BooleanVectorBuilder::with_capacity(3);
326        for v in input {
327            builder.push(v);
328        }
329        let vector = builder.finish();
330        assert_eq!(input.len(), vector.len());
331
332        let res: Vec<_> = vector.iter_data().collect();
333        assert_eq!(input, &res[..]);
334
335        for (i, v) in input.into_iter().enumerate() {
336            assert_eq!(v, vector.get_data(i));
337            assert_eq!(Value::from(v), vector.get(i));
338        }
339    }
340
341    #[test]
342    fn test_boolean_vector_validity() {
343        let vector = BooleanVector::from(vec![Some(true), None, Some(false)]);
344        assert_eq!(1, vector.null_count());
345        let validity = vector.validity();
346        assert_eq!(1, validity.null_count());
347        assert!(!validity.is_set(1));
348
349        let vector = BooleanVector::from(vec![true, false, false]);
350        assert_eq!(0, vector.null_count());
351        assert!(vector.validity().is_all_valid());
352    }
353
354    #[test]
355    fn test_boolean_vector_builder() {
356        let input = BooleanVector::from_slice(&[true, false, true]);
357
358        let mut builder = BooleanType.create_mutable_vector(3);
359        builder.push_value_ref(&ValueRef::Boolean(true));
360        assert!(builder.try_push_value_ref(&ValueRef::Int32(123)).is_err());
361        builder.extend_slice_of(&input, 1, 2).unwrap();
362        assert!(
363            builder
364                .extend_slice_of(&crate::vectors::Int32Vector::from_slice([13]), 0, 1)
365                .is_err()
366        );
367        let vector = builder.to_vector();
368
369        let expect: VectorRef = Arc::new(BooleanVector::from_slice(&[true, false, true]));
370        assert_eq!(expect, vector);
371    }
372
373    #[test]
374    fn test_boolean_vector_builder_finish_cloned() {
375        let mut builder = BooleanVectorBuilder::with_capacity(1024);
376        builder.push(Some(true));
377        builder.push(Some(false));
378        builder.push(Some(true));
379        let vector = builder.finish_cloned();
380        assert!(vector.get_data(0).unwrap());
381        assert_eq!(vector.len(), 3);
382        assert_eq!(builder.len(), 3);
383
384        builder.push(Some(false));
385        let vector = builder.finish_cloned();
386        assert!(!vector.get_data(3).unwrap());
387        assert_eq!(builder.len(), 4);
388    }
389}
datatypes/vectors/boolean.rs

datatypes/vectors/
boolean.rs
