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

datatypes/vectors/
boolean.rs
