opensrv_clickhouse/types/column/
array.rs

1// Copyright 2021 Datafuse Labs.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use chrono_tz::Tz;
18
19use crate::binary::Encoder;
20use crate::binary::ReadEx;
21use crate::errors::Result;
22use crate::types::column::column_data::ArcColumnData;
23use crate::types::column::column_data::BoxColumnData;
24use crate::types::column::list::List;
25use crate::types::column::ArcColumnWrapper;
26use crate::types::column::ColumnData;
27use crate::types::SqlType;
28use crate::types::Value;
29use crate::types::ValueRef;
30
31pub struct ArrayColumnData {
32    pub(crate) inner: ArcColumnData,
33    pub(crate) offsets: List<u64>,
34}
35
36impl ArrayColumnData {
37    #[allow(dead_code)]
38    pub fn create(inner: ArcColumnData, offsets: List<u64>) -> Self {
39        ArrayColumnData { inner, offsets }
40    }
41
42    pub(crate) fn load<R: ReadEx>(
43        reader: &mut R,
44        type_name: &str,
45        rows: usize,
46        tz: Tz,
47    ) -> Result<Self> {
48        let mut offsets = List::with_capacity(rows);
49        offsets.resize(rows, 0_u64);
50        reader.read_bytes(offsets.as_mut())?;
51
52        let size = match rows {
53            0 => 0,
54            _ => offsets.at(rows - 1) as usize,
55        };
56        let inner =
57            <dyn ColumnData>::load_data::<ArcColumnWrapper, _>(reader, type_name, size, tz)?;
58
59        Ok(ArrayColumnData { inner, offsets })
60    }
61}
62
63impl ColumnData for ArrayColumnData {
64    fn sql_type(&self) -> SqlType {
65        let inner_type = self.inner.sql_type();
66        SqlType::Array(inner_type.into())
67    }
68
69    fn save(&self, encoder: &mut Encoder, start: usize, end: usize) {
70        let mut offset = 0_u64;
71
72        for i in start..end {
73            offset = self.offsets.at(i);
74            encoder.write(offset);
75        }
76
77        self.inner.save(encoder, 0, offset as usize);
78    }
79
80    fn len(&self) -> usize {
81        self.offsets.len()
82    }
83
84    fn push(&mut self, value: Value) {
85        if let Value::Array(_, vs) = value {
86            let offsets_len = self.offsets.len();
87            let prev = if offsets_len == 0 {
88                0_usize
89            } else {
90                self.offsets.at(offsets_len - 1) as usize
91            };
92
93            let inner_column = Arc::get_mut(&mut self.inner).unwrap();
94            self.offsets.push((prev + vs.len()) as u64);
95            for v in vs.iter() {
96                inner_column.push(v.clone());
97            }
98        } else {
99            panic!("value should be an array")
100        }
101    }
102
103    fn at(&self, index: usize) -> ValueRef {
104        let sql_type = self.inner.sql_type();
105
106        let start = if index > 0 {
107            self.offsets.at(index - 1) as usize
108        } else {
109            0_usize
110        };
111        let end = self.offsets.at(index) as usize;
112        let mut vs = Vec::with_capacity(end);
113        for i in start..end {
114            let v = self.inner.at(i);
115            vs.push(v);
116        }
117        ValueRef::Array(sql_type.into(), Arc::new(vs))
118    }
119
120    fn clone_instance(&self) -> BoxColumnData {
121        Box::new(Self {
122            inner: self.inner.clone(),
123            offsets: self.offsets.clone(),
124        })
125    }
126
127    unsafe fn get_internal(&self, pointers: &[*mut *const u8], level: u8) -> Result<()> {
128        if level == self.sql_type().level() {
129            *pointers[0] = self.offsets.as_ptr() as *const u8;
130            *(pointers[1] as *mut usize) = self.offsets.len();
131            Ok(())
132        } else {
133            self.inner.get_internal(pointers, level)
134        }
135    }
136
137    fn cast_to(&self, _this: &ArcColumnData, target: &SqlType) -> Option<ArcColumnData> {
138        if let SqlType::Array(inner_target) = target {
139            if let Some(inner) = self.inner.cast_to(&self.inner, inner_target) {
140                return Some(Arc::new(ArrayColumnData {
141                    inner,
142                    offsets: self.offsets.clone(),
143                }));
144            }
145        }
146        None
147    }
148}
opensrv_clickhouse/types/column/array.rs

opensrv_clickhouse/types/column/
array.rs
