arrow_integration_test/
datatype.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow::datatypes::{DataType, Field, Fields, IntervalUnit, TimeUnit, UnionMode};
19use arrow::error::{ArrowError, Result};
20use std::sync::Arc;
21
22/// Read a JSON number as an integer of type `T`.
23fn json_int<T: TryFrom<i64>>(what: &str, value: &serde_json::Value) -> Result<T> {
24    let int = value
25        .as_i64()
26        .ok_or_else(|| ArrowError::ParseError(format!("Expecting {what} to be an integer")))?;
27    T::try_from(int)
28        .map_err(|_| ArrowError::ParseError(format!("{what} is out of range for its type: {int}")))
29}
30
31/// Parse a data type from a JSON representation.
32pub fn data_type_from_json(json: &serde_json::Value) -> Result<DataType> {
33    use serde_json::Value;
34    let default_field = Arc::new(Field::new("", DataType::Boolean, true));
35    match *json {
36        Value::Object(ref map) => match map.get("name") {
37            Some(s) if s == "null" => Ok(DataType::Null),
38            Some(s) if s == "bool" => Ok(DataType::Boolean),
39            Some(s) if s == "binary" => Ok(DataType::Binary),
40            Some(s) if s == "largebinary" => Ok(DataType::LargeBinary),
41            Some(s) if s == "binaryview" => Ok(DataType::BinaryView),
42            Some(s) if s == "utf8view" => Ok(DataType::Utf8View),
43            Some(s) if s == "utf8" => Ok(DataType::Utf8),
44            Some(s) if s == "largeutf8" => Ok(DataType::LargeUtf8),
45            Some(s) if s == "fixedsizebinary" => {
46                // return a list with any type as its child isn't defined in the map
47                if let Some(Value::Number(size)) = map.get("byteWidth") {
48                    Ok(DataType::FixedSizeBinary(json_int(
49                        "byteWidth",
50                        &Value::Number(size.clone()),
51                    )?))
52                } else {
53                    Err(ArrowError::ParseError(
54                        "Expecting a byteWidth for fixedsizebinary".to_string(),
55                    ))
56                }
57            }
58            Some(s) if s == "decimal" => {
59                // return a list with any type as its child isn't defined in the map
60                let precision = match map.get("precision") {
61                    Some(p) => json_int("precision", p),
62                    None => Err(ArrowError::ParseError(
63                        "Expecting a precision for decimal".to_string(),
64                    )),
65                }?;
66                let scale = match map.get("scale") {
67                    Some(s) => json_int("scale", s),
68                    _ => Err(ArrowError::ParseError(
69                        "Expecting a scale for decimal".to_string(),
70                    )),
71                }?;
72                let bit_width: usize = match map.get("bitWidth") {
73                    Some(b) => json_int("bitWidth", b)?,
74                    _ => 128, // Default bit width
75                };
76
77                match bit_width {
78                    32 => Ok(DataType::Decimal32(precision, scale)),
79                    64 => Ok(DataType::Decimal64(precision, scale)),
80                    128 => Ok(DataType::Decimal128(precision, scale)),
81                    256 => Ok(DataType::Decimal256(precision, scale)),
82                    _ => Err(ArrowError::ParseError(
83                        "Decimal bit_width invalid".to_string(),
84                    )),
85                }
86            }
87            Some(s) if s == "floatingpoint" => match map.get("precision") {
88                Some(p) if p == "HALF" => Ok(DataType::Float16),
89                Some(p) if p == "SINGLE" => Ok(DataType::Float32),
90                Some(p) if p == "DOUBLE" => Ok(DataType::Float64),
91                _ => Err(ArrowError::ParseError(
92                    "floatingpoint precision missing or invalid".to_string(),
93                )),
94            },
95            Some(s) if s == "timestamp" => {
96                let unit = match map.get("unit") {
97                    Some(p) if p == "SECOND" => Ok(TimeUnit::Second),
98                    Some(p) if p == "MILLISECOND" => Ok(TimeUnit::Millisecond),
99                    Some(p) if p == "MICROSECOND" => Ok(TimeUnit::Microsecond),
100                    Some(p) if p == "NANOSECOND" => Ok(TimeUnit::Nanosecond),
101                    _ => Err(ArrowError::ParseError(
102                        "timestamp unit missing or invalid".to_string(),
103                    )),
104                };
105                let tz = match map.get("timezone") {
106                    None => Ok(None),
107                    Some(Value::String(tz)) => Ok(Some(tz.as_str().into())),
108                    _ => Err(ArrowError::ParseError(
109                        "timezone must be a string".to_string(),
110                    )),
111                };
112                Ok(DataType::Timestamp(unit?, tz?))
113            }
114            Some(s) if s == "date" => match map.get("unit") {
115                Some(p) if p == "DAY" => Ok(DataType::Date32),
116                Some(p) if p == "MILLISECOND" => Ok(DataType::Date64),
117                _ => Err(ArrowError::ParseError(
118                    "date unit missing or invalid".to_string(),
119                )),
120            },
121            Some(s) if s == "time" => {
122                let unit = match map.get("unit") {
123                    Some(p) if p == "SECOND" => Ok(TimeUnit::Second),
124                    Some(p) if p == "MILLISECOND" => Ok(TimeUnit::Millisecond),
125                    Some(p) if p == "MICROSECOND" => Ok(TimeUnit::Microsecond),
126                    Some(p) if p == "NANOSECOND" => Ok(TimeUnit::Nanosecond),
127                    _ => Err(ArrowError::ParseError(
128                        "time unit missing or invalid".to_string(),
129                    )),
130                };
131                match map.get("bitWidth") {
132                    Some(p) if p == 32 => Ok(DataType::Time32(unit?)),
133                    Some(p) if p == 64 => Ok(DataType::Time64(unit?)),
134                    _ => Err(ArrowError::ParseError(
135                        "time bitWidth missing or invalid".to_string(),
136                    )),
137                }
138            }
139            Some(s) if s == "duration" => match map.get("unit") {
140                Some(p) if p == "SECOND" => Ok(DataType::Duration(TimeUnit::Second)),
141                Some(p) if p == "MILLISECOND" => Ok(DataType::Duration(TimeUnit::Millisecond)),
142                Some(p) if p == "MICROSECOND" => Ok(DataType::Duration(TimeUnit::Microsecond)),
143                Some(p) if p == "NANOSECOND" => Ok(DataType::Duration(TimeUnit::Nanosecond)),
144                _ => Err(ArrowError::ParseError(
145                    "time unit missing or invalid".to_string(),
146                )),
147            },
148            Some(s) if s == "interval" => match map.get("unit") {
149                Some(p) if p == "DAY_TIME" => Ok(DataType::Interval(IntervalUnit::DayTime)),
150                Some(p) if p == "YEAR_MONTH" => Ok(DataType::Interval(IntervalUnit::YearMonth)),
151                Some(p) if p == "MONTH_DAY_NANO" => {
152                    Ok(DataType::Interval(IntervalUnit::MonthDayNano))
153                }
154                _ => Err(ArrowError::ParseError(
155                    "interval unit missing or invalid".to_string(),
156                )),
157            },
158            Some(s) if s == "int" => match map.get("isSigned") {
159                Some(&Value::Bool(true)) => match map.get("bitWidth") {
160                    Some(Value::Number(n)) => match n.as_u64() {
161                        Some(8) => Ok(DataType::Int8),
162                        Some(16) => Ok(DataType::Int16),
163                        Some(32) => Ok(DataType::Int32),
164                        Some(64) => Ok(DataType::Int64),
165                        _ => Err(ArrowError::ParseError(
166                            "int bitWidth missing or invalid".to_string(),
167                        )),
168                    },
169                    _ => Err(ArrowError::ParseError(
170                        "int bitWidth missing or invalid".to_string(),
171                    )),
172                },
173                Some(&Value::Bool(false)) => match map.get("bitWidth") {
174                    Some(Value::Number(n)) => match n.as_u64() {
175                        Some(8) => Ok(DataType::UInt8),
176                        Some(16) => Ok(DataType::UInt16),
177                        Some(32) => Ok(DataType::UInt32),
178                        Some(64) => Ok(DataType::UInt64),
179                        _ => Err(ArrowError::ParseError(
180                            "int bitWidth missing or invalid".to_string(),
181                        )),
182                    },
183                    _ => Err(ArrowError::ParseError(
184                        "int bitWidth missing or invalid".to_string(),
185                    )),
186                },
187                _ => Err(ArrowError::ParseError(
188                    "int signed missing or invalid".to_string(),
189                )),
190            },
191            Some(s) if s == "list" => {
192                // return a list with any type as its child isn't defined in the map
193                Ok(DataType::List(default_field))
194            }
195            Some(s) if s == "largelist" => {
196                // return a largelist with any type as its child isn't defined in the map
197                Ok(DataType::LargeList(default_field))
198            }
199            Some(s) if s == "listview" => {
200                // return a listview with any type as its child isn't defined in the map
201                Ok(DataType::ListView(default_field))
202            }
203            Some(s) if s == "largelistview" => {
204                // return a large listview with any type as its child isn't defined in the map
205                Ok(DataType::LargeListView(default_field))
206            }
207            Some(s) if s == "fixedsizelist" => {
208                // return a list with any type as its child isn't defined in the map
209                if let Some(Value::Number(size)) = map.get("listSize") {
210                    Ok(DataType::FixedSizeList(
211                        default_field,
212                        json_int("listSize", &Value::Number(size.clone()))?,
213                    ))
214                } else {
215                    Err(ArrowError::ParseError(
216                        "Expecting a listSize for fixedsizelist".to_string(),
217                    ))
218                }
219            }
220            Some(s) if s == "struct" => {
221                // return an empty `struct` type as its children aren't defined in the map
222                Ok(DataType::Struct(Fields::empty()))
223            }
224            Some(s) if s == "runendencoded" => {
225                // return a run end encoded with placeholder types as children aren't defined in the map
226                Ok(DataType::RunEndEncoded(
227                    Arc::new(Field::new("run_ends", DataType::Int32, false)),
228                    default_field,
229                ))
230            }
231            Some(s) if s == "map" => {
232                if let Some(Value::Bool(keys_sorted)) = map.get("keysSorted") {
233                    // Return a map with an empty type as its children aren't defined in the map
234                    Ok(DataType::Map(default_field, *keys_sorted))
235                } else {
236                    Err(ArrowError::ParseError(
237                        "Expecting a keysSorted for map".to_string(),
238                    ))
239                }
240            }
241            Some(s) if s == "union" => {
242                if let Some(Value::String(mode)) = map.get("mode") {
243                    let union_mode = if mode == "SPARSE" {
244                        UnionMode::Sparse
245                    } else if mode == "DENSE" {
246                        UnionMode::Dense
247                    } else {
248                        return Err(ArrowError::ParseError(format!(
249                            "Unknown union mode {mode:?} for union"
250                        )));
251                    };
252                    if let Some(values) = map.get("typeIds") {
253                        let values = values.as_array().ok_or_else(|| {
254                            ArrowError::ParseError("Expecting typeIds to be an array".to_string())
255                        })?;
256                        let fields = values
257                            .iter()
258                            .map(|t| Ok((json_int::<i8>("a type id", t)?, default_field.clone())))
259                            .collect::<Result<Vec<_>>>()?
260                            .into_iter()
261                            .collect();
262
263                        Ok(DataType::Union(fields, union_mode))
264                    } else {
265                        Err(ArrowError::ParseError(
266                            "Expecting a typeIds for union ".to_string(),
267                        ))
268                    }
269                } else {
270                    Err(ArrowError::ParseError(
271                        "Expecting a mode for union".to_string(),
272                    ))
273                }
274            }
275            Some(other) => Err(ArrowError::ParseError(format!(
276                "invalid or unsupported type name: {other} in {json:?}"
277            ))),
278            None => Err(ArrowError::ParseError("type name missing".to_string())),
279        },
280        _ => Err(ArrowError::ParseError(
281            "invalid json value type".to_string(),
282        )),
283    }
284}
285
286/// Generate a JSON representation of the data type.
287pub fn data_type_to_json(data_type: &DataType) -> serde_json::Value {
288    use serde_json::json;
289    match data_type {
290        DataType::Null => json!({"name": "null"}),
291        DataType::Boolean => json!({"name": "bool"}),
292        DataType::Int8 => json!({"name": "int", "bitWidth": 8, "isSigned": true}),
293        DataType::Int16 => json!({"name": "int", "bitWidth": 16, "isSigned": true}),
294        DataType::Int32 => json!({"name": "int", "bitWidth": 32, "isSigned": true}),
295        DataType::Int64 => json!({"name": "int", "bitWidth": 64, "isSigned": true}),
296        DataType::UInt8 => json!({"name": "int", "bitWidth": 8, "isSigned": false}),
297        DataType::UInt16 => json!({"name": "int", "bitWidth": 16, "isSigned": false}),
298        DataType::UInt32 => json!({"name": "int", "bitWidth": 32, "isSigned": false}),
299        DataType::UInt64 => json!({"name": "int", "bitWidth": 64, "isSigned": false}),
300        DataType::Float16 => json!({"name": "floatingpoint", "precision": "HALF"}),
301        DataType::Float32 => json!({"name": "floatingpoint", "precision": "SINGLE"}),
302        DataType::Float64 => json!({"name": "floatingpoint", "precision": "DOUBLE"}),
303        DataType::Utf8 => json!({"name": "utf8"}),
304        DataType::LargeUtf8 => json!({"name": "largeutf8"}),
305        DataType::Binary => json!({"name": "binary"}),
306        DataType::LargeBinary => json!({"name": "largebinary"}),
307        DataType::BinaryView => json!({"name": "binaryview"}),
308        DataType::Utf8View => json!({"name": "utf8view"}),
309        DataType::FixedSizeBinary(byte_width) => {
310            json!({"name": "fixedsizebinary", "byteWidth": byte_width})
311        }
312        DataType::Struct(_) => json!({"name": "struct"}),
313        DataType::Union(_, _) => json!({"name": "union"}),
314        DataType::List(_) => json!({ "name": "list"}),
315        DataType::LargeList(_) => json!({ "name": "largelist"}),
316        DataType::ListView(_) => json!({ "name": "listview"}),
317        DataType::LargeListView(_) => json!({ "name": "largelistview"}),
318        DataType::FixedSizeList(_, length) => {
319            json!({"name":"fixedsizelist", "listSize": length})
320        }
321        DataType::Time32(unit) => {
322            json!({"name": "time", "bitWidth": 32, "unit": match unit {
323                TimeUnit::Second => "SECOND",
324                TimeUnit::Millisecond => "MILLISECOND",
325                TimeUnit::Microsecond => "MICROSECOND",
326                TimeUnit::Nanosecond => "NANOSECOND",
327            }})
328        }
329        DataType::Time64(unit) => {
330            json!({"name": "time", "bitWidth": 64, "unit": match unit {
331                TimeUnit::Second => "SECOND",
332                TimeUnit::Millisecond => "MILLISECOND",
333                TimeUnit::Microsecond => "MICROSECOND",
334                TimeUnit::Nanosecond => "NANOSECOND",
335            }})
336        }
337        DataType::Date32 => {
338            json!({"name": "date", "unit": "DAY"})
339        }
340        DataType::Date64 => {
341            json!({"name": "date", "unit": "MILLISECOND"})
342        }
343        DataType::Timestamp(unit, None) => {
344            json!({"name": "timestamp", "unit": match unit {
345                TimeUnit::Second => "SECOND",
346                TimeUnit::Millisecond => "MILLISECOND",
347                TimeUnit::Microsecond => "MICROSECOND",
348                TimeUnit::Nanosecond => "NANOSECOND",
349            }})
350        }
351        DataType::Timestamp(unit, Some(tz)) => {
352            json!({"name": "timestamp", "unit": match unit {
353                    TimeUnit::Second => "SECOND",
354                    TimeUnit::Millisecond => "MILLISECOND",
355                    TimeUnit::Microsecond => "MICROSECOND",
356                    TimeUnit::Nanosecond => "NANOSECOND",
357                }, "timezone": tz})
358        }
359        DataType::Interval(unit) => json!({"name": "interval", "unit": match unit {
360            IntervalUnit::YearMonth => "YEAR_MONTH",
361            IntervalUnit::DayTime => "DAY_TIME",
362            IntervalUnit::MonthDayNano => "MONTH_DAY_NANO",
363        }}),
364        DataType::Duration(unit) => json!({"name": "duration", "unit": match unit {
365            TimeUnit::Second => "SECOND",
366            TimeUnit::Millisecond => "MILLISECOND",
367            TimeUnit::Microsecond => "MICROSECOND",
368            TimeUnit::Nanosecond => "NANOSECOND",
369        }}),
370        DataType::Dictionary(_, _) => json!({ "name": "dictionary"}),
371        DataType::Decimal32(precision, scale) => {
372            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 32})
373        }
374        DataType::Decimal64(precision, scale) => {
375            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 64})
376        }
377        DataType::Decimal128(precision, scale) => {
378            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 128})
379        }
380        DataType::Decimal256(precision, scale) => {
381            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 256})
382        }
383        DataType::Map(_, keys_sorted) => {
384            json!({"name": "map", "keysSorted": keys_sorted})
385        }
386        DataType::RunEndEncoded(_, _) => json!({"name": "runendencoded"}),
387    }
388}
389
390#[cfg(test)]
391mod tests {
392    use super::*;
393    use serde_json::Value;
394
395    #[test]
396    fn parse_utf8_from_json() {
397        let json = "{\"name\":\"utf8\"}";
398        let value: Value = serde_json::from_str(json).unwrap();
399        let dt = data_type_from_json(&value).unwrap();
400        assert_eq!(DataType::Utf8, dt);
401    }
402
403    #[test]
404    fn parse_int32_from_json() {
405        let json = "{\"name\": \"int\", \"isSigned\": true, \"bitWidth\": 32}";
406        let value: Value = serde_json::from_str(json).unwrap();
407        let dt = data_type_from_json(&value).unwrap();
408        assert_eq!(DataType::Int32, dt);
409    }
410}
arrow_integration_test/datatype.rs

arrow_integration_test/
datatype.rs
