Skip to main content

rudb_arrow/
types.rs

1//! Arrow types, and the format strings the C data interface names them with.
2
3use rudb_common::{Error, LogicalType, Result};
4
5/// An Arrow type.
6///
7/// The subset our own types map onto, which is every type the engine can produce a value of today.
8/// The nested types are missing for the same reason `rudb-vector` has no nested vector: a list is
9/// offsets plus a child array, and there is no child array until the storage layer has one.
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub enum DataType {
12    /// No values, all null, no buffers at all.
13    Null,
14    /// One bit per value.
15    Boolean,
16    /// 8 bit signed.
17    Int8,
18    /// 16 bit signed.
19    Int16,
20    /// 32 bit signed.
21    Int32,
22    /// 64 bit signed.
23    Int64,
24    /// 8 bit unsigned.
25    UInt8,
26    /// 16 bit unsigned.
27    UInt16,
28    /// 32 bit unsigned.
29    UInt32,
30    /// 64 bit unsigned.
31    UInt64,
32    /// IEEE 754 binary32.
33    Float32,
34    /// IEEE 754 binary64.
35    Float64,
36    /// UTF-8, with 32 bit offsets.
37    Utf8,
38    /// Bytes, with 32 bit offsets.
39    Binary,
40    /// Days since 1970-01-01, 32 bit.
41    Date32,
42    /// Microseconds since midnight, 64 bit.
43    Time64,
44    /// Microseconds since the epoch, 64 bit, with a time zone when there is one.
45    Timestamp(TimeUnit, Option<String>),
46    /// Months, days and nanoseconds, sixteen bytes.
47    Interval,
48    /// A 128 bit integer with a decimal point in it.
49    Decimal128 {
50        /// How many digits it can hold.
51        precision: u8,
52        /// How many of them are after the point.
53        scale: u8,
54    },
55}
56
57/// How finely a timestamp counts.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum TimeUnit {
60    /// Seconds.
61    Second,
62    /// Milliseconds.
63    Millisecond,
64    /// Microseconds, which is what our own `TIMESTAMP` is.
65    Microsecond,
66    /// Nanoseconds.
67    Nanosecond,
68}
69
70impl TimeUnit {
71    /// The letter the format string uses.
72    fn letter(self) -> char {
73        match self {
74            Self::Second => 's',
75            Self::Millisecond => 'm',
76            Self::Microsecond => 'u',
77            Self::Nanosecond => 'n',
78        }
79    }
80}
81
82impl DataType {
83    /// The Arrow type one of ours becomes.
84    ///
85    /// `HUGEINT` becomes `DECIMAL128(38, 0)`, which is what DuckDB exports it as and the only thing
86    /// it can be: Arrow has no 128 bit integer and a decimal with no fractional digits is one.
87    ///
88    /// # Errors
89    ///
90    /// For a type with no Arrow counterpart yet, which is the nested types, `UHUGEINT`, `BIT` and
91    /// `UUID`.
92    pub fn of(ty: &LogicalType) -> Result<Self> {
93        Ok(match ty {
94            LogicalType::Null => Self::Null,
95            LogicalType::Boolean => Self::Boolean,
96            LogicalType::TinyInt => Self::Int8,
97            LogicalType::SmallInt => Self::Int16,
98            LogicalType::Integer => Self::Int32,
99            LogicalType::BigInt => Self::Int64,
100            LogicalType::HugeInt => Self::Decimal128 { precision: 38, scale: 0 },
101            LogicalType::UTinyInt => Self::UInt8,
102            LogicalType::USmallInt => Self::UInt16,
103            LogicalType::UInteger => Self::UInt32,
104            LogicalType::UBigInt => Self::UInt64,
105            LogicalType::Float => Self::Float32,
106            LogicalType::Double => Self::Float64,
107            LogicalType::Decimal { width, scale } => {
108                Self::Decimal128 { precision: *width, scale: *scale }
109            }
110            LogicalType::Varchar => Self::Utf8,
111            LogicalType::Blob => Self::Binary,
112            LogicalType::Date => Self::Date32,
113            LogicalType::Time => Self::Time64,
114            LogicalType::Timestamp => Self::Timestamp(TimeUnit::Microsecond, None),
115            LogicalType::TimestampS => Self::Timestamp(TimeUnit::Second, None),
116            LogicalType::TimestampMs => Self::Timestamp(TimeUnit::Millisecond, None),
117            LogicalType::TimestampNs => Self::Timestamp(TimeUnit::Nanosecond, None),
118            LogicalType::TimestampTz => {
119                Self::Timestamp(TimeUnit::Microsecond, Some("UTC".to_string()))
120            }
121            LogicalType::Interval => Self::Interval,
122            other => {
123                return Err(Error::not_implemented(format!("exporting {other} to Arrow")));
124            }
125        })
126    }
127
128    /// The format string the Arrow C data interface names this type with.
129    ///
130    /// Written now, before there is an FFI boundary to hand it across, because it is the part of
131    /// the mapping that is defined by somebody else's document and the part a test can check
132    /// against that document. The export itself is then a struct with this string in it.
133    #[must_use]
134    pub fn format(&self) -> String {
135        match self {
136            Self::Null => "n".to_string(),
137            Self::Boolean => "b".to_string(),
138            Self::Int8 => "c".to_string(),
139            Self::Int16 => "s".to_string(),
140            Self::Int32 => "i".to_string(),
141            Self::Int64 => "l".to_string(),
142            Self::UInt8 => "C".to_string(),
143            Self::UInt16 => "S".to_string(),
144            Self::UInt32 => "I".to_string(),
145            Self::UInt64 => "L".to_string(),
146            Self::Float32 => "f".to_string(),
147            Self::Float64 => "g".to_string(),
148            Self::Utf8 => "u".to_string(),
149            Self::Binary => "z".to_string(),
150            Self::Date32 => "tdD".to_string(),
151            Self::Time64 => "ttu".to_string(),
152            Self::Timestamp(unit, zone) => {
153                format!("ts{}:{}", unit.letter(), zone.clone().unwrap_or_default())
154            }
155            Self::Interval => "tin".to_string(),
156            Self::Decimal128 { precision, scale } => format!("d:{precision},{scale}"),
157        }
158    }
159
160    /// How many buffers an array of this type has, which the C data interface also has to say.
161    ///
162    /// Two for anything fixed width, which is the validity bitmap and the values. Three for a
163    /// variable width type, which puts the offsets in between. None at all for the null type, which
164    /// has no values to be valid or invalid.
165    #[must_use]
166    pub fn buffer_count(&self) -> usize {
167        match self {
168            Self::Null => 0,
169            Self::Utf8 | Self::Binary => 3,
170            _ => 2,
171        }
172    }
173
174    /// How wide one value is, for the fixed width types.
175    #[must_use]
176    pub fn width(&self) -> Option<usize> {
177        Some(match self {
178            Self::Null | Self::Utf8 | Self::Binary => return None,
179            // A boolean is a bit rather than a byte, and the caller that asks this is asking about
180            // bytes, so it is not a fixed width type for this purpose either.
181            Self::Boolean => return None,
182            Self::Int8 | Self::UInt8 => 1,
183            Self::Int16 | Self::UInt16 => 2,
184            Self::Int32 | Self::UInt32 | Self::Float32 | Self::Date32 => 4,
185            Self::Int64 | Self::UInt64 | Self::Float64 | Self::Time64 | Self::Timestamp(_, _) => 8,
186            Self::Interval | Self::Decimal128 { .. } => 16,
187        })
188    }
189}
190
191/// One column's name and type.
192#[derive(Debug, Clone, PartialEq, Eq)]
193pub struct Field {
194    /// What the column is called.
195    pub name: String,
196    /// What it holds.
197    pub data_type: DataType,
198    /// Whether it may hold nulls. Everything a query produces may, so this is true unless somebody
199    /// building a schema by hand says otherwise.
200    pub nullable: bool,
201}
202
203impl Field {
204    /// A nullable field, which is what a query result column is.
205    #[must_use]
206    pub fn new(name: impl Into<String>, data_type: DataType) -> Self {
207        Self { name: name.into(), data_type, nullable: true }
208    }
209}
210
211/// The columns of a record batch, in order.
212#[derive(Debug, Clone, Default, PartialEq, Eq)]
213pub struct Schema {
214    /// The fields, left to right.
215    pub fields: Vec<Field>,
216}
217
218impl Schema {
219    /// A schema of these fields.
220    #[must_use]
221    pub fn new(fields: Vec<Field>) -> Self {
222        Self { fields }
223    }
224
225    /// How many columns.
226    #[must_use]
227    pub fn len(&self) -> usize {
228        self.fields.len()
229    }
230
231    /// Whether there are no columns.
232    #[must_use]
233    pub fn is_empty(&self) -> bool {
234        self.fields.is_empty()
235    }
236}