1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

use datafusion::{
    arrow::datatypes::{DataType, Field},
    scalar::ScalarValue,
};
use serde_json::Value as JsonValue;
use toml::Value as TomlValue;

/// Extension trait for [`ScalarValue`](datafusion::scalar::ScalarValue).
pub trait ScalarValueExt {
    /// Constructs an instance from a TOML value.
    fn from_toml_value(value: TomlValue) -> ScalarValue;

    /// Constructs an instance from a JSON value.
    fn from_json_value(value: JsonValue) -> ScalarValue;
}

impl ScalarValueExt for ScalarValue {
    fn from_toml_value(value: TomlValue) -> ScalarValue {
        match value {
            TomlValue::Boolean(b) => Self::Boolean(Some(b)),
            TomlValue::Integer(i) => Self::Int64(Some(i)),
            TomlValue::Float(f) => Self::Float64(Some(f)),
            TomlValue::String(s) => Self::Utf8(Some(s)),
            TomlValue::Datetime(dt) => Self::Utf8(Some(dt.to_string())),
            TomlValue::Array(vec) => {
                let mut data_type = DataType::Null;
                let scalars = vec
                    .into_iter()
                    .map(|value| {
                        let scalar = Self::from_toml_value(value);
                        if data_type == DataType::Null {
                            data_type = scalar.get_datatype();
                        }
                        scalar
                    })
                    .collect::<Vec<_>>();
                Self::new_list(Some(scalars), data_type)
            }
            TomlValue::Table(table) => {
                let mut fields = Vec::with_capacity(table.len());
                let scalars = table
                    .into_iter()
                    .map(|(key, value)| {
                        let scalar = Self::from_toml_value(value);
                        let field = Field::new(key, scalar.get_datatype(), true);
                        fields.push(field);
                        scalar
                    })
                    .collect::<Vec<_>>();
                Self::Struct(Some(scalars), Box::new(fields))
            }
        }
    }

    /// Constructs an instance from a JSON value.
    fn from_json_value(value: JsonValue) -> ScalarValue {
        match value {
            JsonValue::Null => Self::Null,
            JsonValue::Bool(b) => Self::Boolean(Some(b)),
            JsonValue::Number(n) => {
                if let Some(i) = n.as_u64() {
                    Self::UInt64(Some(i))
                } else if let Some(i) = n.as_i64() {
                    Self::Int64(Some(i))
                } else if let Some(f) = n.as_f64() {
                    Self::Float64(Some(f))
                } else {
                    Self::Null
                }
            }
            JsonValue::String(s) => Self::Utf8(Some(s)),
            JsonValue::Array(vec) => {
                let mut data_type = DataType::Null;
                let scalars = vec
                    .into_iter()
                    .map(|value| {
                        let scalar = Self::from_json_value(value);
                        if data_type == DataType::Null {
                            data_type = scalar.get_datatype();
                        }
                        scalar
                    })
                    .collect::<Vec<_>>();
                Self::new_list(Some(scalars), data_type)
            }
            JsonValue::Object(map) => {
                let mut fields = Vec::with_capacity(map.len());
                let scalars = map
                    .into_iter()
                    .map(|(key, value)| {
                        let scalar = Self::from_json_value(value);
                        let field = Field::new(key, scalar.get_datatype(), true);
                        fields.push(field);
                        scalar
                    })
                    .collect::<Vec<_>>();
                Self::Struct(Some(scalars), Box::new(fields))
            }
        }
    }
}