clickhouse_arrow/native/convert/
raw_row.rs

1use std::borrow::Cow;
2
3use crate::{Error, FromSql, Result, Row, ToSql, Type, Value};
4
5/// A row of raw data returned from the database by a query.
6/// Or an unstructured runtime-defined row to upload to the server.
7#[derive(Debug, Default, Clone)]
8pub struct RawRow(Vec<Option<(String, Type, Value)>>);
9
10impl Row for RawRow {
11    const COLUMN_COUNT: Option<usize> = None;
12
13    fn column_names() -> Option<Vec<Cow<'static, str>>> { None }
14
15    fn to_schema() -> Option<Vec<(String, Type, Option<Value>)>> { None }
16
17    fn deserialize_row(map: Vec<(&str, &Type, Value)>) -> Result<Self> {
18        Ok(Self(
19            map.into_iter()
20                .map(|(name, type_, value)| Some((name.to_string(), type_.clone(), value)))
21                .collect(),
22        ))
23    }
24
25    fn serialize_row(
26        self,
27        _type_hints: &[(String, Type)],
28    ) -> Result<Vec<(Cow<'static, str>, Value)>> {
29        Ok(self
30            .0
31            .into_iter()
32            .map(|x| x.expect("cannot serialize a Row which has been retrieved from"))
33            .map(|(name, _, value)| (Cow::Owned(name), value))
34            .collect())
35    }
36}
37
38pub trait RowIndex {
39    fn get<'a, I: IntoIterator<Item = &'a str>>(&self, columns: I) -> Option<usize>;
40}
41
42impl RowIndex for usize {
43    fn get<'a, I: IntoIterator<Item = &'a str>>(&self, columns: I) -> Option<usize> {
44        let count = columns.into_iter().count();
45        if count >= *self { Some(*self) } else { None }
46    }
47}
48
49impl RowIndex for str {
50    fn get<'a, I: IntoIterator<Item = &'a str>>(&self, columns: I) -> Option<usize> {
51        columns.into_iter().position(|x| x == self)
52    }
53}
54
55impl<T: RowIndex + ?Sized> RowIndex for &T {
56    fn get<'a, I: IntoIterator<Item = &'a str>>(&self, columns: I) -> Option<usize> {
57        (*self).get(columns)
58    }
59}
60
61impl RawRow {
62    /// Determines if the row contains no values.
63    pub fn is_empty(&self) -> bool { self.0.is_empty() }
64
65    /// Returns the number of values in the row.
66    pub fn len(&self) -> usize { self.0.len() }
67
68    /// # Panics
69    ///
70    /// Panics if any of the values are `None`
71    pub fn into_values(self) -> Vec<(Type, Value)> {
72        self.0.into_iter().map(|x| x.map(|(_, t, v)| (t, v)).unwrap()).collect()
73    }
74
75    /// Like [`RawRow::get`], but returns a [`Result`] rather than panicking.
76    ///
77    /// # Errors
78    ///
79    /// Returns an error if the index is out of bounds or if the value cannot be converted to the
80    /// specified
81    ///
82    /// # Panics
83    ///
84    /// Shouldn't panic, bounds checked
85    pub fn try_get<I: RowIndex, T: FromSql>(&mut self, index: I) -> Result<T> {
86        let index = index
87            .get(self.0.iter().map(|x| x.as_ref().map_or("", |x| &*x.0)))
88            .ok_or(Error::OutOfBounds)?;
89        let (_, type_, value) = self.0.get_mut(index).unwrap().take().ok_or(Error::DoubleFetch)?;
90        T::from_sql(&type_, value)
91    }
92
93    /// Deserializes a value from the row.
94    /// The value can be specified either by its numeric index in the row, or by its column name.
95    /// # Panics
96    /// Panics if the index is out of bounds or if the value cannot be converted to the specified
97    /// type.
98    pub fn get<I: RowIndex, T: FromSql>(&mut self, index: I) -> T {
99        self.try_get(index).expect("failed to convert column")
100    }
101
102    /// Sets or inserts a column value with a given name. `type_` is inferred if `None`. Index is
103    /// defined on insertion order.
104    ///
105    /// # Errors
106    ///
107    /// Returns an error if type conversion fails.
108    ///
109    /// # Panics
110    ///
111    /// Shouldn't panic as `current_position` checks for existence of element
112    pub fn try_set_typed(
113        &mut self,
114        name: &impl ToString,
115        type_: Option<Type>,
116        value: impl ToSql,
117    ) -> Result<()> {
118        let name = name.to_string();
119        let value = value.to_sql(type_.as_ref())?;
120        let type_ = type_.unwrap_or_else(|| value.guess_type());
121
122        let current_position =
123            self.0.iter().map(|x| x.as_ref().map_or("", |x| &*x.0)).position(|x| x == &*name);
124
125        if let Some(current_position) = current_position {
126            self.0[current_position].as_mut().unwrap().1 = type_;
127            self.0[current_position].as_mut().unwrap().2 = value;
128        } else {
129            self.0.push(Some((name, type_, value)));
130        }
131        Ok(())
132    }
133
134    /// Same as `try_set_typed`, but always infers the type
135    ///
136    /// # Errors
137    ///
138    /// Returns an error if type conversion fails.
139    pub fn try_set(&mut self, name: &impl ToString, value: impl ToSql) -> Result<()> {
140        self.try_set_typed(name, None, value)
141    }
142
143    /// Same as `try_set`, but panics on type conversion failure.
144    ///
145    /// # Panics
146    ///
147    /// Panics on type conversion failure
148    pub fn set(&mut self, name: &impl ToString, value: impl ToSql) {
149        self.try_set(name, value).expect("failed to convert column");
150    }
151
152    /// Same as `try_set_typed`, but panics on type conversion failure.
153    /// # Panics
154    ///
155    /// Panics on type conversion failure
156    pub fn set_typed(&mut self, name: &impl ToString, type_: Option<Type>, value: impl ToSql) {
157        self.try_set_typed(name, type_, value).expect("failed to convert column");
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    #[test]
166    fn test_raw_row_default() {
167        let row = RawRow::default();
168        assert!(row.is_empty());
169        assert_eq!(row.len(), 0);
170    }
171
172    #[test]
173    fn test_raw_row_basic_operations() {
174        let mut row = RawRow::default();
175
176        // Test empty row
177        assert!(row.is_empty());
178        assert_eq!(row.len(), 0);
179
180        // Test setting values
181        row.set(&"test_col", 42i32);
182        assert!(!row.is_empty());
183        assert_eq!(row.len(), 1);
184
185        // Test try_set
186        row.try_set(&"test_col2", "hello").unwrap();
187        assert_eq!(row.len(), 2);
188    }
189
190    #[test]
191    fn test_raw_row_set_typed() {
192        let mut row = RawRow::default();
193
194        // Test with explicit type
195        row.set_typed(&"int_col", Some(Type::Int32), 123i32);
196        assert_eq!(row.len(), 1);
197
198        // Test try_set_typed
199        row.try_set_typed(&"str_col", Some(Type::String), "test").unwrap();
200        assert_eq!(row.len(), 2);
201
202        // Test updating existing column
203        row.try_set_typed(&"int_col", Some(Type::Int32), 456i32).unwrap();
204        assert_eq!(row.len(), 2); // Should still be 2, not 3
205    }
206
207    #[test]
208    fn test_raw_row_get_by_index() {
209        let mut row = RawRow::default();
210        row.set(&"col1", 42i32);
211        row.set(&"col2", "hello");
212
213        // Test getting by numeric index
214        let val: i32 = row.try_get(0usize).unwrap();
215        assert_eq!(val, 42);
216    }
217
218    #[test]
219    fn test_raw_row_get_by_name() {
220        let mut row = RawRow::default();
221        row.set(&"test_column", 123i64);
222
223        // Test getting by column name
224        let val: i64 = row.try_get("test_column").unwrap();
225        assert_eq!(val, 123);
226    }
227
228    #[test]
229    fn test_raw_row_get_panics() {
230        let mut row = RawRow::default();
231        row.set(&"test", 42i32);
232
233        // This should work
234        let _val: i32 = row.get(0usize);
235
236        // Second get should panic due to double fetch
237        drop(
238            std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
239                let _val2: i32 = row.get(0usize);
240            }))
241            .unwrap_err(),
242        );
243    }
244
245    #[test]
246    fn test_raw_row_errors() {
247        let mut row = RawRow::default();
248        row.set(&"test", 42i32);
249
250        // Test out of bounds error
251        let result: Result<i32> = row.try_get(10usize);
252        assert!(matches!(result, Err(Error::OutOfBounds)));
253
254        // Test double fetch error
255        let _val: i32 = row.try_get(0usize).unwrap();
256        let result2: Result<i32> = row.try_get(0usize);
257        assert!(matches!(result2, Err(Error::DoubleFetch)));
258    }
259
260    #[test]
261    fn test_raw_row_into_values() {
262        let mut row = RawRow::default();
263        row.set(&"col1", 42i32);
264        row.set(&"col2", "test");
265
266        let values = row.into_values();
267        assert_eq!(values.len(), 2);
268        assert_eq!(values[0].0, Type::Int32);
269        assert_eq!(values[1].0, Type::String);
270    }
271
272    #[test]
273    fn test_row_index_usize() {
274        let columns = ["col1", "col2", "col3"];
275
276        // Valid index
277        assert_eq!(1usize.get(columns.iter().copied()), Some(1));
278
279        // Invalid index
280        assert_eq!(5usize.get(columns.iter().copied()), None);
281    }
282
283    #[test]
284    fn test_row_index_str() {
285        let columns = ["col1", "col2", "col3"];
286
287        // Valid column name
288        assert_eq!(RowIndex::get("col2", columns.iter().copied()), Some(1));
289
290        // Invalid column name
291        assert_eq!(RowIndex::get("nonexistent", columns.iter().copied()), None);
292    }
293
294    #[test]
295    fn test_row_index_ref() {
296        let columns = ["col1", "col2", "col3"];
297        let name = "col2";
298
299        // Test reference implementation
300        assert_eq!((&name).get(columns.iter().copied()), Some(1));
301        assert_eq!((1usize).get(columns.iter().copied()), Some(1));
302    }
303
304    #[test]
305    fn test_raw_row_deserialize() {
306        let map = vec![
307            ("col1", &Type::Int32, Value::Int32(42)),
308            ("col2", &Type::String, Value::String("test".to_string().into_bytes())),
309        ];
310
311        let row = RawRow::deserialize_row(map).unwrap();
312        assert_eq!(row.len(), 2);
313        assert!(!row.is_empty());
314    }
315
316    #[test]
317    fn test_raw_row_serialize() {
318        let mut row = RawRow::default();
319        row.set(&"col1", 42i32);
320        row.set(&"col2", "test");
321
322        let serialized = row.serialize_row(&[]).unwrap();
323        assert_eq!(serialized.len(), 2);
324        assert_eq!(serialized[0].0, "col1");
325        assert_eq!(serialized[1].0, "col2");
326    }
327
328    #[test]
329    fn test_raw_row_serialize_panic() {
330        // Create a row with a None value (which should cause serialize to panic)
331        let raw_row = RawRow(vec![None]);
332
333        // This should panic because we have a None value
334        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
335            raw_row.serialize_row(&[]).unwrap()
336        }));
337        assert!(result.is_err());
338    }
339
340    #[test]
341    fn test_raw_row_static_methods() {
342        // Test Row trait static methods
343        assert_eq!(RawRow::COLUMN_COUNT, None);
344        assert_eq!(RawRow::column_names(), None);
345        assert_eq!(RawRow::to_schema(), None);
346    }
347}
clickhouse_arrow/native/convert/raw_row.rs

clickhouse_arrow/native/convert/
raw_row.rs
