prusto_rs/types/
mod.rs

1mod boolean;
2mod data_set;
3mod date_time;
4mod decimal;
5mod fixed_char;
6mod float;
7mod integer;
8mod interval_day_to_second;
9mod interval_year_to_month;
10mod ip_address;
11mod map;
12mod option;
13mod row;
14mod seq;
15mod string;
16pub(self) mod util;
17pub mod uuid;
18
19pub use self::uuid::*;
20pub use boolean::*;
21pub use data_set::*;
22pub use date_time::*;
23pub use decimal::*;
24pub use fixed_char::*;
25pub use float::*;
26pub use integer::*;
27pub use integer::*;
28pub use interval_day_to_second::*;
29pub use interval_year_to_month::*;
30pub use ip_address::*;
31pub use map::*;
32pub use option::*;
33pub use row::*;
34pub use seq::*;
35pub use string::*;
36
37//mod str;
38//pub use self::str::*;
39
40use std::borrow::Cow;
41use std::collections::HashMap;
42use std::iter::FromIterator;
43use std::sync::Arc;
44
45use derive_more::Display;
46use iterable::*;
47use serde::de::DeserializeSeed;
48use serde::Serialize;
49
50use crate::{
51    ClientTypeSignatureParameter, Column, NamedTypeSignature, RawPrestoTy, RowFieldName,
52    TypeSignature,
53};
54
55//TODO: refine it
56#[derive(Display, Debug)]
57pub enum Error {
58    InvalidPrestoType,
59    InvalidColumn,
60    InvalidTypeSignature,
61    ParseDecimalFailed(String),
62    ParseIntervalMonthFailed,
63    ParseIntervalDayFailed,
64    EmptyInPrestoRow,
65    NonePrestoRow,
66}
67
68pub trait Presto {
69    type ValueType<'a>: Serialize
70    where
71        Self: 'a;
72    type Seed<'a, 'de>: DeserializeSeed<'de, Value = Self>;
73
74    fn value(&self) -> Self::ValueType<'_>;
75
76    fn ty() -> PrestoTy;
77
78    /// caller must provide a valid context
79    fn seed<'a, 'de>(ctx: &'a Context<'a>) -> Self::Seed<'a, 'de>;
80
81    fn empty() -> Self;
82}
83
84pub trait PrestoMapKey: Presto {}
85
86#[derive(Debug)]
87pub struct Context<'a> {
88    ty: &'a PrestoTy,
89    map: Arc<HashMap<usize, Vec<usize>>>,
90}
91
92impl<'a> Context<'a> {
93    pub fn new<T: Presto>(provided: &'a PrestoTy) -> Result<Self, Error> {
94        let target = T::ty();
95        let ret = extract(&target, provided)?;
96        let map = HashMap::from_iter(ret.into_iter());
97        Ok(Context {
98            ty: provided,
99            map: Arc::new(map),
100        })
101    }
102
103    pub fn with_ty(&'a self, ty: &'a PrestoTy) -> Context<'a> {
104        Context {
105            ty,
106            map: self.map.clone(),
107        }
108    }
109
110    pub fn ty(&self) -> &PrestoTy {
111        self.ty
112    }
113
114    pub fn row_map(&self) -> Option<&[usize]> {
115        let key = self.ty as *const PrestoTy as usize;
116        self.map.get(&key).map(|r| &**r)
117    }
118}
119
120fn extract(target: &PrestoTy, provided: &PrestoTy) -> Result<Vec<(usize, Vec<usize>)>, Error> {
121    use PrestoTy::*;
122
123    match (target, provided) {
124        (Unknown, _) => Ok(vec![]),
125        (Decimal(p1, s1), Decimal(p2, s2)) if p1 == p2 && s1 == s2 => Ok(vec![]),
126        (Option(ty), provided) => extract(ty, provided),
127        (Boolean, Boolean) => Ok(vec![]),
128        (Date, Date) => Ok(vec![]),
129        (Time, Time) => Ok(vec![]),
130        (Timestamp, Timestamp) => Ok(vec![]),
131        (IntervalYearToMonth, IntervalYearToMonth) => Ok(vec![]),
132        (IntervalDayToSecond, IntervalDayToSecond) => Ok(vec![]),
133        (PrestoInt(_), PrestoInt(_)) => Ok(vec![]),
134        (PrestoFloat(_), PrestoFloat(_)) => Ok(vec![]),
135        (Varchar, Varchar) => Ok(vec![]),
136        (Char(a), Char(b)) if a == b => Ok(vec![]),
137        (Tuple(t1), Tuple(t2)) => {
138            if t1.len() != t2.len() {
139                Err(Error::InvalidPrestoType)
140            } else {
141                t1.lazy_zip(t2).try_flat_map(|(l, r)| extract(l, r))
142            }
143        }
144        (Row(t1), Row(t2)) => {
145            if t1.len() != t2.len() {
146                Err(Error::InvalidPrestoType)
147            } else {
148                // create a vector of the original element's reference
149                let t1k = t1.sorted_by(|t1, t2| Ord::cmp(&t1.0, &t2.0));
150                let t2k = t2.sorted_by(|t1, t2| Ord::cmp(&t1.0, &t2.0));
151
152                let ret = t1k.lazy_zip(t2k).try_flat_map(|(l, r)| {
153                    if l.0 == r.0 {
154                        extract(&l.1, &r.1)
155                    } else {
156                        Err(Error::InvalidPrestoType)
157                    }
158                })?;
159
160                let map = t2.map(|provided| t1.position(|target| provided.0 == target.0).unwrap());
161                let key = provided as *const PrestoTy as usize;
162                Ok(ret.add_one((key, map)))
163            }
164        }
165        (Array(t1), Array(t2)) => extract(t1, t2),
166        (Map(t1k, t1v), Map(t2k, t2v)) => Ok(extract(t1k, t2k)?.chain(extract(t1v, t2v)?)),
167        (IpAddress, IpAddress) => Ok(vec![]),
168        (Uuid, Uuid) => Ok(vec![]),
169        _ => Err(Error::InvalidPrestoType),
170    }
171}
172
173// TODO:
174// VarBinary Json
175// TimestampWithTimeZone TimeWithTimeZone
176// HyperLogLog P4HyperLogLog
177// QDigest
178#[derive(Clone, Debug, Eq, PartialEq)]
179pub enum PrestoTy {
180    Date,
181    Time,
182    Timestamp,
183    Uuid,
184    IntervalYearToMonth,
185    IntervalDayToSecond,
186    Option(Box<PrestoTy>),
187    Boolean,
188    PrestoInt(PrestoInt),
189    PrestoFloat(PrestoFloat),
190    Varchar,
191    Char(usize),
192    Tuple(Vec<PrestoTy>),
193    Row(Vec<(String, PrestoTy)>),
194    Array(Box<PrestoTy>),
195    Map(Box<PrestoTy>, Box<PrestoTy>),
196    Decimal(usize, usize),
197    IpAddress,
198    Unknown,
199}
200
201#[derive(Clone, Debug, Eq, PartialEq)]
202pub enum PrestoInt {
203    I8,
204    I16,
205    I32,
206    I64,
207}
208
209#[derive(Clone, Debug, Eq, PartialEq)]
210pub enum PrestoFloat {
211    F32,
212    F64,
213}
214
215impl PrestoTy {
216    pub fn from_type_signature(mut sig: TypeSignature) -> Result<Self, Error> {
217        use PrestoFloat::*;
218        use PrestoInt::*;
219
220        let ty = match sig.raw_type {
221            RawPrestoTy::Date => PrestoTy::Date,
222            RawPrestoTy::Time => PrestoTy::Time,
223            RawPrestoTy::Timestamp => PrestoTy::Timestamp,
224            RawPrestoTy::IntervalYearToMonth => PrestoTy::IntervalYearToMonth,
225            RawPrestoTy::IntervalDayToSecond => PrestoTy::IntervalDayToSecond,
226            RawPrestoTy::Unknown => PrestoTy::Unknown,
227            RawPrestoTy::Decimal if sig.arguments.len() == 2 => {
228                let s_sig = sig.arguments.pop().unwrap();
229                let p_sig = sig.arguments.pop().unwrap();
230                if let (
231                    ClientTypeSignatureParameter::LongLiteral(p),
232                    ClientTypeSignatureParameter::LongLiteral(s),
233                ) = (p_sig, s_sig)
234                {
235                    PrestoTy::Decimal(p as usize, s as usize)
236                } else {
237                    return Err(Error::InvalidTypeSignature);
238                }
239            }
240            RawPrestoTy::Boolean => PrestoTy::Boolean,
241            RawPrestoTy::TinyInt => PrestoTy::PrestoInt(I8),
242            RawPrestoTy::SmallInt => PrestoTy::PrestoInt(I16),
243            RawPrestoTy::Integer => PrestoTy::PrestoInt(I32),
244            RawPrestoTy::BigInt => PrestoTy::PrestoInt(I64),
245            RawPrestoTy::Real => PrestoTy::PrestoFloat(F32),
246            RawPrestoTy::Double => PrestoTy::PrestoFloat(F64),
247            RawPrestoTy::VarChar => PrestoTy::Varchar,
248            RawPrestoTy::Char if sig.arguments.len() == 1 => {
249                if let ClientTypeSignatureParameter::LongLiteral(p) = sig.arguments.pop().unwrap() {
250                    PrestoTy::Char(p as usize)
251                } else {
252                    return Err(Error::InvalidTypeSignature);
253                }
254            }
255            RawPrestoTy::Array if sig.arguments.len() == 1 => {
256                let sig = sig.arguments.pop().unwrap();
257                if let ClientTypeSignatureParameter::TypeSignature(sig) = sig {
258                    let inner = Self::from_type_signature(sig)?;
259                    PrestoTy::Array(Box::new(inner))
260                } else {
261                    return Err(Error::InvalidTypeSignature);
262                }
263            }
264            RawPrestoTy::Map if sig.arguments.len() == 2 => {
265                let v_sig = sig.arguments.pop().unwrap();
266                let k_sig = sig.arguments.pop().unwrap();
267                if let (
268                    ClientTypeSignatureParameter::TypeSignature(k_sig),
269                    ClientTypeSignatureParameter::TypeSignature(v_sig),
270                ) = (k_sig, v_sig)
271                {
272                    let k_inner = Self::from_type_signature(k_sig)?;
273                    let v_inner = Self::from_type_signature(v_sig)?;
274                    PrestoTy::Map(Box::new(k_inner), Box::new(v_inner))
275                } else {
276                    return Err(Error::InvalidTypeSignature);
277                }
278            }
279            RawPrestoTy::Row if !sig.arguments.is_empty() => {
280                let ir = sig.arguments.try_map(|arg| match arg {
281                    ClientTypeSignatureParameter::NamedTypeSignature(sig) => {
282                        let name = sig.field_name.map(|n| n.name);
283                        let ty = Self::from_type_signature(sig.type_signature)?;
284                        Ok((name, ty))
285                    }
286                    _ => Err(Error::InvalidTypeSignature),
287                })?;
288
289                let is_named = ir[0].0.is_some();
290
291                if is_named {
292                    let row = ir.try_map(|(name, ty)| match name {
293                        Some(n) => Ok((n, ty)),
294                        None => Err(Error::InvalidTypeSignature),
295                    })?;
296                    PrestoTy::Row(row)
297                } else {
298                    let tuple = ir.try_map(|(name, ty)| match name {
299                        Some(_) => Err(Error::InvalidTypeSignature),
300                        None => Ok(ty),
301                    })?;
302                    PrestoTy::Tuple(tuple)
303                }
304            }
305            RawPrestoTy::IpAddress => PrestoTy::IpAddress,
306            RawPrestoTy::Uuid => PrestoTy::Uuid,
307            _ => return Err(Error::InvalidTypeSignature),
308        };
309
310        Ok(ty)
311    }
312
313    pub fn from_column(column: Column) -> Result<(String, Self), Error> {
314        let name = column.name;
315        if let Some(sig) = column.type_signature {
316            let ty = Self::from_type_signature(sig)?;
317            Ok((name, ty))
318        } else {
319            Err(Error::InvalidColumn)
320        }
321    }
322
323    pub fn from_columns(columns: Vec<Column>) -> Result<Self, Error> {
324        let row = columns.try_map(Self::from_column)?;
325        Ok(PrestoTy::Row(row))
326    }
327
328    pub fn into_type_signature(self) -> TypeSignature {
329        use PrestoTy::*;
330
331        let raw_ty = self.raw_type();
332
333        let params = match self {
334            Unknown => vec![],
335            Decimal(p, s) => vec![
336                ClientTypeSignatureParameter::LongLiteral(p as u64),
337                ClientTypeSignatureParameter::LongLiteral(s as u64),
338            ],
339            Date => vec![],
340            Time => vec![],
341            Timestamp => vec![],
342            IntervalYearToMonth => vec![],
343            IntervalDayToSecond => vec![],
344            Option(t) => return t.into_type_signature(),
345            Boolean => vec![],
346            PrestoInt(_) => vec![],
347            PrestoFloat(_) => vec![],
348            Varchar => vec![ClientTypeSignatureParameter::LongLiteral(2147483647)],
349            Char(a) => vec![ClientTypeSignatureParameter::LongLiteral(a as u64)],
350            Tuple(ts) => ts.map(|ty| {
351                ClientTypeSignatureParameter::NamedTypeSignature(NamedTypeSignature {
352                    field_name: None,
353                    type_signature: ty.into_type_signature(),
354                })
355            }),
356            Row(ts) => ts.map(|(name, ty)| {
357                ClientTypeSignatureParameter::NamedTypeSignature(NamedTypeSignature {
358                    field_name: Some(RowFieldName::new(name)),
359                    type_signature: ty.into_type_signature(),
360                })
361            }),
362            Array(t) => vec![ClientTypeSignatureParameter::TypeSignature(
363                t.into_type_signature(),
364            )],
365            Map(t1, t2) => vec![
366                ClientTypeSignatureParameter::TypeSignature(t1.into_type_signature()),
367                ClientTypeSignatureParameter::TypeSignature(t2.into_type_signature()),
368            ],
369            IpAddress => vec![],
370            Uuid => vec![],
371        };
372
373        TypeSignature::new(raw_ty, params)
374    }
375
376    pub fn full_type(&self) -> Cow<'static, str> {
377        use PrestoTy::*;
378
379        match self {
380            Unknown => RawPrestoTy::Unknown.to_str().into(),
381            Decimal(p, s) => format!("{}({},{})", RawPrestoTy::Decimal.to_str(), p, s).into(),
382            Option(t) => t.full_type(),
383            Date => RawPrestoTy::Date.to_str().into(),
384            Time => RawPrestoTy::Time.to_str().into(),
385            Timestamp => RawPrestoTy::Timestamp.to_str().into(),
386            IntervalYearToMonth => RawPrestoTy::IntervalYearToMonth.to_str().into(),
387            IntervalDayToSecond => RawPrestoTy::IntervalDayToSecond.to_str().into(),
388            Boolean => RawPrestoTy::Boolean.to_str().into(),
389            PrestoInt(ty) => ty.raw_type().to_str().into(),
390            PrestoFloat(ty) => ty.raw_type().to_str().into(),
391            Varchar => RawPrestoTy::VarChar.to_str().into(),
392            Char(a) => format!("{}({})", RawPrestoTy::Char.to_str(), a).into(),
393            Tuple(ts) => format!(
394                "{}({})",
395                RawPrestoTy::Row.to_str(),
396                ts.lazy_map(|ty| ty.full_type()).join(",")
397            )
398            .into(),
399            Row(ts) => format!(
400                "{}({})",
401                RawPrestoTy::Row.to_str(),
402                ts.lazy_map(|(name, ty)| format!("{} {}", name, ty.full_type()))
403                    .join(",")
404            )
405            .into(),
406            Array(t) => format!("{}({})", RawPrestoTy::Array.to_str(), t.full_type()).into(),
407            Map(t1, t2) => format!(
408                "{}({},{})",
409                RawPrestoTy::Map.to_str(),
410                t1.full_type(),
411                t2.full_type()
412            )
413            .into(),
414            IpAddress => RawPrestoTy::IpAddress.to_str().into(),
415            Uuid => RawPrestoTy::Uuid.to_str().into(),
416        }
417    }
418
419    pub fn raw_type(&self) -> RawPrestoTy {
420        use PrestoTy::*;
421
422        match self {
423            Unknown => RawPrestoTy::Unknown,
424            Date => RawPrestoTy::Date,
425            Time => RawPrestoTy::Time,
426            Timestamp => RawPrestoTy::Timestamp,
427            IntervalYearToMonth => RawPrestoTy::IntervalYearToMonth,
428            IntervalDayToSecond => RawPrestoTy::IntervalDayToSecond,
429            Decimal(_, _) => RawPrestoTy::Decimal,
430            Option(ty) => ty.raw_type(),
431            Boolean => RawPrestoTy::Boolean,
432            PrestoInt(ty) => ty.raw_type(),
433            PrestoFloat(ty) => ty.raw_type(),
434            Varchar => RawPrestoTy::VarChar,
435            Char(_) => RawPrestoTy::Char,
436            Tuple(_) => RawPrestoTy::Row,
437            Row(_) => RawPrestoTy::Row,
438            Array(_) => RawPrestoTy::Array,
439            Map(_, _) => RawPrestoTy::Map,
440            IpAddress => RawPrestoTy::IpAddress,
441            Uuid => RawPrestoTy::Uuid,
442        }
443    }
444}
445
446impl PrestoInt {
447    pub fn raw_type(&self) -> RawPrestoTy {
448        use PrestoInt::*;
449        match self {
450            I8 => RawPrestoTy::TinyInt,
451            I16 => RawPrestoTy::SmallInt,
452            I32 => RawPrestoTy::Integer,
453            I64 => RawPrestoTy::BigInt,
454        }
455    }
456}
457
458impl PrestoFloat {
459    pub fn raw_type(&self) -> RawPrestoTy {
460        use PrestoFloat::*;
461        match self {
462            F32 => RawPrestoTy::Real,
463            F64 => RawPrestoTy::Double,
464        }
465    }
466}
prusto_rs/types/mod.rs

prusto_rs/types/
mod.rs
