use crate::ast::ddl::{DataType, VectorMetric};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ResolvedType {
Integer,
BigInt,
Float,
Double,
Text,
Blob,
Boolean,
Timestamp,
Date,
Time,
Interval,
Decimal { precision: u8, scale: u8 },
Json,
Array(Box<ResolvedType>),
Map {
key: Box<ResolvedType>,
value: Box<ResolvedType>,
},
Struct(Vec<(String, ResolvedType)>),
Vector {
dimension: u32,
metric: VectorMetric,
},
Null,
}
impl ResolvedType {
pub fn from_ast(dt: &DataType) -> Self {
match dt {
DataType::Integer | DataType::Int => Self::Integer,
DataType::BigInt => Self::BigInt,
DataType::Float => Self::Float,
DataType::Double => Self::Double,
DataType::Text => Self::Text,
DataType::Blob => Self::Blob,
DataType::Boolean | DataType::Bool => Self::Boolean,
DataType::Timestamp => Self::Timestamp,
DataType::Date => Self::Date,
DataType::Time => Self::Time,
DataType::Interval => Self::Interval,
DataType::Decimal { precision, scale } => Self::Decimal {
precision: *precision,
scale: *scale,
},
DataType::Json => Self::Json,
DataType::Array { element } => Self::Array(Box::new(Self::from_ast(element))),
DataType::Map { key, value } => Self::Map {
key: Box::new(Self::from_ast(key)),
value: Box::new(Self::from_ast(value)),
},
DataType::Struct { fields } => Self::Struct(
fields
.iter()
.map(|field| (field.name.clone(), Self::from_ast(&field.data_type)))
.collect(),
),
DataType::Vector { dimension, metric } => Self::Vector {
dimension: *dimension,
metric: metric.unwrap_or(VectorMetric::Cosine),
},
}
}
pub fn can_cast_to(&self, target: &ResolvedType) -> bool {
use ResolvedType::*;
match (self, target) {
(a, b) if a == b => true,
(Null, _) => true,
(Integer, BigInt | Float | Double | Decimal { .. }) => true,
(BigInt, Double) => true,
(Float, Double) => true,
(BigInt | Float | Double | Text | Decimal { .. }, Decimal { .. }) => true,
(Double, Float) => true,
(Text | Integer | BigInt | Float | Double, Timestamp) => true,
(Text, Date | Time | Interval) => true,
(Text, Json) | (Json, Text) => true,
(Array(source), Array(target)) => source.can_cast_to(target),
(
Map {
key: source_key,
value: source_value,
},
Map {
key: target_key,
value: target_value,
},
) => source_key.can_cast_to(target_key) && source_value.can_cast_to(target_value),
(Struct(source), Struct(target)) if source.len() == target.len() => source
.iter()
.zip(target)
.all(|((source_name, source_type), (target_name, target_type))| {
source_name == target_name && source_type.can_cast_to(target_type)
}),
(Vector { .. }, Vector { .. }) => false,
_ => false,
}
}
pub fn type_name(&self) -> &'static str {
match self {
Self::Integer => "Integer",
Self::BigInt => "BigInt",
Self::Float => "Float",
Self::Double => "Double",
Self::Text => "Text",
Self::Blob => "Blob",
Self::Boolean => "Boolean",
Self::Timestamp => "Timestamp",
Self::Date => "Date",
Self::Time => "Time",
Self::Interval => "Interval",
Self::Decimal { .. } => "Decimal",
Self::Json => "Json",
Self::Array(_) => "Array",
Self::Map { .. } => "Map",
Self::Struct(_) => "Struct",
Self::Vector { .. } => "Vector",
Self::Null => "Null",
}
}
}
impl std::fmt::Display for ResolvedType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Integer => write!(f, "INTEGER"),
Self::BigInt => write!(f, "BIGINT"),
Self::Float => write!(f, "FLOAT"),
Self::Double => write!(f, "DOUBLE"),
Self::Text => write!(f, "TEXT"),
Self::Blob => write!(f, "BLOB"),
Self::Boolean => write!(f, "BOOLEAN"),
Self::Timestamp => write!(f, "TIMESTAMP"),
Self::Date => write!(f, "DATE"),
Self::Time => write!(f, "TIME"),
Self::Interval => write!(f, "INTERVAL"),
Self::Decimal { precision, scale } => write!(f, "DECIMAL({precision},{scale})"),
Self::Json => write!(f, "JSON"),
Self::Array(element) => write!(f, "ARRAY<{element}>"),
Self::Map { key, value } => write!(f, "MAP<{key},{value}>"),
Self::Struct(fields) => {
write!(f, "STRUCT<")?;
for (index, (name, data_type)) in fields.iter().enumerate() {
if index > 0 {
write!(f, ",")?;
}
write!(f, "{name} {data_type}")?;
}
write!(f, ">")
}
Self::Vector { dimension, metric } => {
write!(f, "VECTOR({}, {:?})", dimension, metric)
}
Self::Null => write!(f, "NULL"),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_from_ast_integer() {
assert_eq!(
ResolvedType::from_ast(&DataType::Integer),
ResolvedType::Integer
);
assert_eq!(
ResolvedType::from_ast(&DataType::Int),
ResolvedType::Integer
);
}
#[test]
fn test_from_ast_boolean() {
assert_eq!(
ResolvedType::from_ast(&DataType::Boolean),
ResolvedType::Boolean
);
assert_eq!(
ResolvedType::from_ast(&DataType::Bool),
ResolvedType::Boolean
);
}
#[test]
fn test_from_ast_vector_with_metric() {
let dt = DataType::Vector {
dimension: 128,
metric: Some(VectorMetric::L2),
};
assert_eq!(
ResolvedType::from_ast(&dt),
ResolvedType::Vector {
dimension: 128,
metric: VectorMetric::L2,
}
);
}
#[test]
fn test_from_ast_vector_default_metric() {
let dt = DataType::Vector {
dimension: 256,
metric: None,
};
assert_eq!(
ResolvedType::from_ast(&dt),
ResolvedType::Vector {
dimension: 256,
metric: VectorMetric::Cosine,
}
);
}
#[test]
fn test_can_cast_same_type() {
assert!(ResolvedType::Integer.can_cast_to(&ResolvedType::Integer));
assert!(ResolvedType::Text.can_cast_to(&ResolvedType::Text));
}
#[test]
fn test_can_cast_null() {
assert!(ResolvedType::Null.can_cast_to(&ResolvedType::Integer));
assert!(ResolvedType::Null.can_cast_to(&ResolvedType::Text));
assert!(ResolvedType::Null.can_cast_to(&ResolvedType::Boolean));
}
#[test]
fn test_can_cast_numeric_widening() {
assert!(ResolvedType::Integer.can_cast_to(&ResolvedType::BigInt));
assert!(ResolvedType::Integer.can_cast_to(&ResolvedType::Float));
assert!(ResolvedType::Integer.can_cast_to(&ResolvedType::Double));
assert!(ResolvedType::BigInt.can_cast_to(&ResolvedType::Double));
assert!(ResolvedType::Float.can_cast_to(&ResolvedType::Double));
assert!(
ResolvedType::Decimal {
precision: 5,
scale: 3,
}
.can_cast_to(&ResolvedType::Decimal {
precision: 10,
scale: 2,
})
);
}
#[test]
fn test_can_cast_incompatible() {
assert!(!ResolvedType::Text.can_cast_to(&ResolvedType::Integer));
assert!(!ResolvedType::BigInt.can_cast_to(&ResolvedType::Integer));
}
#[test]
fn double_narrows_to_float_because_decimal_literals_are_double() {
assert!(ResolvedType::Double.can_cast_to(&ResolvedType::Float));
}
#[test]
fn test_can_cast_vector() {
let vec1 = ResolvedType::Vector {
dimension: 128,
metric: VectorMetric::Cosine,
};
let vec2 = ResolvedType::Vector {
dimension: 128,
metric: VectorMetric::L2,
};
assert!(!vec1.can_cast_to(&vec2));
}
#[test]
fn test_display() {
assert_eq!(format!("{}", ResolvedType::Integer), "INTEGER");
assert_eq!(format!("{}", ResolvedType::Text), "TEXT");
assert_eq!(
format!(
"{}",
ResolvedType::Vector {
dimension: 128,
metric: VectorMetric::Cosine
}
),
"VECTOR(128, Cosine)"
);
}
}