1use std::fmt;
2
3use arrow::array::Array as _;
4use arrow::datatypes::{DataType, Field, TimeUnit};
5use datafusion::common::ScalarValue;
6
7use crate::SqlError;
8
9#[derive(Clone, Debug, Eq, PartialEq)]
11pub enum SqlType {
12 Null,
14 Boolean,
16 SignedInteger(u8),
18 UnsignedInteger(u8),
20 Float(u8),
22 TimestampMillis(Option<String>),
24 Text,
26 Bytes,
28 List(Box<Self>),
30 Other(String),
32}
33
34impl SqlType {
35 pub(crate) fn from_arrow(data_type: &DataType) -> Self {
36 match data_type {
37 DataType::Null => Self::Null,
38 DataType::Boolean => Self::Boolean,
39 DataType::Int8 => Self::SignedInteger(8),
40 DataType::Int16 => Self::SignedInteger(16),
41 DataType::Int32 => Self::SignedInteger(32),
42 DataType::Int64 => Self::SignedInteger(64),
43 DataType::UInt8 => Self::UnsignedInteger(8),
44 DataType::UInt16 => Self::UnsignedInteger(16),
45 DataType::UInt32 => Self::UnsignedInteger(32),
46 DataType::UInt64 => Self::UnsignedInteger(64),
47 DataType::Float16 => Self::Float(16),
48 DataType::Float32 => Self::Float(32),
49 DataType::Float64 => Self::Float(64),
50 DataType::Timestamp(TimeUnit::Millisecond, timezone) => {
51 Self::TimestampMillis(timezone.as_ref().map(ToString::to_string))
52 }
53 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Self::Text,
54 DataType::Binary
55 | DataType::LargeBinary
56 | DataType::BinaryView
57 | DataType::FixedSizeBinary(_) => Self::Bytes,
58 DataType::List(field) | DataType::LargeList(field) => {
59 Self::List(Box::new(Self::from_arrow(field.data_type())))
60 }
61 other => Self::Other(other.to_string()),
62 }
63 }
64}
65
66#[derive(Clone, Debug, Eq, PartialEq)]
68pub struct SqlColumn {
69 pub name: String,
71 pub data_type: SqlType,
73 pub nullable: bool,
75}
76
77impl SqlColumn {
78 pub(crate) fn from_arrow(field: &Field) -> Self {
79 Self {
80 name: field.name().clone(),
81 data_type: SqlType::from_arrow(field.data_type()),
82 nullable: field.is_nullable(),
83 }
84 }
85}
86
87#[derive(Clone, Debug, PartialEq)]
89pub enum SqlValue {
90 Null,
92 Boolean(bool),
94 Integer(i64),
96 Unsigned(u64),
98 Float(f64),
100 TimestampMillis(i64),
102 Unspecified(String),
105 Text(String),
107 Bytes(Vec<u8>),
109 List(Vec<Self>),
112 Other(String),
114}
115
116pub type SqlRow = Vec<SqlValue>;
118
119impl SqlValue {
120 pub(crate) fn to_scalar(&self) -> Result<ScalarValue, SqlError> {
121 Ok(match self {
122 Self::Null => ScalarValue::Null,
123 Self::Boolean(value) => ScalarValue::Boolean(Some(*value)),
124 Self::Integer(value) => ScalarValue::Int64(Some(*value)),
125 Self::Unsigned(value) => ScalarValue::UInt64(Some(*value)),
126 Self::Float(value) => ScalarValue::Float64(Some(*value)),
127 Self::TimestampMillis(value) => {
128 ScalarValue::TimestampMillisecond(Some(*value), Some("UTC".into()))
129 }
130 Self::Unspecified(value) | Self::Text(value) | Self::Other(value) => {
131 ScalarValue::Utf8(Some(value.clone()))
132 }
133 Self::Bytes(value) => ScalarValue::Binary(Some(value.clone())),
134 Self::List(values) => {
135 let scalars = values
136 .iter()
137 .map(Self::to_scalar)
138 .collect::<Result<Vec<_>, _>>()?;
139 let data_type = scalars
140 .iter()
141 .find(|value| !value.is_null())
142 .map_or(DataType::Null, ScalarValue::data_type);
143 ScalarValue::List(ScalarValue::new_list(&scalars, &data_type, true))
144 }
145 })
146 }
147
148 pub(crate) fn from_scalar(value: ScalarValue) -> Result<Self, SqlError> {
149 Ok(match value {
150 ScalarValue::Null
151 | ScalarValue::Boolean(None)
152 | ScalarValue::Int8(None)
153 | ScalarValue::Int16(None)
154 | ScalarValue::Int32(None)
155 | ScalarValue::Int64(None)
156 | ScalarValue::UInt8(None)
157 | ScalarValue::UInt16(None)
158 | ScalarValue::UInt32(None)
159 | ScalarValue::UInt64(None)
160 | ScalarValue::Float16(None)
161 | ScalarValue::Float32(None)
162 | ScalarValue::Float64(None)
163 | ScalarValue::Utf8(None)
164 | ScalarValue::Utf8View(None)
165 | ScalarValue::LargeUtf8(None)
166 | ScalarValue::Binary(None)
167 | ScalarValue::BinaryView(None)
168 | ScalarValue::LargeBinary(None)
169 | ScalarValue::TimestampMillisecond(None, _)
170 | ScalarValue::FixedSizeBinary(_, None) => Self::Null,
171 ScalarValue::Boolean(Some(value)) => Self::Boolean(value),
172 ScalarValue::Int8(Some(value)) => Self::Integer(i64::from(value)),
173 ScalarValue::Int16(Some(value)) => Self::Integer(i64::from(value)),
174 ScalarValue::Int32(Some(value)) => Self::Integer(i64::from(value)),
175 ScalarValue::Int64(Some(value)) => Self::Integer(value),
176 ScalarValue::UInt8(Some(value)) => Self::Unsigned(u64::from(value)),
177 ScalarValue::UInt16(Some(value)) => Self::Unsigned(u64::from(value)),
178 ScalarValue::UInt32(Some(value)) => Self::Unsigned(u64::from(value)),
179 ScalarValue::UInt64(Some(value)) => Self::Unsigned(value),
180 ScalarValue::Float16(Some(value)) => Self::Float(f64::from(value)),
181 ScalarValue::Float32(Some(value)) => Self::Float(f64::from(value)),
182 ScalarValue::Float64(Some(value)) => Self::Float(value),
183 ScalarValue::TimestampMillisecond(Some(value), _) => Self::TimestampMillis(value),
184 ScalarValue::Utf8(Some(value))
185 | ScalarValue::Utf8View(Some(value))
186 | ScalarValue::LargeUtf8(Some(value)) => Self::Text(value),
187 ScalarValue::Binary(Some(value))
188 | ScalarValue::BinaryView(Some(value))
189 | ScalarValue::LargeBinary(Some(value))
190 | ScalarValue::FixedSizeBinary(_, Some(value)) => Self::Bytes(value),
191 ScalarValue::List(array) => {
192 if array.is_null(0) {
193 Self::Null
194 } else {
195 let values = array.value(0);
196 let mut result = Vec::with_capacity(values.len());
197 for index in 0..values.len() {
198 result.push(Self::from_scalar(ScalarValue::try_from_array(
199 values.as_ref(),
200 index,
201 )?)?);
202 }
203 Self::List(result)
204 }
205 }
206 other => Self::Other(other.to_string()),
207 })
208 }
209}
210
211impl fmt::Display for SqlValue {
212 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
213 match self {
214 Self::Null => formatter.write_str("NULL"),
215 Self::Boolean(value) => write!(formatter, "{value}"),
216 Self::Integer(value) | Self::TimestampMillis(value) => write!(formatter, "{value}"),
217 Self::Unsigned(value) => write!(formatter, "{value}"),
218 Self::Float(value) => write!(formatter, "{value}"),
219 Self::Unspecified(value) | Self::Text(value) | Self::Other(value) => {
220 formatter.write_str(value)
221 }
222 Self::Bytes(value) => {
223 for byte in value {
224 write!(formatter, "{byte:02x}")?;
225 }
226 Ok(())
227 }
228 Self::List(values) => {
229 formatter.write_str("[")?;
230 for (index, value) in values.iter().enumerate() {
231 if index > 0 {
232 formatter.write_str(", ")?;
233 }
234 write!(formatter, "{value}")?;
235 }
236 formatter.write_str("]")
237 }
238 }
239 }
240}