taos_query/common/raw/views/
tinyint_view.rs

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