Skip to main content

clickhouse_native_client/column/
enum_column.rs

1//! Enum8 and Enum16 column implementations.
2//!
3//! ClickHouse enums map string names to integer values. `Enum8` uses `Int8`
4//! storage (up to 127 distinct values) and `Enum16` uses `Int16` storage
5//! (up to 32767 distinct values). The name-to-value mapping is stored in
6//! the column's [`Type`].
7
8use super::{
9    Column,
10    ColumnRef,
11};
12use crate::{
13    types::Type,
14    Error,
15    Result,
16};
17use bytes::BytesMut;
18use std::sync::Arc;
19
20/// Column for Enum8 type (stored as Int8 with name-value mapping in Type).
21pub struct ColumnEnum8 {
22    type_: Type,
23    data: Vec<i8>,
24}
25
26impl ColumnEnum8 {
27    /// Create a new empty Enum8 column.
28    ///
29    /// # Panics
30    ///
31    /// Panics if `type_` is not `Type::Enum8`.
32    pub fn new(type_: Type) -> Self {
33        match &type_ {
34            Type::Enum8 { .. } => Self { type_, data: Vec::new() },
35            _ => panic!("ColumnEnum8 requires Enum8 type"),
36        }
37    }
38
39    /// Set the column data from a vector of raw `i8` enum values.
40    pub fn with_data(mut self, data: Vec<i8>) -> Self {
41        self.data = data;
42        self
43    }
44
45    /// Append enum by numeric value
46    pub fn append_value(&mut self, value: i8) {
47        self.data.push(value);
48    }
49
50    /// Append enum by name (looks up value in Type).
51    ///
52    /// # Errors
53    ///
54    /// Returns an error if `name` is not a known variant in this enum type.
55    pub fn append_name(&mut self, name: &str) -> Result<()> {
56        let value = self.type_.get_enum_value(name).ok_or_else(|| {
57            Error::Protocol(format!("Unknown enum name: {}", name))
58        })?;
59
60        self.data.push(value as i8);
61        Ok(())
62    }
63
64    /// Get numeric value at index.
65    pub fn at(&self, index: usize) -> i8 {
66        self.data[index]
67    }
68
69    /// Get enum name at index (looks up in Type).
70    pub fn name_at(&self, index: usize) -> Option<&str> {
71        let value = self.data[index] as i16;
72        self.type_.get_enum_name(value)
73    }
74
75    /// Returns the number of values in this column.
76    pub fn len(&self) -> usize {
77        self.data.len()
78    }
79
80    /// Returns `true` if the column contains no values.
81    pub fn is_empty(&self) -> bool {
82        self.data.is_empty()
83    }
84}
85
86impl Column for ColumnEnum8 {
87    fn column_type(&self) -> &Type {
88        &self.type_
89    }
90
91    fn size(&self) -> usize {
92        self.data.len()
93    }
94
95    fn clear(&mut self) {
96        self.data.clear();
97    }
98
99    fn reserve(&mut self, new_cap: usize) {
100        self.data.reserve(new_cap);
101    }
102
103    fn append_column(&mut self, other: ColumnRef) -> Result<()> {
104        let other =
105            other.as_any().downcast_ref::<ColumnEnum8>().ok_or_else(|| {
106                Error::TypeMismatch {
107                    expected: self.type_.name(),
108                    actual: other.column_type().name(),
109                }
110            })?;
111
112        self.data.extend_from_slice(&other.data);
113        Ok(())
114    }
115
116    fn load_from_buffer(
117        &mut self,
118        buffer: &mut &[u8],
119        rows: usize,
120    ) -> Result<()> {
121        let bytes_needed = rows;
122        if buffer.len() < bytes_needed {
123            return Err(Error::Protocol(format!(
124                "Buffer underflow: need {} bytes for Enum8, have {}",
125                bytes_needed,
126                buffer.len()
127            )));
128        }
129
130        // Use bulk copy for performance
131        self.data.reserve(rows);
132        let current_len = self.data.len();
133        unsafe {
134            // Set length first to claim ownership of the memory
135            self.data.set_len(current_len + rows);
136            let dest_ptr =
137                (self.data.as_mut_ptr() as *mut u8).add(current_len);
138            std::ptr::copy_nonoverlapping(
139                buffer.as_ptr(),
140                dest_ptr,
141                bytes_needed,
142            );
143        }
144
145        use bytes::Buf;
146        buffer.advance(bytes_needed);
147        Ok(())
148    }
149
150    fn save_to_buffer(&self, buffer: &mut BytesMut) -> Result<()> {
151        if !self.data.is_empty() {
152            let byte_slice = unsafe {
153                std::slice::from_raw_parts(
154                    self.data.as_ptr() as *const u8,
155                    self.data.len(),
156                )
157            };
158            buffer.extend_from_slice(byte_slice);
159        }
160        Ok(())
161    }
162
163    fn clone_empty(&self) -> ColumnRef {
164        Arc::new(ColumnEnum8::new(self.type_.clone()))
165    }
166
167    fn slice(&self, begin: usize, len: usize) -> Result<ColumnRef> {
168        if begin + len > self.data.len() {
169            return Err(Error::InvalidArgument(format!(
170                "Slice out of bounds: begin={}, len={}, size={}",
171                begin,
172                len,
173                self.data.len()
174            )));
175        }
176
177        let sliced_data = self.data[begin..begin + len].to_vec();
178        Ok(Arc::new(
179            ColumnEnum8::new(self.type_.clone()).with_data(sliced_data),
180        ))
181    }
182
183    fn as_any(&self) -> &dyn std::any::Any {
184        self
185    }
186
187    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
188        self
189    }
190}
191
192/// Column for Enum16 type (stored as Int16 with name-value mapping in Type).
193pub struct ColumnEnum16 {
194    type_: Type,
195    data: Vec<i16>,
196}
197
198impl ColumnEnum16 {
199    /// Create a new empty Enum16 column.
200    ///
201    /// # Panics
202    ///
203    /// Panics if `type_` is not `Type::Enum16`.
204    pub fn new(type_: Type) -> Self {
205        match &type_ {
206            Type::Enum16 { .. } => Self { type_, data: Vec::new() },
207            _ => panic!("ColumnEnum16 requires Enum16 type"),
208        }
209    }
210
211    /// Set the column data from a vector of raw `i16` enum values.
212    pub fn with_data(mut self, data: Vec<i16>) -> Self {
213        self.data = data;
214        self
215    }
216
217    /// Append enum by numeric value.
218    pub fn append_value(&mut self, value: i16) {
219        self.data.push(value);
220    }
221
222    /// Append enum by name (looks up value in Type).
223    ///
224    /// # Errors
225    ///
226    /// Returns an error if `name` is not a known variant in this enum type.
227    pub fn append_name(&mut self, name: &str) -> Result<()> {
228        let value = self.type_.get_enum_value(name).ok_or_else(|| {
229            Error::Protocol(format!("Unknown enum name: {}", name))
230        })?;
231
232        self.data.push(value);
233        Ok(())
234    }
235
236    /// Get numeric value at index.
237    pub fn at(&self, index: usize) -> i16 {
238        self.data[index]
239    }
240
241    /// Get enum name at index (looks up in Type).
242    pub fn name_at(&self, index: usize) -> Option<&str> {
243        let value = self.data[index];
244        self.type_.get_enum_name(value)
245    }
246
247    /// Returns the number of values in this column.
248    pub fn len(&self) -> usize {
249        self.data.len()
250    }
251
252    /// Returns `true` if the column contains no values.
253    pub fn is_empty(&self) -> bool {
254        self.data.is_empty()
255    }
256}
257
258impl Column for ColumnEnum16 {
259    fn column_type(&self) -> &Type {
260        &self.type_
261    }
262
263    fn size(&self) -> usize {
264        self.data.len()
265    }
266
267    fn clear(&mut self) {
268        self.data.clear();
269    }
270
271    fn reserve(&mut self, new_cap: usize) {
272        self.data.reserve(new_cap);
273    }
274
275    fn append_column(&mut self, other: ColumnRef) -> Result<()> {
276        let other = other.as_any().downcast_ref::<ColumnEnum16>().ok_or_else(
277            || Error::TypeMismatch {
278                expected: self.type_.name(),
279                actual: other.column_type().name(),
280            },
281        )?;
282
283        self.data.extend_from_slice(&other.data);
284        Ok(())
285    }
286
287    fn load_from_buffer(
288        &mut self,
289        buffer: &mut &[u8],
290        rows: usize,
291    ) -> Result<()> {
292        let bytes_needed = rows * 2;
293        if buffer.len() < bytes_needed {
294            return Err(Error::Protocol(format!(
295                "Buffer underflow: need {} bytes for Enum16, have {}",
296                bytes_needed,
297                buffer.len()
298            )));
299        }
300
301        // Use bulk copy for performance
302        self.data.reserve(rows);
303        let current_len = self.data.len();
304        unsafe {
305            // Set length first to claim ownership of the memory
306            self.data.set_len(current_len + rows);
307            let dest_ptr =
308                (self.data.as_mut_ptr() as *mut u8).add(current_len * 2);
309            std::ptr::copy_nonoverlapping(
310                buffer.as_ptr(),
311                dest_ptr,
312                bytes_needed,
313            );
314        }
315
316        use bytes::Buf;
317        buffer.advance(bytes_needed);
318        Ok(())
319    }
320
321    fn save_to_buffer(&self, buffer: &mut BytesMut) -> Result<()> {
322        if !self.data.is_empty() {
323            let byte_slice = unsafe {
324                std::slice::from_raw_parts(
325                    self.data.as_ptr() as *const u8,
326                    self.data.len() * 2,
327                )
328            };
329            buffer.extend_from_slice(byte_slice);
330        }
331        Ok(())
332    }
333
334    fn clone_empty(&self) -> ColumnRef {
335        Arc::new(ColumnEnum16::new(self.type_.clone()))
336    }
337
338    fn slice(&self, begin: usize, len: usize) -> Result<ColumnRef> {
339        if begin + len > self.data.len() {
340            return Err(Error::InvalidArgument(format!(
341                "Slice out of bounds: begin={}, len={}, size={}",
342                begin,
343                len,
344                self.data.len()
345            )));
346        }
347
348        let sliced_data = self.data[begin..begin + len].to_vec();
349        Ok(Arc::new(
350            ColumnEnum16::new(self.type_.clone()).with_data(sliced_data),
351        ))
352    }
353
354    fn as_any(&self) -> &dyn std::any::Any {
355        self
356    }
357
358    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
359        self
360    }
361}
362
363#[cfg(test)]
364#[cfg_attr(coverage_nightly, coverage(off))]
365mod tests {
366    use super::*;
367    use crate::types::EnumItem;
368
369    #[test]
370    fn test_enum8_append_value() {
371        let items = vec![
372            EnumItem { name: "Red".to_string(), value: 1 },
373            EnumItem { name: "Green".to_string(), value: 2 },
374        ];
375        let mut col = ColumnEnum8::new(Type::enum8(items));
376
377        col.append_value(1);
378        col.append_value(2);
379
380        assert_eq!(col.len(), 2);
381        assert_eq!(col.at(0), 1);
382        assert_eq!(col.at(1), 2);
383    }
384
385    #[test]
386    fn test_enum8_append_name() {
387        let items = vec![
388            EnumItem { name: "Red".to_string(), value: 1 },
389            EnumItem { name: "Green".to_string(), value: 2 },
390        ];
391        let mut col = ColumnEnum8::new(Type::enum8(items));
392
393        col.append_name("Red").unwrap();
394        col.append_name("Green").unwrap();
395
396        assert_eq!(col.len(), 2);
397        assert_eq!(col.at(0), 1);
398        assert_eq!(col.at(1), 2);
399        assert_eq!(col.name_at(0), Some("Red"));
400        assert_eq!(col.name_at(1), Some("Green"));
401    }
402
403    #[test]
404    fn test_enum16() {
405        let items = vec![
406            EnumItem { name: "Small".to_string(), value: 100 },
407            EnumItem { name: "Large".to_string(), value: 1000 },
408        ];
409        let mut col = ColumnEnum16::new(Type::enum16(items));
410
411        col.append_value(100);
412        col.append_name("Large").unwrap();
413
414        assert_eq!(col.len(), 2);
415        assert_eq!(col.at(0), 100);
416        assert_eq!(col.at(1), 1000);
417        assert_eq!(col.name_at(0), Some("Small"));
418        assert_eq!(col.name_at(1), Some("Large"));
419    }
420}