taos_query/common/raw/views/
var_char_view.rs

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

taos_query/common/raw/views/
var_char_view.rs
