taos_query/common/raw/views/
varbinary_view.rs

1use std::{borrow::Cow, ffi::c_void, fmt::Debug};
2
3use super::{IsColumnView, Offsets};
4use crate::{
5    common::{BorrowedValue, Ty},
6    prelude::InlinableWrite,
7    util::InlineBytes,
8};
9
10use bytes::Bytes;
11use itertools::Itertools;
12
13#[derive(Debug, Clone)]
14pub struct VarBinaryView {
15    // version: Version,
16    pub(crate) offsets: Offsets,
17    pub(crate) data: Bytes,
18}
19
20impl IsColumnView for VarBinaryView {
21    fn ty(&self) -> Ty {
22        Ty::VarBinary
23    }
24    fn from_borrowed_value_iter<'b>(iter: impl Iterator<Item = BorrowedValue<'b>>) -> Self {
25        Self::from_iter::<Bytes, _, _, _>(iter.map(|v| v.to_bytes()).collect_vec())
26    }
27}
28
29impl VarBinaryView {
30    pub fn len(&self) -> usize {
31        self.offsets.len()
32    }
33
34    /// A iterator only decide if the value at some row index is NULL or not.
35    pub fn is_null_iter(&self) -> VarBinaryNullsIter {
36        VarBinaryNullsIter { view: self, row: 0 }
37    }
38
39    /// Build a nulls vector.
40    pub fn to_nulls_vec(&self) -> Vec<bool> {
41        self.is_null_iter().collect()
42    }
43
44    /// Check if the value at `row` index is NULL or not.
45    ///
46    /// Returns null when `row` index out of bound.
47    pub fn is_null(&self, row: usize) -> bool {
48        if row < self.len() {
49            unsafe { self.is_null_unchecked(row) }
50        } else {
51            false
52        }
53    }
54
55    /// Unsafe version for [is_null](#method.is_null)
56    pub(crate) unsafe fn is_null_unchecked(&self, row: usize) -> bool {
57        self.offsets.get_unchecked(row) < 0
58    }
59
60    pub(crate) unsafe fn get_unchecked(&self, row: usize) -> Option<&InlineBytes> {
61        let offset = self.offsets.get_unchecked(row);
62        if offset >= 0 {
63            Some(InlineBytes::<u16>::from_ptr(
64                self.data.as_ptr().offset(offset as isize),
65            ))
66        } else {
67            None
68        }
69    }
70
71    pub(crate) unsafe fn get_value_unchecked(&self, row: usize) -> BorrowedValue {
72        self.get_unchecked(row)
73            .map(|s| BorrowedValue::VarBinary(Cow::Borrowed(s.as_bytes())))
74            .unwrap_or(BorrowedValue::Null(Ty::VarBinary))
75    }
76
77    pub(crate) unsafe fn get_raw_value_unchecked(&self, row: usize) -> (Ty, u32, *const c_void) {
78        match self.get_unchecked(row) {
79            Some(s) => (Ty::VarBinary, s.len() as _, s.as_ptr() as _),
80            None => (Ty::VarBinary, 0, std::ptr::null()),
81        }
82    }
83
84    #[inline]
85    pub unsafe fn get_length_unchecked(&self, row: usize) -> Option<usize> {
86        let offset = self.offsets.get_unchecked(row);
87        if offset >= 0 {
88            Some(InlineBytes::<u16>::from_ptr(self.data.as_ptr().offset(offset as isize)).len())
89        } else {
90            None
91        }
92    }
93
94    #[inline]
95    pub fn lengths(&self) -> Vec<Option<usize>> {
96        (0..self.len())
97            .map(|i| unsafe { self.get_length_unchecked(i) })
98            .collect()
99    }
100
101    #[inline]
102    pub fn max_length(&self) -> usize {
103        (0..self.len())
104            .filter_map(|i| unsafe { self.get_length_unchecked(i) })
105            .min()
106            .unwrap_or(0)
107    }
108
109    pub fn iter(&self) -> VarBinaryIter {
110        VarBinaryIter { view: self, row: 0 }
111    }
112
113    pub fn to_vec(&self) -> Vec<Option<Vec<u8>>> {
114        (0..self.len())
115            .map(|row| {
116                unsafe { self.get_unchecked(row) }
117                    .map(|s| s.as_bytes())
118                    .map(|s| s.to_vec())
119            })
120            .collect()
121    }
122    // pub fn iter_as_bytes(&self) -> impl Iterator<Item = Option<&[u8]>> {
123    //     (0..self.len()).map(|row| unsafe { self.get_unchecked(row) }.map(|s| s.as_bytes()))
124    // }
125
126    // pub fn to_bytes_vec(&self) -> Vec<Option<&[u8]>> {
127    //     self.iter_as_bytes().collect_vec()
128    // }
129
130    /// Write column data as raw bytes.
131    pub(crate) fn write_raw_into<W: std::io::Write>(&self, mut wtr: W) -> std::io::Result<usize> {
132        let mut offsets = Vec::with_capacity(self.len());
133        let mut bytes: Vec<u8> = Vec::new();
134        for v in self.iter() {
135            if let Some(v) = v {
136                offsets.push(bytes.len() as i32);
137                bytes.write_inlined_bytes::<2>(v.as_bytes()).unwrap();
138            } else {
139                offsets.push(-1);
140            }
141        }
142        unsafe {
143            let offsets_bytes = std::slice::from_raw_parts(
144                offsets.as_ptr() as *const u8,
145                offsets.len() * std::mem::size_of::<i32>(),
146            );
147            wtr.write_all(offsets_bytes)?;
148            wtr.write_all(&bytes)?;
149            Ok(offsets_bytes.len() + bytes.len())
150        }
151        // let offsets = self.offsets.as_bytes();
152        // dbg!(self, offsets);
153        // wtr.write_all(offsets)?;
154        // wtr.write_all(&self.data)?;
155        // Ok(offsets.len() + self.data.len())
156    }
157
158    pub fn from_iter<
159        S: AsRef<[u8]>,
160        T: Into<Option<S>>,
161        I: ExactSizeIterator<Item = T>,
162        V: IntoIterator<Item = T, IntoIter = I>,
163    >(
164        iter: V,
165    ) -> Self {
166        let iter = iter.into_iter();
167        let mut offsets = Vec::with_capacity(iter.len());
168        let mut data = Vec::new();
169
170        for i in iter.map(|v| v.into()) {
171            if let Some(s) = i {
172                let s: &[u8] = s.as_ref();
173                offsets.push(data.len() as i32);
174                data.write_inlined_bytes::<2>(s).unwrap();
175            } else {
176                offsets.push(-1);
177            }
178        }
179        // dbg!(&offsets);
180        let offsets_bytes = unsafe {
181            Vec::from_raw_parts(
182                offsets.as_mut_ptr() as *mut u8,
183                offsets.len() * 4,
184                offsets.capacity() * 4,
185            )
186        };
187        std::mem::forget(offsets);
188        Self {
189            offsets: Offsets(offsets_bytes.into()),
190            data: data.into(),
191        }
192    }
193
194    // pub fn concat(&self, rhs: &Self) -> Self {
195    //     Self::from_iter::<&InlineJson, _, _, _>(self.iter().chain(rhs.iter()).collect_vec())
196    // }
197}
198
199pub struct VarBinaryIter<'a> {
200    view: &'a VarBinaryView,
201    row: usize,
202}
203
204impl<'a> Iterator for VarBinaryIter<'a> {
205    type Item = Option<&'a InlineBytes>;
206
207    fn next(&mut self) -> Option<Self::Item> {
208        if self.row < self.view.len() {
209            let row = self.row;
210            self.row += 1;
211            Some(unsafe { self.view.get_unchecked(row) })
212        } else {
213            None
214        }
215    }
216
217    fn size_hint(&self) -> (usize, Option<usize>) {
218        let len = self.view.len() - self.row;
219        (len, Some(len))
220    }
221}
222
223impl<'a> ExactSizeIterator for VarBinaryIter<'a> {}
224
225pub struct VarBinaryNullsIter<'a> {
226    view: &'a VarBinaryView,
227    row: usize,
228}
229
230impl<'a> Iterator for VarBinaryNullsIter<'a> {
231    type Item = bool;
232
233    fn next(&mut self) -> Option<Self::Item> {
234        if self.row < self.view.len() {
235            let row = self.row;
236            self.row += 1;
237            Some(unsafe { self.view.is_null_unchecked(row) })
238        } else {
239            None
240        }
241    }
242}
243
244impl<'a> ExactSizeIterator for VarBinaryNullsIter<'a> {
245    fn len(&self) -> usize {
246        self.view.len() - self.row
247    }
248}
249
250// #[test]
251// fn test_slice() {
252//     let data = [None, Some(""), Some("abc"), Some("中文"), None, None, Some("a loooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooog string")];
253//     let view = VarBinaryView::from_iter::<&str, _, _, _>(data);
254//     let slice = view.slice(0..0);
255//     assert!(slice.is_none());
256//     let slice = view.slice(100..1000);
257//     assert!(slice.is_none());
258
259//     for start in 0..data.len() {
260//         let end = start + 1;
261//         for end in end..data.len() {
262//             let slice = view.slice(start..end).unwrap();
263//             assert_eq!(
264//                 slice.to_vec().as_slice(),
265//                 &data[start..end]
266//                     .iter()
267//                     .map(|s| s.map(ToString::to_string))
268//                     .collect_vec()
269//             );
270//         }
271//     }
272// }
taos_query/common/raw/views/varbinary_view.rs

taos_query/common/raw/views/
varbinary_view.rs
