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