taos_query/common/raw/views/
bool_view.rs

1use std::{ffi::c_void, io::Write, ops::Range};
2
3use crate::common::{BorrowedValue, Ty};
4
5use super::{IsColumnView, NullBits, NullsIter};
6
7use bytes::Bytes;
8
9type View = BoolView;
10
11#[derive(Debug, Clone)]
12pub struct BoolView {
13    pub(crate) nulls: NullBits,
14    pub(crate) data: Bytes,
15}
16
17impl IsColumnView for BoolView {
18    fn ty(&self) -> Ty {
19        Ty::Bool
20    }
21    fn from_borrowed_value_iter<'b>(iter: impl Iterator<Item = BorrowedValue<'b>>) -> Self {
22        Self::from_iter(iter.map(|v| v.to_bool()))
23    }
24}
25
26impl std::ops::Add for BoolView {
27    type Output = Self;
28
29    fn add(self, rhs: Self) -> Self::Output {
30        &self + &rhs
31    }
32}
33impl std::ops::Add for &BoolView {
34    type Output = BoolView;
35
36    fn add(self, rhs: Self) -> Self::Output {
37        let nulls = NullBits::from_iter(
38            self.nulls
39                .iter()
40                .take(self.len())
41                .chain(rhs.nulls.iter().take(rhs.len())),
42        );
43        let data: Bytes = self
44            .data
45            .as_ref()
46            .iter()
47            .chain(rhs.data.as_ref().iter())
48            .copied()
49            .collect();
50
51        BoolView { nulls, data }
52    }
53}
54
55impl std::ops::Add<BoolView> for &BoolView {
56    type Output = BoolView;
57
58    fn add(self, rhs: BoolView) -> Self::Output {
59        self + &rhs
60    }
61}
62
63impl std::ops::Add<&BoolView> for BoolView {
64    type Output = BoolView;
65
66    fn add(self, rhs: &BoolView) -> Self::Output {
67        &self + rhs
68    }
69}
70
71impl BoolView {
72    /// Rows
73    pub fn len(&self) -> usize {
74        self.data.len()
75    }
76
77    // /// Raw slice of `bool` type.
78    // pub fn as_raw_slice(&self) -> &[bool] {
79    //     unsafe { std::slice::from_raw_parts(self.data.as_ptr() as *const bool, self.len()) }
80    // }
81
82    /// Raw pointer of the slice.
83    pub fn as_raw_ptr(&self) -> *const bool {
84        self.data.as_ptr() as *const bool
85    }
86
87    /// A iterator only decide if the value at some row index is NULL or not.
88    pub fn is_null_iter(&self) -> NullsIter {
89        NullsIter {
90            nulls: &self.nulls,
91            row: 0,
92            len: self.len(),
93        }
94    }
95
96    /// Build a nulls vector.
97    pub fn to_nulls_vec(&self) -> Vec<bool> {
98        self.is_null_iter().collect()
99    }
100
101    /// Check if the value at `row` index is NULL or not.
102    pub fn is_null(&self, row: usize) -> bool {
103        if row < self.len() {
104            unsafe { self.is_null_unchecked(row) }
105        } else {
106            false
107        }
108    }
109
110    /// Unsafe version for [methods.is_null]
111    pub unsafe fn is_null_unchecked(&self, row: usize) -> bool {
112        self.nulls.is_null_unchecked(row)
113    }
114
115    /// Get nullable value at `row` index.
116    pub fn get(&self, row: usize) -> Option<bool> {
117        if row < self.len() {
118            unsafe { self.get_unchecked(row) }
119        } else {
120            None
121        }
122    }
123
124    #[inline(always)]
125    unsafe fn get_raw_at(&self, index: usize) -> *const bool {
126        self.data.as_ptr().add(index) as _
127    }
128
129    /// Get nullable value at `row` index.
130    pub unsafe fn get_unchecked(&self, row: usize) -> Option<bool> {
131        if self.nulls.is_null_unchecked(row) {
132            None
133        } else {
134            Some(*self.get_raw_at(row))
135        }
136    }
137
138    pub unsafe fn get_ref_unchecked(&self, row: usize) -> Option<*const bool> {
139        if self.nulls.is_null_unchecked(row) {
140            None
141        } else {
142            Some(self.get_raw_at(row))
143        }
144    }
145
146    pub unsafe fn get_value_unchecked(&self, row: usize) -> BorrowedValue {
147        self.get_unchecked(row)
148            .map(BorrowedValue::Bool)
149            .unwrap_or(BorrowedValue::Null(Ty::Bool))
150    }
151
152    pub unsafe fn get_raw_value_unchecked(&self, row: usize) -> (Ty, u32, *const c_void) {
153        if self.nulls.is_null_unchecked(row) {
154            (Ty::Bool, std::mem::size_of::<bool>() as _, std::ptr::null())
155        } else {
156            (
157                Ty::Bool,
158                std::mem::size_of::<bool>() as _,
159                self.get_raw_at(row) as _,
160            )
161        }
162    }
163
164    pub fn slice(&self, mut range: Range<usize>) -> Option<Self> {
165        if range.start >= self.len() {
166            return None;
167        }
168        if range.end >= self.len() {
169            range.end = self.len();
170        }
171        if range.is_empty() {
172            return None;
173        }
174
175        let nulls = unsafe { self.nulls.slice(range.clone()) };
176        let data = self.data.slice(range);
177        Some(Self { nulls, data })
178    }
179
180    /// A iterator to nullable values of current row.
181    pub fn iter(&self) -> BoolViewIter {
182        BoolViewIter { view: self, row: 0 }
183    }
184
185    /// Convert data to a vector of all nullable values.
186    pub fn to_vec(&self) -> Vec<Option<bool>> {
187        self.iter().collect()
188    }
189
190    /// Write column data as raw bytes.
191    pub(crate) fn write_raw_into<W: Write>(&self, wtr: &mut W) -> std::io::Result<usize> {
192        let nulls = self.nulls.0.as_ref();
193        debug_assert_eq!(nulls.len(), (self.len() + 7) / 8);
194        wtr.write_all(nulls)?;
195        wtr.write_all(&self.data)?;
196        Ok(nulls.len() + self.data.len())
197    }
198
199    pub fn concat(&self, rhs: &View) -> View {
200        let nulls = NullBits::from_iter(
201            self.nulls
202                .iter()
203                .take(self.len())
204                .chain(rhs.nulls.iter().take(rhs.len())),
205        );
206        let data: Bytes = self
207            .data
208            .as_ref()
209            .iter()
210            .chain(rhs.data.as_ref().iter())
211            .copied()
212            .collect();
213
214        View { nulls, data }
215    }
216}
217
218pub struct BoolViewIter<'a> {
219    view: &'a BoolView,
220    row: usize,
221}
222
223impl<'a> Iterator for BoolViewIter<'a> {
224    type Item = Option<bool>;
225
226    fn next(&mut self) -> Option<Self::Item> {
227        if self.row < self.view.len() {
228            let row = self.row;
229            self.row += 1;
230            Some(unsafe { self.view.get_unchecked(row) })
231        } else {
232            None
233        }
234    }
235
236    #[inline]
237    fn size_hint(&self) -> (usize, Option<usize>) {
238        if self.row < self.view.len() {
239            let len = self.view.len() - self.row;
240            (len, Some(len))
241        } else {
242            (0, Some(0))
243        }
244    }
245}
246
247impl<'a> ExactSizeIterator for BoolViewIter<'a> {
248    fn len(&self) -> usize {
249        self.view.len() - self.row
250    }
251}
252
253impl<A: Into<Option<bool>>> FromIterator<A> for BoolView {
254    fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> Self {
255        let (nulls, mut values): (Vec<bool>, Vec<_>) = iter
256            .into_iter()
257            .map(|v| match v.into() {
258                Some(v) => (false, v),
259                None => (true, false),
260            })
261            .unzip();
262        Self {
263            nulls: NullBits::from_iter(nulls),
264            data: Bytes::from({
265                let (ptr, len, cap) = (values.as_mut_ptr(), values.len(), values.capacity());
266                std::mem::forget(values);
267                unsafe { Vec::from_raw_parts(ptr as *mut u8, len, cap) }
268            }),
269        }
270    }
271}
272
273#[test]
274fn test_bool_slice() {
275    let data = [true, false, false, true];
276    let view = BoolView::from_iter(data);
277    dbg!(&view);
278    let slice = view.slice(1..3);
279    dbg!(&slice);
280
281    let data = [None, Some(false), Some(true), None];
282    let view = BoolView::from_iter(data);
283    dbg!(&view);
284    let slice = view.slice(1..4);
285    dbg!(&slice);
286
287    let chain = &view + slice.as_ref().unwrap();
288    dbg!(&chain);
289    assert!(chain.len() == 7);
290    assert!(chain
291        .slice(0..4)
292        .unwrap()
293        .iter()
294        .zip(view.iter())
295        .all(|(l, r)| dbg!(l) == r));
296}