taos_query/common/raw/views/
big_int_unsigned_view.rs

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

taos_query/common/raw/views/
big_int_unsigned_view.rs
