duckdb/types/
value_ref.rs

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/// An absolute length of time in seconds, milliseconds, microseconds or nanoseconds.
16/// Copy from arrow::datatypes::TimeUnit
17#[derive(Copy, Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
18pub enum TimeUnit {
19    /// Time in seconds.
20    Second,
21    /// Time in milliseconds.
22    Millisecond,
23    /// Time in microseconds.
24    Microsecond,
25    /// Time in nanoseconds.
26    Nanosecond,
27}
28
29impl TimeUnit {
30    /// Convert a number of `TimeUnit` to microseconds.
31    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/// A non-owning [static type value](https://duckdb.org/docs/sql/data_types/overview). Typically the
42/// memory backing this value is owned by SQLite.
43///
44/// See [`Value`](Value) for an owning dynamic type value.
45#[derive(Copy, Clone, Debug, PartialEq)]
46pub enum ValueRef<'a> {
47    /// The value is a `NULL` value.
48    Null,
49    /// The value is a boolean.
50    Boolean(bool),
51    /// The value is a signed tiny integer.
52    TinyInt(i8),
53    /// The value is a signed small integer.
54    SmallInt(i16),
55    /// The value is a signed integer.
56    Int(i32),
57    /// The value is a signed big integer.
58    BigInt(i64),
59    /// The value is a signed huge integer.
60    HugeInt(i128),
61    /// The value is a unsigned tiny integer.
62    UTinyInt(u8),
63    /// The value is a unsigned small integer.
64    USmallInt(u16),
65    /// The value is a unsigned integer.
66    UInt(u32),
67    /// The value is a unsigned big integer.
68    UBigInt(u64),
69    /// The value is a f32.
70    Float(f32),
71    /// The value is a f64.
72    Double(f64),
73    /// The value is a decimal
74    Decimal(Decimal),
75    /// The value is a timestamp.
76    Timestamp(TimeUnit, i64),
77    /// The value is a text string.
78    Text(&'a [u8]),
79    /// The value is a blob of data
80    Blob(&'a [u8]),
81    /// The value is a date32
82    Date32(i32),
83    /// The value is a time64
84    Time64(TimeUnit, i64),
85    /// The value is an interval (month, day, nano)
86    Interval {
87        /// months
88        months: i32,
89        /// days
90        days: i32,
91        /// nanos
92        nanos: i64,
93    },
94    /// The value is a list
95    List(ListType<'a>, usize),
96    /// The value is an enum
97    Enum(EnumType<'a>, usize),
98    /// The value is a struct
99    Struct(&'a StructArray, usize),
100    /// The value is an array
101    Array(&'a FixedSizeListArray, usize),
102    /// The value is a map
103    Map(&'a MapArray, usize),
104    /// The value is a union
105    Union(&'a ArrayRef, usize),
106}
107
108/// Wrapper type for different list sizes
109#[derive(Debug, Copy, Clone, PartialEq)]
110pub enum ListType<'a> {
111    /// The underlying list is a `ListArray`
112    Regular(&'a ListArray),
113    /// The underlying list is a `LargeListArray`
114    Large(&'a LargeListArray),
115}
116
117/// Wrapper type for different enum sizes
118#[derive(Debug, Copy, Clone, PartialEq)]
119pub enum EnumType<'a> {
120    /// The underlying enum type is u8
121    UInt8(&'a DictionaryArray<UInt8Type>),
122    /// The underlying enum type is u16
123    UInt16(&'a DictionaryArray<UInt16Type>),
124    /// The underlying enum type is u32
125    UInt32(&'a DictionaryArray<UInt32Type>),
126}
127
128impl ValueRef<'_> {
129    /// Returns DuckDB fundamental datatype.
130    #[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    /// Returns an owned version of this ValueRef
166    pub fn to_owned(&self) -> Value {
167        (*self).into()
168    }
169}
170
171impl<'a> ValueRef<'a> {
172    /// If `self` is case `Text`, returns the string value. Otherwise, returns
173    /// [`Err(Error::InvalidColumnType)`](crate::Error::InvalidColumnType).
174    #[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    /// If `self` is case `Blob`, returns the byte slice. Otherwise, returns
183    /// [`Err(Error::InvalidColumnType)`](crate::Error::InvalidColumnType).
184    #[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}