taos_query/common/raw/views/
smallint_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 = i16;
10type View = SmallIntView;
11const ITEM_SIZE: usize = std::mem::size_of::<Item>();
12#[derive(Debug, Clone)]
13pub struct SmallIntView {
14    pub(crate) nulls: NullBits,
15    pub(crate) data: Bytes,
16}
17impl IsColumnView for View {
18    fn ty(&self) -> Ty {
19        Ty::USmallInt
20    }
21    fn from_borrowed_value_iter<'b>(iter: impl Iterator<Item = BorrowedValue<'b>>) -> Self {
22        Self::from_iter(iter.map(|v| v.to_i16()))
23    }
24}
25
26impl std::ops::Add for View {
27    type Output = Self;
28
29    fn add(self, rhs: Self) -> Self::Output {
30        &self + &rhs
31    }
32}
33impl std::ops::Add for &View {
34    type Output = View;
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        View { nulls, data }
52    }
53}
54
55impl std::ops::Add<View> for &View {
56    type Output = View;
57
58    fn add(self, rhs: View) -> Self::Output {
59        self + &rhs
60    }
61}
62
63impl std::ops::Add<&View> for View {
64    type Output = View;
65
66    fn add(self, rhs: &View) -> Self::Output {
67        &self + rhs
68    }
69}
70
71impl SmallIntView {
72    /// Rows
73    pub fn len(&self) -> usize {
74        self.data.len() / std::mem::size_of::<Item>()
75    }
76
77    /// Raw slice of target type.
78    pub fn as_raw_slice(&self) -> &[Item] {
79        unsafe { std::slice::from_raw_parts(self.data.as_ptr() as *const Item, self.len()) }
80    }
81
82    /// Raw pointer of the slice.
83    pub fn as_raw_ptr(&self) -> *const Item {
84        self.data.as_ptr() as *const Item
85    }
86
87    /// Build a nulls vector.
88    pub fn to_nulls_vec(&self) -> Vec<bool> {
89        self.is_null_iter().collect()
90    }
91
92    /// A iterator only decide if the value at some row index is NULL or not.
93    pub fn is_null_iter(&self) -> NullsIter {
94        NullsIter {
95            nulls: &self.nulls,
96            row: 0,
97            len: self.len(),
98        }
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<Item> {
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 Item {
126        self.data.as_ptr().add(index * ITEM_SIZE) as _
127    }
128
129    /// Get nullable value at `row` index.
130    pub unsafe fn get_unchecked(&self, row: usize) -> Option<Item> {
131        if self.nulls.is_null_unchecked(row) {
132            None
133        } else {
134            Some(self.get_raw_at(row).read_unaligned())
135        }
136    }
137
138    pub unsafe fn get_ref_unchecked(&self, row: usize) -> Option<*const Item> {
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::SmallInt)
149            .unwrap_or(BorrowedValue::Null(Ty::SmallInt))
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            (
155                Ty::SmallInt,
156                std::mem::size_of::<Item>() as _,
157                std::ptr::null(),
158            )
159        } else {
160            (
161                Ty::SmallInt,
162                std::mem::size_of::<Item>() as _,
163                self.get_raw_at(row) as *const Item as _,
164            )
165        }
166    }
167
168    /// Create a slice of view.
169    pub fn slice(&self, mut range: std::ops::Range<usize>) -> Option<Self> {
170        if range.start >= self.len() {
171            return None;
172        }
173        if range.end > self.len() {
174            range.end = self.len();
175        }
176        if range.is_empty() {
177            return None;
178        }
179
180        let nulls = unsafe { self.nulls.slice(range.clone()) };
181        let data = self
182            .data
183            .slice(range.start * ITEM_SIZE..range.end * ITEM_SIZE);
184        Some(Self { nulls, data })
185    }
186
187    /// A iterator to nullable values of current row.
188    pub fn iter(&self) -> SmallIntViewIter {
189        SmallIntViewIter { view: self, row: 0 }
190    }
191
192    /// Convert data to a vector of all nullable values.
193    pub fn to_vec(&self) -> Vec<Option<Item>> {
194        self.iter().collect()
195    }
196
197    /// Write column data as raw bytes.
198    pub(crate) fn write_raw_into<W: std::io::Write>(&self, mut wtr: W) -> std::io::Result<usize> {
199        let nulls = self.nulls.0.as_ref();
200        wtr.write_all(nulls)?;
201        wtr.write_all(&self.data)?;
202        Ok(nulls.len() + self.data.len())
203    }
204
205    pub fn concat(&self, rhs: &View) -> View {
206        let nulls = NullBits::from_iter(
207            self.nulls
208                .iter()
209                .take(self.len())
210                .chain(rhs.nulls.iter().take(rhs.len())),
211        );
212        let data: Bytes = self
213            .data
214            .as_ref()
215            .iter()
216            .chain(rhs.data.as_ref().iter())
217            .copied()
218            .collect();
219
220        View { nulls, data }
221    }
222}
223
224pub struct SmallIntViewIter<'a> {
225    view: &'a SmallIntView,
226    row: usize,
227}
228
229impl<'a> Iterator for SmallIntViewIter<'a> {
230    type Item = Option<Item>;
231
232    fn next(&mut self) -> Option<Self::Item> {
233        if self.row < self.view.len() {
234            let row = self.row;
235            self.row += 1;
236            Some(unsafe { self.view.get_unchecked(row) })
237        } else {
238            None
239        }
240    }
241
242    #[inline]
243    fn size_hint(&self) -> (usize, Option<usize>) {
244        if self.row < self.view.len() {
245            let len = self.view.len() - self.row;
246            (len, Some(len))
247        } else {
248            (0, Some(0))
249        }
250    }
251}
252
253impl<'a> ExactSizeIterator for SmallIntViewIter<'a> {
254    fn len(&self) -> usize {
255        self.view.len() - self.row
256    }
257}
258
259impl<A: Into<Option<Item>>> FromIterator<A> for View {
260    fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> Self {
261        let (nulls, mut values): (Vec<bool>, Vec<_>) = iter
262            .into_iter()
263            .map(|v| match v.into() {
264                Some(v) => (false, v),
265                None => (true, Item::default()),
266            })
267            .unzip();
268        Self {
269            nulls: NullBits::from_iter(nulls),
270            data: Bytes::from({
271                let (ptr, len, cap) = (values.as_mut_ptr(), values.len(), values.capacity());
272                std::mem::forget(values);
273                unsafe { Vec::from_raw_parts(ptr as *mut u8, len * ITEM_SIZE, cap * ITEM_SIZE) }
274            }),
275        }
276    }
277}
278
279#[test]
280fn test_slice() {
281    let data = [0, 1, Item::MIN, Item::MAX];
282    let view = View::from_iter(data);
283    dbg!(&view);
284    let slice = view.slice(1..3);
285    dbg!(&slice);
286
287    let data = [None, Some(Item::MIN), Some(Item::MAX), None];
288    let view = View::from_iter(data);
289    dbg!(&view);
290    let range = 1..4;
291    let slice = view.slice(range.clone()).unwrap();
292    for (v, i) in slice.iter().zip(range) {
293        assert_eq!(v, data[i]);
294    }
295}
taos_query/common/raw/views/smallint_view.rs

taos_query/common/raw/views/
smallint_view.rs
