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

taos_query/common/raw/views/
json_view.rs
