use crate::ast::ddl::{DataType, VectorMetric};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ResolvedType {
Integer,
BigInt,
Float,
Double,
Text,
Blob,
Boolean,
Timestamp,
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::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) => true,
(BigInt, Double) => true,
(Float, Double) => true,
(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::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::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));
}
#[test]
fn test_can_cast_incompatible() {
assert!(!ResolvedType::Text.can_cast_to(&ResolvedType::Integer));
assert!(!ResolvedType::BigInt.can_cast_to(&ResolvedType::Integer));
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)"
);
}
}