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(
228                        Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
229                        DataType::Int32,
230                        false,
231                    )),
232                    default_field,
233                ))
234            }
235            Some(s) if s == "map" => {
236                if let Some(Value::Bool(keys_sorted)) = map.get("keysSorted") {
237                    // Return a map with an empty type as its children aren't defined in the map
238                    Ok(DataType::Map(default_field, *keys_sorted))
239                } else {
240                    Err(ArrowError::ParseError(
241                        "Expecting a keysSorted for map".to_string(),
242                    ))
243                }
244            }
245            Some(s) if s == "union" => {
246                if let Some(Value::String(mode)) = map.get("mode") {
247                    let union_mode = if mode == "SPARSE" {
248                        UnionMode::Sparse
249                    } else if mode == "DENSE" {
250                        UnionMode::Dense
251                    } else {
252                        return Err(ArrowError::ParseError(format!(
253                            "Unknown union mode {mode:?} for union"
254                        )));
255                    };
256                    if let Some(values) = map.get("typeIds") {
257                        let values = values.as_array().ok_or_else(|| {
258                            ArrowError::ParseError("Expecting typeIds to be an array".to_string())
259                        })?;
260                        let fields = values
261                            .iter()
262                            .map(|t| Ok((json_int::<i8>("a type id", t)?, default_field.clone())))
263                            .collect::<Result<Vec<_>>>()?
264                            .into_iter()
265                            .collect();
266
267                        Ok(DataType::Union(fields, union_mode))
268                    } else {
269                        Err(ArrowError::ParseError(
270                            "Expecting a typeIds for union ".to_string(),
271                        ))
272                    }
273                } else {
274                    Err(ArrowError::ParseError(
275                        "Expecting a mode for union".to_string(),
276                    ))
277                }
278            }
279            Some(other) => Err(ArrowError::ParseError(format!(
280                "invalid or unsupported type name: {other} in {json:?}"
281            ))),
282            None => Err(ArrowError::ParseError("type name missing".to_string())),
283        },
284        _ => Err(ArrowError::ParseError(
285            "invalid json value type".to_string(),
286        )),
287    }
288}
289
290/// Generate a JSON representation of the data type.
291pub fn data_type_to_json(data_type: &DataType) -> serde_json::Value {
292    use serde_json::json;
293    match data_type {
294        DataType::Null => json!({"name": "null"}),
295        DataType::Boolean => json!({"name": "bool"}),
296        DataType::Int8 => json!({"name": "int", "bitWidth": 8, "isSigned": true}),
297        DataType::Int16 => json!({"name": "int", "bitWidth": 16, "isSigned": true}),
298        DataType::Int32 => json!({"name": "int", "bitWidth": 32, "isSigned": true}),
299        DataType::Int64 => json!({"name": "int", "bitWidth": 64, "isSigned": true}),
300        DataType::UInt8 => json!({"name": "int", "bitWidth": 8, "isSigned": false}),
301        DataType::UInt16 => json!({"name": "int", "bitWidth": 16, "isSigned": false}),
302        DataType::UInt32 => json!({"name": "int", "bitWidth": 32, "isSigned": false}),
303        DataType::UInt64 => json!({"name": "int", "bitWidth": 64, "isSigned": false}),
304        DataType::Float16 => json!({"name": "floatingpoint", "precision": "HALF"}),
305        DataType::Float32 => json!({"name": "floatingpoint", "precision": "SINGLE"}),
306        DataType::Float64 => json!({"name": "floatingpoint", "precision": "DOUBLE"}),
307        DataType::Utf8 => json!({"name": "utf8"}),
308        DataType::LargeUtf8 => json!({"name": "largeutf8"}),
309        DataType::Binary => json!({"name": "binary"}),
310        DataType::LargeBinary => json!({"name": "largebinary"}),
311        DataType::BinaryView => json!({"name": "binaryview"}),
312        DataType::Utf8View => json!({"name": "utf8view"}),
313        DataType::FixedSizeBinary(byte_width) => {
314            json!({"name": "fixedsizebinary", "byteWidth": byte_width})
315        }
316        DataType::Struct(_) => json!({"name": "struct"}),
317        DataType::Union(_, _) => json!({"name": "union"}),
318        DataType::List(_) => json!({ "name": "list"}),
319        DataType::LargeList(_) => json!({ "name": "largelist"}),
320        DataType::ListView(_) => json!({ "name": "listview"}),
321        DataType::LargeListView(_) => json!({ "name": "largelistview"}),
322        DataType::FixedSizeList(_, length) => {
323            json!({"name":"fixedsizelist", "listSize": length})
324        }
325        DataType::Time32(unit) => {
326            json!({"name": "time", "bitWidth": 32, "unit": match unit {
327                TimeUnit::Second => "SECOND",
328                TimeUnit::Millisecond => "MILLISECOND",
329                TimeUnit::Microsecond => "MICROSECOND",
330                TimeUnit::Nanosecond => "NANOSECOND",
331            }})
332        }
333        DataType::Time64(unit) => {
334            json!({"name": "time", "bitWidth": 64, "unit": match unit {
335                TimeUnit::Second => "SECOND",
336                TimeUnit::Millisecond => "MILLISECOND",
337                TimeUnit::Microsecond => "MICROSECOND",
338                TimeUnit::Nanosecond => "NANOSECOND",
339            }})
340        }
341        DataType::Date32 => {
342            json!({"name": "date", "unit": "DAY"})
343        }
344        DataType::Date64 => {
345            json!({"name": "date", "unit": "MILLISECOND"})
346        }
347        DataType::Timestamp(unit, None) => {
348            json!({"name": "timestamp", "unit": match unit {
349                TimeUnit::Second => "SECOND",
350                TimeUnit::Millisecond => "MILLISECOND",
351                TimeUnit::Microsecond => "MICROSECOND",
352                TimeUnit::Nanosecond => "NANOSECOND",
353            }})
354        }
355        DataType::Timestamp(unit, Some(tz)) => {
356            json!({"name": "timestamp", "unit": match unit {
357                    TimeUnit::Second => "SECOND",
358                    TimeUnit::Millisecond => "MILLISECOND",
359                    TimeUnit::Microsecond => "MICROSECOND",
360                    TimeUnit::Nanosecond => "NANOSECOND",
361                }, "timezone": tz})
362        }
363        DataType::Interval(unit) => json!({"name": "interval", "unit": match unit {
364            IntervalUnit::YearMonth => "YEAR_MONTH",
365            IntervalUnit::DayTime => "DAY_TIME",
366            IntervalUnit::MonthDayNano => "MONTH_DAY_NANO",
367        }}),
368        DataType::Duration(unit) => json!({"name": "duration", "unit": match unit {
369            TimeUnit::Second => "SECOND",
370            TimeUnit::Millisecond => "MILLISECOND",
371            TimeUnit::Microsecond => "MICROSECOND",
372            TimeUnit::Nanosecond => "NANOSECOND",
373        }}),
374        DataType::Dictionary(_, _) => json!({ "name": "dictionary"}),
375        DataType::Decimal32(precision, scale) => {
376            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 32})
377        }
378        DataType::Decimal64(precision, scale) => {
379            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 64})
380        }
381        DataType::Decimal128(precision, scale) => {
382            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 128})
383        }
384        DataType::Decimal256(precision, scale) => {
385            json!({"name": "decimal", "precision": precision, "scale": scale, "bitWidth": 256})
386        }
387        DataType::Map(_, keys_sorted) => {
388            json!({"name": "map", "keysSorted": keys_sorted})
389        }
390        DataType::RunEndEncoded(_, _) => json!({"name": "runendencoded"}),
391    }
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397    use serde_json::Value;
398
399    #[test]
400    fn parse_utf8_from_json() {
401        let json = "{\"name\":\"utf8\"}";
402        let value: Value = serde_json::from_str(json).unwrap();
403        let dt = data_type_from_json(&value).unwrap();
404        assert_eq!(DataType::Utf8, dt);
405    }
406
407    #[test]
408    fn parse_int32_from_json() {
409        let json = "{\"name\": \"int\", \"isSigned\": true, \"bitWidth\": 32}";
410        let value: Value = serde_json::from_str(json).unwrap();
411        let dt = data_type_from_json(&value).unwrap();
412        assert_eq!(DataType::Int32, dt);
413    }
414}
arrow_integration_test/datatype.rs

arrow_integration_test/
datatype.rs
