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

datatypes/vectors/
boolean.rs
