opensrv_clickhouse/types/column/
enums.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::BoxColumnData;
23use crate::types::column::column_data::ColumnData;
24use crate::types::column::list::List;
25use crate::types::column::nullable::NullableColumnData;
26use crate::types::column::BoxColumnWrapper;
27use crate::types::column::ColumnFrom;
28use crate::types::column::ColumnWrapper;
29use crate::types::column::Either;
30use crate::types::column::VectorColumnData;
31use crate::types::enums::Enum16;
32use crate::types::enums::Enum8;
33use crate::types::from_sql::FromSql;
34use crate::types::Column;
35use crate::types::ColumnType;
36use crate::types::SqlType;
37use crate::types::Value;
38use crate::types::ValueRef;
39
40pub(crate) struct Enum16ColumnData {
41    pub(crate) enum_values: Vec<(String, i16)>,
42    pub(crate) inner: Box<dyn ColumnData + Send + Sync>,
43}
44
45pub(crate) struct Enum16Adapter<K: ColumnType> {
46    pub(crate) column: Column<K>,
47    pub(crate) enum_values: Vec<(String, i16)>,
48}
49
50pub(crate) struct NullableEnum16Adapter<K: ColumnType> {
51    pub(crate) column: Column<K>,
52    pub(crate) enum_values: Vec<(String, i16)>,
53}
54
55impl Enum16ColumnData {
56    pub(crate) fn load<T: ReadEx>(
57        reader: &mut T,
58        enum_values: Vec<(String, i16)>,
59        size: usize,
60        tz: Tz,
61    ) -> Result<Self> {
62        let type_name = "Int16";
63
64        let inner =
65            <dyn ColumnData>::load_data::<BoxColumnWrapper, _>(reader, type_name, size, tz)?;
66
67        Ok(Enum16ColumnData { enum_values, inner })
68    }
69}
70
71impl ColumnData for Enum16ColumnData {
72    fn sql_type(&self) -> SqlType {
73        SqlType::Enum16(self.enum_values.clone())
74    }
75
76    fn save(&self, encoder: &mut Encoder, start: usize, end: usize) {
77        self.inner.save(encoder, start, end)
78    }
79
80    fn len(&self) -> usize {
81        self.inner.len()
82    }
83
84    fn push(&mut self, value: Value) {
85        if let Value::Enum16(_values, enum_value) = value {
86            self.inner.push(Value::Int16(enum_value.internal()))
87        } else {
88            panic!("value should be Enum ({:?})", value);
89        }
90    }
91
92    fn at(&self, index: usize) -> ValueRef {
93        let enum_value = i16::from(self.inner.at(index));
94        ValueRef::Enum16(self.enum_values.clone(), Enum16(enum_value))
95    }
96
97    fn clone_instance(&self) -> BoxColumnData {
98        Box::new(Self {
99            inner: self.inner.clone_instance(),
100            enum_values: self.enum_values.clone(),
101        })
102    }
103}
104
105impl<K: ColumnType> ColumnData for Enum16Adapter<K> {
106    fn sql_type(&self) -> SqlType {
107        SqlType::Enum16(self.enum_values.clone())
108    }
109
110    fn save(&self, encoder: &mut Encoder, start: usize, end: usize) {
111        for i in start..end {
112            if let ValueRef::Enum16(_enum_values, value) = self.at(i) {
113                encoder.write(value.internal())
114            } else {
115                panic!("should be Enum");
116            }
117        }
118    }
119
120    fn len(&self) -> usize {
121        self.column.len()
122    }
123
124    fn push(&mut self, _: Value) {
125        unimplemented!()
126    }
127
128    fn at(&self, index: usize) -> ValueRef {
129        if let ValueRef::Enum16(enum_values, value) = self.column.at(index) {
130            ValueRef::Enum16(enum_values, value)
131        } else {
132            panic!("should be Enum");
133        }
134    }
135
136    fn clone_instance(&self) -> BoxColumnData {
137        unimplemented!()
138    }
139}
140
141impl<K: ColumnType> ColumnData for NullableEnum16Adapter<K> {
142    fn sql_type(&self) -> SqlType {
143        SqlType::Nullable(SqlType::Enum16(self.enum_values.clone()).into())
144    }
145
146    fn save(&self, encoder: &mut Encoder, start: usize, end: usize) {
147        let size = end - start;
148        let mut nulls = vec![0; size];
149        let mut values: Vec<Option<i16>> = vec![None; size];
150
151        for (i, index) in (start..end).enumerate() {
152            let value: Option<Enum16> = Option::from_sql(self.at(index)).unwrap();
153            match value {
154                Some(v) => values[i] = Some(v.internal()),
155                None => nulls[i] = 1,
156            }
157        }
158        encoder.write_bytes(nulls.as_ref());
159
160        for value in values {
161            let data_value = value.unwrap_or(0);
162            encoder.write(data_value);
163        }
164    }
165
166    fn len(&self) -> usize {
167        self.column.len()
168    }
169
170    fn push(&mut self, _: Value) {
171        unimplemented!()
172    }
173
174    fn at(&self, index: usize) -> ValueRef {
175        let value: Option<Enum16> = Option::from_sql(self.column.at(index)).unwrap();
176        match value {
177            None => ValueRef::Nullable(Either::Left(self.sql_type().into())),
178            Some(v) => {
179                let inner = ValueRef::Enum16(self.enum_values.clone(), v);
180                ValueRef::Nullable(Either::Right(Box::new(inner)))
181            }
182        }
183    }
184
185    fn clone_instance(&self) -> BoxColumnData {
186        unimplemented!()
187    }
188}
189
190impl ColumnFrom for Vec<Enum16> {
191    fn column_from<W: ColumnWrapper>(source: Self) -> W::Wrapper {
192        let mut data = List::<i16>::with_capacity(source.len());
193
194        for s in source {
195            data.push(s.internal());
196        }
197        let inner = Box::new(VectorColumnData { data });
198
199        let column = Enum16ColumnData {
200            inner,
201            enum_values: vec![],
202        };
203
204        W::wrap(column)
205    }
206}
207
208impl ColumnFrom for Vec<Option<Enum16>> {
209    fn column_from<W: ColumnWrapper>(source: Self) -> W::Wrapper {
210        let mut nulls: Vec<u8> = Vec::with_capacity(source.len());
211        let mut data = List::<i16>::with_capacity(source.len());
212
213        for os in source {
214            if let Some(s) = os {
215                data.push(s.internal());
216                nulls.push(0);
217            } else {
218                data.push(0);
219                nulls.push(1);
220            }
221        }
222        let inner = Box::new(VectorColumnData { data });
223
224        let inner = Enum16ColumnData {
225            enum_values: vec![],
226            inner,
227        };
228
229        W::wrap(NullableColumnData {
230            nulls,
231            inner: Arc::new(inner),
232        })
233    }
234}
235
236pub(crate) struct Enum8ColumnData {
237    pub(crate) enum_values: Vec<(String, i8)>,
238    pub(crate) inner: Box<dyn ColumnData + Send + Sync>,
239}
240
241pub(crate) struct Enum8Adapter<K: ColumnType> {
242    pub(crate) column: Column<K>,
243    pub(crate) enum_values: Vec<(String, i8)>,
244}
245
246pub(crate) struct NullableEnum8Adapter<K: ColumnType> {
247    pub(crate) column: Column<K>,
248    pub(crate) enum_values: Vec<(String, i8)>,
249}
250
251impl Enum8ColumnData {
252    pub(crate) fn load<T: ReadEx>(
253        reader: &mut T,
254        enum_values: Vec<(String, i8)>,
255        size: usize,
256        tz: Tz,
257    ) -> Result<Self> {
258        let type_name = "Int8";
259
260        let inner =
261            <dyn ColumnData>::load_data::<BoxColumnWrapper, _>(reader, type_name, size, tz)?;
262
263        Ok(Enum8ColumnData { enum_values, inner })
264    }
265}
266
267impl ColumnData for Enum8ColumnData {
268    fn sql_type(&self) -> SqlType {
269        SqlType::Enum8(self.enum_values.clone())
270    }
271
272    fn save(&self, encoder: &mut Encoder, start: usize, end: usize) {
273        self.inner.save(encoder, start, end)
274    }
275
276    fn len(&self) -> usize {
277        self.inner.len()
278    }
279
280    fn push(&mut self, value: Value) {
281        if let Value::Enum8(_values, enum_value) = value {
282            self.inner.push(Value::Int8(enum_value.internal()))
283        } else {
284            panic!("value should be Enum ({:?})", value);
285        }
286    }
287
288    fn at(&self, index: usize) -> ValueRef {
289        let enum_value = i8::from(self.inner.at(index));
290        ValueRef::Enum8(self.enum_values.clone(), Enum8(enum_value))
291    }
292
293    fn clone_instance(&self) -> BoxColumnData {
294        Box::new(Self {
295            inner: self.inner.clone_instance(),
296            enum_values: self.enum_values.clone(),
297        })
298    }
299}
300
301impl<K: ColumnType> ColumnData for Enum8Adapter<K> {
302    fn sql_type(&self) -> SqlType {
303        SqlType::Enum8(self.enum_values.clone())
304    }
305
306    fn save(&self, encoder: &mut Encoder, start: usize, end: usize) {
307        for i in start..end {
308            if let ValueRef::Enum8(_enum_values, value) = self.at(i) {
309                encoder.write(value.internal())
310            } else {
311                panic!("should be Enum");
312            }
313        }
314    }
315
316    fn len(&self) -> usize {
317        self.column.len()
318    }
319
320    fn push(&mut self, _: Value) {
321        unimplemented!()
322    }
323
324    fn at(&self, index: usize) -> ValueRef {
325        if let ValueRef::Enum8(enum_values, value) = self.column.at(index) {
326            ValueRef::Enum8(enum_values, value)
327        } else {
328            panic!("should be Enum");
329        }
330    }
331
332    fn clone_instance(&self) -> BoxColumnData {
333        unimplemented!()
334    }
335}
336
337impl<K: ColumnType> ColumnData for NullableEnum8Adapter<K> {
338    fn sql_type(&self) -> SqlType {
339        SqlType::Nullable(SqlType::Enum8(self.enum_values.clone()).into())
340    }
341
342    fn save(&self, encoder: &mut Encoder, start: usize, end: usize) {
343        let size = end - start;
344        let mut nulls = vec![0; size];
345        let mut values: Vec<Option<i8>> = vec![None; size];
346
347        for (i, index) in (start..end).enumerate() {
348            let value: Option<Enum8> = Option::from_sql(self.at(index)).unwrap();
349            match value {
350                Some(v) => values[i] = Some(v.internal()),
351                None => nulls[i] = 1,
352            }
353        }
354
355        encoder.write_bytes(nulls.as_ref());
356
357        for value in values {
358            let data_value = value.unwrap_or(0);
359            encoder.write(data_value);
360        }
361    }
362
363    fn len(&self) -> usize {
364        self.column.len()
365    }
366
367    fn push(&mut self, _: Value) {
368        unimplemented!()
369    }
370
371    fn at(&self, index: usize) -> ValueRef {
372        let value: Option<Enum8> = Option::from_sql(self.column.at(index)).unwrap();
373        match value {
374            None => ValueRef::Nullable(Either::Left(self.sql_type().into())),
375            Some(v) => {
376                let inner = ValueRef::Enum8(self.enum_values.clone(), v);
377                ValueRef::Nullable(Either::Right(Box::new(inner)))
378            }
379        }
380    }
381
382    fn clone_instance(&self) -> BoxColumnData {
383        unimplemented!()
384    }
385}
386
387impl ColumnFrom for Vec<Enum8> {
388    fn column_from<W: ColumnWrapper>(source: Self) -> W::Wrapper {
389        let mut data = List::<i8>::with_capacity(source.len());
390
391        for s in source {
392            data.push(s.internal());
393        }
394        let inner = Box::new(VectorColumnData { data });
395
396        let column = Enum8ColumnData {
397            inner,
398            enum_values: vec![],
399        };
400
401        W::wrap(column)
402    }
403}
404
405impl ColumnFrom for Vec<Option<Enum8>> {
406    fn column_from<W: ColumnWrapper>(source: Self) -> W::Wrapper {
407        let mut nulls: Vec<u8> = Vec::with_capacity(source.len());
408        let mut data = List::<i8>::with_capacity(source.len());
409
410        for os in source {
411            if let Some(s) = os {
412                data.push(s.internal());
413                nulls.push(0);
414            } else {
415                data.push(0);
416                nulls.push(1);
417            }
418        }
419        let inner = Box::new(VectorColumnData { data });
420
421        let inner = Enum8ColumnData {
422            enum_values: vec![],
423            inner,
424        };
425
426        W::wrap(NullableColumnData {
427            nulls,
428            inner: Arc::new(inner),
429        })
430    }
431}