1use super::{Type, Value};
2use crate::types::{FromSqlError, FromSqlResult, OrderedMap};
3
4use crate::Row;
5use rust_decimal::prelude::*;
6
7use arrow::{
8 array::{
9 Array, ArrayRef, DictionaryArray, FixedSizeListArray, LargeListArray, ListArray, MapArray, StringArray,
10 StructArray, UnionArray,
11 },
12 datatypes::{UInt16Type, UInt32Type, UInt8Type},
13};
14
15#[derive(Copy, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
18pub enum TimeUnit {
19 Second,
21 Millisecond,
23 Microsecond,
25 Nanosecond,
27}
28
29impl TimeUnit {
30 pub fn to_micros(&self, value: i64) -> i64 {
32 match self {
33 Self::Second => value * 1_000_000,
34 Self::Millisecond => value * 1000,
35 Self::Microsecond => value,
36 Self::Nanosecond => value / 1000,
37 }
38 }
39}
40
41#[derive(Copy, Clone, Debug, PartialEq)]
46pub enum ValueRef<'a> {
47 Null,
49 Boolean(bool),
51 TinyInt(i8),
53 SmallInt(i16),
55 Int(i32),
57 BigInt(i64),
59 HugeInt(i128),
61 UTinyInt(u8),
63 USmallInt(u16),
65 UInt(u32),
67 UBigInt(u64),
69 Float(f32),
71 Double(f64),
73 Decimal(Decimal),
75 Timestamp(TimeUnit, i64),
77 Text(&'a [u8]),
79 Blob(&'a [u8]),
81 Date32(i32),
83 Time64(TimeUnit, i64),
85 Interval {
87 months: i32,
89 days: i32,
91 nanos: i64,
93 },
94 List(ListType<'a>, usize),
96 Enum(EnumType<'a>, usize),
98 Struct(&'a StructArray, usize),
100 Array(&'a FixedSizeListArray, usize),
102 Map(&'a MapArray, usize),
104 Union(&'a ArrayRef, usize),
106}
107
108#[derive(Debug, Copy, Clone, PartialEq)]
110pub enum ListType<'a> {
111 Regular(&'a ListArray),
113 Large(&'a LargeListArray),
115}
116
117#[derive(Debug, Copy, Clone, PartialEq)]
119pub enum EnumType<'a> {
120 UInt8(&'a DictionaryArray<UInt8Type>),
122 UInt16(&'a DictionaryArray<UInt16Type>),
124 UInt32(&'a DictionaryArray<UInt32Type>),
126}
127
128impl ValueRef<'_> {
129 #[inline]
131 pub fn data_type(&self) -> Type {
132 match *self {
133 ValueRef::Null => Type::Null,
134 ValueRef::Boolean(_) => Type::Boolean,
135 ValueRef::TinyInt(_) => Type::TinyInt,
136 ValueRef::SmallInt(_) => Type::SmallInt,
137 ValueRef::Int(_) => Type::Int,
138 ValueRef::BigInt(_) => Type::BigInt,
139 ValueRef::HugeInt(_) => Type::HugeInt,
140 ValueRef::UTinyInt(_) => Type::UTinyInt,
141 ValueRef::USmallInt(_) => Type::USmallInt,
142 ValueRef::UInt(_) => Type::UInt,
143 ValueRef::UBigInt(_) => Type::UBigInt,
144 ValueRef::Float(_) => Type::Float,
145 ValueRef::Double(_) => Type::Double,
146 ValueRef::Decimal(_) => Type::Decimal,
147 ValueRef::Timestamp(..) => Type::Timestamp,
148 ValueRef::Text(_) => Type::Text,
149 ValueRef::Blob(_) => Type::Blob,
150 ValueRef::Date32(_) => Type::Date32,
151 ValueRef::Time64(..) => Type::Time64,
152 ValueRef::Interval { .. } => Type::Interval,
153 ValueRef::Struct(arr, _) => arr.data_type().into(),
154 ValueRef::Map(arr, _) => arr.data_type().into(),
155 ValueRef::Array(arr, _) => arr.data_type().into(),
156 ValueRef::List(arr, _) => match arr {
157 ListType::Large(arr) => arr.data_type().into(),
158 ListType::Regular(arr) => arr.data_type().into(),
159 },
160 ValueRef::Enum(..) => Type::Enum,
161 ValueRef::Union(arr, _) => arr.data_type().into(),
162 }
163 }
164
165 pub fn to_owned(&self) -> Value {
167 (*self).into()
168 }
169}
170
171impl<'a> ValueRef<'a> {
172 #[inline]
175 pub fn as_str(&self) -> FromSqlResult<&'a str> {
176 match *self {
177 ValueRef::Text(t) => std::str::from_utf8(t).map_err(|e| FromSqlError::Other(Box::new(e))),
178 _ => Err(FromSqlError::InvalidType),
179 }
180 }
181
182 #[inline]
185 pub fn as_blob(&self) -> FromSqlResult<&'a [u8]> {
186 match *self {
187 ValueRef::Blob(b) => Ok(b),
188 ValueRef::Text(t) => Ok(t),
189 _ => Err(FromSqlError::InvalidType),
190 }
191 }
192}
193
194impl From<ValueRef<'_>> for Value {
195 #[inline]
196 fn from(borrowed: ValueRef<'_>) -> Self {
197 match borrowed {
198 ValueRef::Null => Self::Null,
199 ValueRef::Boolean(i) => Self::Boolean(i),
200 ValueRef::TinyInt(i) => Self::TinyInt(i),
201 ValueRef::SmallInt(i) => Self::SmallInt(i),
202 ValueRef::Int(i) => Self::Int(i),
203 ValueRef::BigInt(i) => Self::BigInt(i),
204 ValueRef::HugeInt(i) => Self::HugeInt(i),
205 ValueRef::UTinyInt(i) => Self::UTinyInt(i),
206 ValueRef::USmallInt(i) => Self::USmallInt(i),
207 ValueRef::UInt(i) => Self::UInt(i),
208 ValueRef::UBigInt(i) => Self::UBigInt(i),
209 ValueRef::Float(i) => Self::Float(i),
210 ValueRef::Double(i) => Self::Double(i),
211 ValueRef::Decimal(i) => Self::Decimal(i),
212 ValueRef::Timestamp(tu, t) => Self::Timestamp(tu, t),
213 ValueRef::Text(s) => {
214 let s = std::str::from_utf8(s).expect("invalid UTF-8");
215 Self::Text(s.to_string())
216 }
217 ValueRef::Blob(b) => Self::Blob(b.to_vec()),
218 ValueRef::Date32(d) => Self::Date32(d),
219 ValueRef::Time64(t, d) => Self::Time64(t, d),
220 ValueRef::Interval { months, days, nanos } => Self::Interval { months, days, nanos },
221 ValueRef::List(items, idx) => match items {
222 ListType::Regular(items) => {
223 let offsets = items.offsets();
224 from_list(
225 offsets[idx].try_into().unwrap(),
226 offsets[idx + 1].try_into().unwrap(),
227 idx,
228 items.values(),
229 )
230 }
231 ListType::Large(items) => {
232 let offsets = items.offsets();
233 from_list(
234 offsets[idx].try_into().unwrap(),
235 offsets[idx + 1].try_into().unwrap(),
236 idx,
237 items.values(),
238 )
239 }
240 },
241 ValueRef::Enum(items, idx) => {
242 let dict_values = match items {
243 EnumType::UInt8(res) => res.values(),
244 EnumType::UInt16(res) => res.values(),
245 EnumType::UInt32(res) => res.values(),
246 }
247 .as_any()
248 .downcast_ref::<StringArray>()
249 .expect("Enum value is not a string");
250 let dict_key = match items {
251 EnumType::UInt8(res) => res.key(idx),
252 EnumType::UInt16(res) => res.key(idx),
253 EnumType::UInt32(res) => res.key(idx),
254 }
255 .unwrap();
256 Self::Enum(dict_values.value(dict_key).to_string())
257 }
258 ValueRef::Struct(items, idx) => {
259 let capacity = items.columns().len();
260 let mut value = Vec::with_capacity(capacity);
261 value.extend(
262 items
263 .columns()
264 .iter()
265 .zip(items.fields().iter().map(|f| f.name().to_owned()))
266 .map(|(column, name)| -> (String, Self) {
267 (name, Row::value_ref_internal(idx, 0, column).to_owned())
268 }),
269 );
270 Self::Struct(OrderedMap::from(value))
271 }
272 ValueRef::Map(arr, idx) => {
273 let keys = arr.keys();
274 let values = arr.values();
275 let offsets = arr.offsets();
276 let range = offsets[idx]..offsets[idx + 1];
277 let capacity = range.len();
278 let mut map_vec = Vec::with_capacity(capacity);
279 map_vec.extend(range.map(|row| {
280 let row = row.try_into().unwrap();
281 let key = Row::value_ref_internal(row, idx, keys).to_owned();
282 let value = Row::value_ref_internal(row, idx, values).to_owned();
283 (key, value)
284 }));
285 Self::Map(OrderedMap::from(map_vec))
286 }
287 ValueRef::Array(items, idx) => {
288 let value_length = usize::try_from(items.value_length()).unwrap();
289 let range = (idx * value_length)..((idx + 1) * value_length);
290 let capacity = value_length;
291 let mut array_vec = Vec::with_capacity(capacity);
292 array_vec.extend(range.map(|row| Row::value_ref_internal(row, idx, items.values()).to_owned()));
293 Self::Array(array_vec)
294 }
295 ValueRef::Union(column, idx) => {
296 let column = column.as_any().downcast_ref::<UnionArray>().unwrap();
297 let type_id = column.type_id(idx);
298 let value_offset = column.value_offset(idx);
299
300 let tag = Row::value_ref_internal(idx, value_offset, column.child(type_id));
301 Self::Union(Box::new(tag.to_owned()))
302 }
303 }
304 }
305}
306
307fn from_list(start: usize, end: usize, idx: usize, values: &ArrayRef) -> Value {
308 let capacity = end - start;
309 let mut list_vec = Vec::with_capacity(capacity);
310 list_vec.extend((start..end).map(|row| Row::value_ref_internal(row, idx, values).to_owned()));
311 Value::List(list_vec)
312}
313
314impl<'a> From<&'a str> for ValueRef<'a> {
315 #[inline]
316 fn from(s: &str) -> ValueRef<'_> {
317 ValueRef::Text(s.as_bytes())
318 }
319}
320
321impl<'a> From<&'a [u8]> for ValueRef<'a> {
322 #[inline]
323 fn from(s: &[u8]) -> ValueRef<'_> {
324 ValueRef::Blob(s)
325 }
326}
327
328impl<'a> From<&'a Value> for ValueRef<'a> {
329 #[inline]
330 fn from(value: &'a Value) -> Self {
331 match *value {
332 Value::Null => ValueRef::Null,
333 Value::Boolean(i) => ValueRef::Boolean(i),
334 Value::TinyInt(i) => ValueRef::TinyInt(i),
335 Value::SmallInt(i) => ValueRef::SmallInt(i),
336 Value::Int(i) => ValueRef::Int(i),
337 Value::BigInt(i) => ValueRef::BigInt(i),
338 Value::HugeInt(i) => ValueRef::HugeInt(i),
339 Value::UTinyInt(i) => ValueRef::UTinyInt(i),
340 Value::USmallInt(i) => ValueRef::USmallInt(i),
341 Value::UInt(i) => ValueRef::UInt(i),
342 Value::UBigInt(i) => ValueRef::UBigInt(i),
343 Value::Float(i) => ValueRef::Float(i),
344 Value::Double(i) => ValueRef::Double(i),
345 Value::Decimal(i) => ValueRef::Decimal(i),
346 Value::Timestamp(tu, t) => ValueRef::Timestamp(tu, t),
347 Value::Text(ref s) => ValueRef::Text(s.as_bytes()),
348 Value::Blob(ref b) => ValueRef::Blob(b),
349 Value::Date32(d) => ValueRef::Date32(d),
350 Value::Time64(t, d) => ValueRef::Time64(t, d),
351 Value::Interval { months, days, nanos } => ValueRef::Interval { months, days, nanos },
352 Value::Enum(..) => todo!(),
353 Value::List(..) | Value::Struct(..) | Value::Map(..) | Value::Array(..) | Value::Union(..) => {
354 unimplemented!()
355 }
356 }
357 }
358}
359
360impl<'a, T> From<Option<T>> for ValueRef<'a>
361where
362 T: Into<Self>,
363{
364 #[inline]
365 fn from(s: Option<T>) -> Self {
366 match s {
367 Some(x) => x.into(),
368 None => ValueRef::Null,
369 }
370 }
371}
372
373#[cfg(test)]
374mod tests {
375 use crate::types::Type;
376 use crate::{Connection, Result};
377
378 #[test]
379 fn test_list_types() -> Result<()> {
380 let conn = Connection::open_in_memory()?;
381 conn.execute(
382 "CREATE TABLE test_table (float_list FLOAT[], double_list DOUBLE[], int_list INT[])",
383 [],
384 )?;
385 conn.execute("INSERT INTO test_table VALUES ([1.5, 2.5], [3.5, 4.5], [1, 2])", [])?;
386
387 let mut stmt = conn.prepare("SELECT float_list, double_list, int_list FROM test_table")?;
388 let mut rows = stmt.query([])?;
389 let row = rows.next()?.unwrap();
390
391 let float_list = row.get_ref_unwrap(0);
392 assert!(
393 matches!(float_list.data_type(), Type::List(ref inner_type) if **inner_type == Type::Float),
394 "Expected Type::List(Type::Float), got {:?}",
395 float_list.data_type()
396 );
397
398 let double_list = row.get_ref_unwrap(1);
399 assert!(
400 matches!(double_list.data_type(), Type::List(ref inner_type) if **inner_type == Type::Double),
401 "Expected Type::List(Type::Double), got {:?}",
402 double_list.data_type()
403 );
404
405 let int_list = row.get_ref_unwrap(2);
406 assert!(
407 matches!(int_list.data_type(), Type::List(ref inner_type) if **inner_type == Type::Int),
408 "Expected Type::List(Type::Int), got {:?}",
409 int_list.data_type()
410 );
411
412 Ok(())
413 }
414}