use crate::{DuckResult, duck_error};
use quack_rs::prelude::{LogicalType, TypeId};
use super::dynamic_value::DuckDynamicValue;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DuckTypeDesc {
Scalar(TypeId),
Decimal {
width: u8,
scale: u8,
},
List(Box<DuckTypeDesc>),
Struct(Vec<(String, DuckTypeDesc)>),
Map(Box<DuckTypeDesc>, Box<DuckTypeDesc>),
}
impl DuckTypeDesc {
#[must_use]
pub fn scalar(type_id: TypeId) -> Self {
Self::Scalar(type_id)
}
#[must_use]
pub fn list(element: Self) -> Self {
Self::List(Box::new(element))
}
#[must_use]
pub fn struct_type<I: IntoIterator<Item = (String, Self)>>(fields: I) -> Self {
Self::Struct(fields.into_iter().collect())
}
#[must_use]
pub fn map(key: Self, value: Self) -> Self {
Self::Map(Box::new(key), Box::new(value))
}
#[must_use]
pub fn type_id(&self) -> TypeId {
match self {
Self::Scalar(type_id) => *type_id,
Self::Decimal { .. } => TypeId::Decimal,
Self::List(_) => TypeId::List,
Self::Struct(_) => TypeId::Struct,
Self::Map(_, _) => TypeId::Map,
}
}
#[must_use]
pub fn to_logical_type(&self) -> LogicalType {
match self {
Self::Scalar(type_id) => LogicalType::new(*type_id),
Self::Decimal { width, scale } => LogicalType::decimal(*width, *scale),
Self::List(element) => LogicalType::list_from_logical(&element.to_logical_type()),
Self::Struct(fields) => {
let children: Vec<(&str, LogicalType)> = fields
.iter()
.map(|(name, desc)| (name.as_str(), desc.to_logical_type()))
.collect();
LogicalType::struct_type_from_logical(&children)
}
Self::Map(key, value) => {
LogicalType::map_from_logical(&key.to_logical_type(), &value.to_logical_type())
}
}
}
pub fn from_logical_type(logical_type: &LogicalType) -> DuckResult<Self> {
let raw = unsafe { libduckdb_sys::duckdb_get_type_id(logical_type.as_raw()) };
let Some(type_id) = TypeId::try_from_duckdb_type(raw) else {
return Err(duck_error(format!(
"dynamic column: unknown DuckDB type id {raw}; cannot build a \
DuckTypeDesc from this logical type"
)));
};
match type_id {
TypeId::List => Ok(Self::list(Self::from_logical_type(&unsafe {
logical_type.list_child_type()
})?)),
TypeId::Map => {
let key = Self::from_logical_type(&unsafe { logical_type.map_key_type() })?;
let value = Self::from_logical_type(&unsafe { logical_type.map_value_type() })?;
Ok(Self::map(key, value))
}
TypeId::Struct => {
let count = unsafe { logical_type.struct_child_count() };
let mut fields = Vec::with_capacity(count as usize);
for index in 0..count {
fields.push((
unsafe { logical_type.struct_child_name(index) },
Self::from_logical_type(&unsafe { logical_type.struct_child_type(index) })?,
));
}
Ok(Self::Struct(fields))
}
TypeId::Decimal => Ok(Self::Decimal {
width: unsafe { logical_type.decimal_width() },
scale: unsafe { logical_type.decimal_scale() },
}),
other => Self::from_scalar_type_id(other).ok_or_else(|| {
duck_error(format!(
"dynamic column: the DuckDB logical type `{other:?}` cannot be turned \
into a DuckTypeDesc (its parameters are not expressible); use a supported \
scalar, DECIMAL, LIST, STRUCT or MAP instead"
))
}),
}
}
fn from_scalar_type_id(type_id: TypeId) -> Option<Self> {
let supported = matches!(
type_id,
TypeId::Boolean
| TypeId::TinyInt
| TypeId::SmallInt
| TypeId::Integer
| TypeId::BigInt
| TypeId::UTinyInt
| TypeId::USmallInt
| TypeId::UInteger
| TypeId::UBigInt
| TypeId::HugeInt
| TypeId::UHugeInt
| TypeId::Float
| TypeId::Double
| TypeId::Varchar
| TypeId::Blob
| TypeId::Uuid
| TypeId::Date
| TypeId::Time
| TypeId::TimeTz
| TypeId::Timestamp
| TypeId::TimestampTz
| TypeId::TimestampS
| TypeId::TimestampMs
| TypeId::TimestampNs
| TypeId::Interval
);
supported.then_some(Self::Scalar(type_id))
}
#[must_use]
pub fn from_value(value: &DuckDynamicValue) -> Self {
match value {
DuckDynamicValue::List(items) => Self::list(
items
.iter()
.flatten()
.next()
.map_or(Self::Scalar(TypeId::Varchar), Self::from_value),
),
DuckDynamicValue::Struct(fields) => Self::Struct(
fields
.iter()
.enumerate()
.map(|(index, field)| {
let name = format!("field_{index}");
(name, field.as_ref().map_or(Self::Scalar(TypeId::Varchar), Self::from_value))
})
.collect(),
),
DuckDynamicValue::Map(pairs) => {
let key = pairs
.first()
.map_or(Self::Scalar(TypeId::Varchar), |(key, _)| Self::from_value(key));
let value = pairs
.first()
.map_or(Self::Scalar(TypeId::Varchar), |(_, value)| Self::from_value(value));
Self::map(key, value)
}
other => other.scalar_type_desc(),
}
}
#[must_use]
pub fn matches_value(&self, value: &DuckDynamicValue) -> bool {
match (self, value) {
(Self::Scalar(type_id), value) => value.scalar_type_id() == Some(*type_id),
(
Self::Decimal { width, scale },
DuckDynamicValue::Decimal {
width: value_width,
scale: value_scale,
..
},
) => width == value_width && scale == value_scale,
(Self::List(element), DuckDynamicValue::List(items)) => items
.iter()
.all(|item| item.as_ref().is_none_or(|item| element.matches_value(item))),
(Self::Struct(fields), DuckDynamicValue::Struct(values)) => {
fields.len() == values.len()
&& fields.iter().zip(values).all(|((_, desc), value)| {
value.as_ref().is_none_or(|value| desc.matches_value(value))
})
}
(Self::Map(key, map_value), DuckDynamicValue::Map(pairs)) => pairs
.iter()
.all(|(k, v)| key.matches_value(k) && map_value.matches_value(v)),
_ => false,
}
}
}
impl std::fmt::Display for DuckTypeDesc {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Scalar(type_id) => f.write_str(type_id.sql_name()),
Self::Decimal { width, scale } => write!(f, "DECIMAL({width}, {scale})"),
Self::List(element) => write!(f, "{element}[]"),
Self::Struct(fields) => {
f.write_str("STRUCT(")?;
for (index, (name, desc)) in fields.iter().enumerate() {
if index > 0 {
f.write_str(", ")?;
}
write!(f, "{name} {desc}")?;
}
f.write_str(")")
}
Self::Map(key, value) => write!(f, "MAP({key}, {value})"),
}
}
}