Skip to main content

corium_sql/
value.rs

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/// SQL result type independent of `DataFusion` and Arrow versions.
10#[derive(Clone, Debug, Eq, PartialEq)]
11pub enum SqlType {
12    /// SQL NULL.
13    Null,
14    /// Boolean.
15    Boolean,
16    /// Signed integer, with its bit width.
17    SignedInteger(u8),
18    /// Unsigned integer, with its bit width.
19    UnsignedInteger(u8),
20    /// Floating-point value, with its bit width.
21    Float(u8),
22    /// Milliseconds since the Unix epoch, with an optional timezone.
23    TimestampMillis(Option<String>),
24    /// UTF-8 text.
25    Text,
26    /// Arbitrary bytes.
27    Bytes,
28    /// Variable-length list.
29    List(Box<Self>),
30    /// An engine type not yet represented structurally by this API.
31    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/// One result column.
67#[derive(Clone, Debug, Eq, PartialEq)]
68pub struct SqlColumn {
69    /// Projected column name.
70    pub name: String,
71    /// SQL result type.
72    pub data_type: SqlType,
73    /// Whether the result may contain NULL.
74    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/// One SQL value independent of `DataFusion` and Arrow versions.
88#[derive(Clone, Debug, PartialEq)]
89pub enum SqlValue {
90    /// SQL NULL.
91    Null,
92    /// Boolean.
93    Boolean(bool),
94    /// Signed integer normalized to 64 bits.
95    Integer(i64),
96    /// Unsigned integer normalized to 64 bits.
97    Unsigned(u64),
98    /// Floating-point value normalized to 64 bits.
99    Float(f64),
100    /// Milliseconds since the Unix epoch.
101    TimestampMillis(i64),
102    /// Untyped `PostgreSQL` text awaiting coercion by the SQL expression or
103    /// mutation target that consumes it.
104    Unspecified(String),
105    /// UTF-8 text.
106    Text(String),
107    /// Arbitrary bytes.
108    Bytes(Vec<u8>),
109    /// Variable-length list. Corium cardinality-many columns use unique,
110    /// deterministically ordered lists with set semantics.
111    List(Vec<Self>),
112    /// Text rendering of an engine value not yet represented structurally.
113    Other(String),
114}
115
116/// One SQL result row.
117pub 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}