substrait_explain/parser/
types.rs

1use pest::iterators::Pair;
2use substrait::proto::r#type::{Kind, Nullability, Parameter};
3use substrait::proto::{self, Type};
4
5use super::{ParsePair, Rule, ScopedParsePair, iter_pairs, unwrap_single_pair};
6use crate::extensions::SimpleExtensions;
7use crate::extensions::simple::ExtensionKind;
8use crate::parser::{ErrorKind, MessageParseError};
9
10// Given a name and an optional anchor, get the anchor and validate it. Errors will be pushed to the Scope error accumulator,
11// and the anchor will be returned if it is valid.
12pub(crate) fn get_and_validate_anchor(
13    extensions: &SimpleExtensions,
14    kind: ExtensionKind,
15    anchor: Option<u32>,
16    name: &str,
17    span: pest::Span,
18) -> Result<u32, MessageParseError> {
19    match anchor {
20        Some(a) => match extensions.is_name_unique(kind, a, name) {
21            Ok(_) => Ok(a),
22            Err(e) => {
23                let message = "Error matching name to anchor".to_string();
24                let error = MessageParseError {
25                    message: kind.name(),
26                    kind: ErrorKind::Lookup(e),
27                    error: Box::new(pest::error::Error::new_from_span(
28                        pest::error::ErrorVariant::CustomError { message },
29                        span,
30                    )),
31                };
32                Err(error)
33            }
34        },
35        None => match extensions.find_by_name(kind, name) {
36            Ok(a) => Ok(a),
37            Err(e) => {
38                let message = "Error finding extension for name".to_string();
39                let error = MessageParseError {
40                    message: kind.name(),
41                    kind: ErrorKind::Lookup(e),
42                    error: Box::new(pest::error::Error::new_from_span(
43                        pest::error::ErrorVariant::CustomError { message },
44                        span,
45                    )),
46                };
47                Err(error)
48            }
49        },
50    }
51}
52
53impl ParsePair for Nullability {
54    fn rule() -> Rule {
55        Rule::nullability
56    }
57
58    fn message() -> &'static str {
59        "Nullability"
60    }
61
62    fn parse_pair(pair: Pair<Rule>) -> Self {
63        assert_eq!(pair.as_rule(), Rule::nullability);
64        match pair.as_str() {
65            "?" => Nullability::Nullable,
66            "" => Nullability::Required,
67            "⁉" => Nullability::Unspecified,
68            _ => panic!("Invalid nullability: {}", pair.as_str()),
69        }
70    }
71}
72
73impl ScopedParsePair for Parameter {
74    fn rule() -> Rule {
75        Rule::parameter
76    }
77
78    fn message() -> &'static str {
79        "Parameter"
80    }
81
82    fn parse_pair(
83        extensions: &SimpleExtensions,
84        pair: Pair<Rule>,
85    ) -> Result<Self, MessageParseError> {
86        assert_eq!(pair.as_rule(), Rule::parameter);
87        let inner = unwrap_single_pair(pair);
88        match inner.as_rule() {
89            Rule::r#type => Ok(Parameter {
90                parameter: Some(proto::r#type::parameter::Parameter::DataType(
91                    Type::parse_pair(extensions, inner)?,
92                )),
93            }),
94            _ => unimplemented!("{:?}", inner.as_rule()),
95        }
96    }
97}
98
99fn parse_simple_type(pair: Pair<Rule>) -> Type {
100    assert_eq!(pair.as_rule(), Rule::simple_type);
101    let mut iter = iter_pairs(pair.into_inner());
102    let name = iter.pop(Rule::simple_type_name).as_str();
103    let nullability = iter.parse_next::<Nullability>();
104    iter.done();
105
106    let kind = match name {
107        "boolean" => Kind::Bool(proto::r#type::Boolean {
108            nullability: nullability.into(),
109            type_variation_reference: 0,
110        }),
111        "i64" => Kind::I64(proto::r#type::I64 {
112            nullability: nullability.into(),
113            type_variation_reference: 0,
114        }),
115        "i32" => Kind::I32(proto::r#type::I32 {
116            nullability: nullability.into(),
117            type_variation_reference: 0,
118        }),
119        "i16" => Kind::I16(proto::r#type::I16 {
120            nullability: nullability.into(),
121            type_variation_reference: 0,
122        }),
123        "i8" => Kind::I8(proto::r#type::I8 {
124            nullability: nullability.into(),
125            type_variation_reference: 0,
126        }),
127        "fp32" => Kind::Fp32(proto::r#type::Fp32 {
128            nullability: nullability.into(),
129            type_variation_reference: 0,
130        }),
131        "fp64" => Kind::Fp64(proto::r#type::Fp64 {
132            nullability: nullability.into(),
133            type_variation_reference: 0,
134        }),
135        "string" => Kind::String(proto::r#type::String {
136            nullability: nullability.into(),
137            type_variation_reference: 0,
138        }),
139        "binary" => Kind::Binary(proto::r#type::Binary {
140            nullability: nullability.into(),
141            type_variation_reference: 0,
142        }),
143        #[allow(deprecated)]
144        "timestamp" => Kind::Timestamp(proto::r#type::Timestamp {
145            nullability: nullability.into(),
146            type_variation_reference: 0,
147        }),
148        #[allow(deprecated)]
149        "timestamp_tz" => Kind::TimestampTz(proto::r#type::TimestampTz {
150            nullability: nullability.into(),
151            type_variation_reference: 0,
152        }),
153        "date" => Kind::Date(proto::r#type::Date {
154            nullability: nullability.into(),
155            type_variation_reference: 0,
156        }),
157        "time" => Kind::Time(proto::r#type::Time {
158            nullability: nullability.into(),
159            type_variation_reference: 0,
160        }),
161        "interval_year" => Kind::IntervalYear(proto::r#type::IntervalYear {
162            nullability: nullability.into(),
163            type_variation_reference: 0,
164        }),
165        "uuid" => Kind::Uuid(proto::r#type::Uuid {
166            nullability: nullability.into(),
167            type_variation_reference: 0,
168        }),
169        _ => unreachable!("Type {} exists in parser but not implemented in code", name),
170    };
171    Type { kind: Some(kind) }
172}
173
174fn parse_compound_type(
175    extensions: &SimpleExtensions,
176    pair: Pair<Rule>,
177) -> Result<Type, MessageParseError> {
178    assert_eq!(pair.as_rule(), Rule::compound_type);
179    let inner = unwrap_single_pair(pair);
180    match inner.as_rule() {
181        Rule::list_type => parse_list_type(extensions, inner),
182        // Rule::map_type => parse_map_type(inner),
183        // Rule::struct_type => parse_struct_type(inner),
184        _ => unimplemented!("{:?}", inner.as_rule()),
185    }
186}
187
188fn parse_list_type(
189    extensions: &SimpleExtensions,
190    pair: Pair<Rule>,
191) -> Result<Type, MessageParseError> {
192    assert_eq!(pair.as_rule(), Rule::list_type);
193    let mut iter = iter_pairs(pair.into_inner());
194    let nullability = iter.parse_next::<Nullability>();
195    let inner = iter.parse_next_scoped::<Type>(extensions)?;
196    iter.done();
197
198    Ok(Type {
199        kind: Some(Kind::List(Box::new(proto::r#type::List {
200            nullability: nullability.into(),
201            r#type: Some(Box::new(inner)),
202            type_variation_reference: 0,
203        }))),
204    })
205}
206
207fn parse_parameters(
208    extensions: &SimpleExtensions,
209    pair: Pair<Rule>,
210) -> Result<Vec<Parameter>, MessageParseError> {
211    assert_eq!(pair.as_rule(), Rule::parameters);
212    let mut iter = iter_pairs(pair.into_inner());
213    let mut params = Vec::new();
214    while let Some(param) = iter.parse_if_next_scoped::<Parameter>(extensions) {
215        params.push(param?);
216    }
217    iter.done();
218    Ok(params)
219}
220
221fn parse_user_defined_type(
222    extensions: &SimpleExtensions,
223    pair: Pair<Rule>,
224) -> Result<Type, MessageParseError> {
225    let span = pair.as_span();
226    assert_eq!(pair.as_rule(), Rule::user_defined_type);
227    let mut iter = iter_pairs(pair.into_inner());
228    let name = iter.pop(Rule::name).as_str();
229    let anchor = iter
230        .try_pop(Rule::anchor)
231        .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());
232
233    // TODO: Handle urn_anchor; validate that it matches the anchor
234    let _urn_anchor = iter
235        .try_pop(Rule::urn_anchor)
236        .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());
237
238    let nullability = iter.parse_next::<Nullability>();
239    let parameters = match iter.try_pop(Rule::parameters) {
240        Some(p) => parse_parameters(extensions, p)?,
241        None => Vec::new(),
242    };
243    iter.done();
244
245    let anchor = get_and_validate_anchor(extensions, ExtensionKind::Type, anchor, name, span)?;
246
247    Ok(Type {
248        kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
249            type_reference: anchor,
250            nullability: nullability.into(),
251            type_parameters: parameters,
252            type_variation_reference: 0,
253        })),
254    })
255}
256
257impl ScopedParsePair for Type {
258    fn rule() -> Rule {
259        Rule::r#type
260    }
261
262    fn message() -> &'static str {
263        "Type"
264    }
265
266    fn parse_pair(
267        extensions: &SimpleExtensions,
268        pair: Pair<Rule>,
269    ) -> Result<Self, MessageParseError> {
270        assert_eq!(pair.as_rule(), Rule::r#type);
271        let inner = unwrap_single_pair(pair);
272        match inner.as_rule() {
273            Rule::simple_type => Ok(parse_simple_type(inner)),
274            Rule::compound_type => parse_compound_type(extensions, inner),
275            Rule::user_defined_type => parse_user_defined_type(extensions, inner),
276            _ => unimplemented!("{:?}", inner.as_rule()),
277        }
278    }
279}
280
281#[cfg(test)]
282mod tests {
283    use pest::Parser;
284    use substrait::proto::r#type::{I64, Kind, Nullability};
285
286    use super::*;
287    use crate::parser::ExpressionParser;
288
289    #[test]
290    fn test_parse_simple_type() {
291        let mut pairs = ExpressionParser::parse(Rule::simple_type, "i64").unwrap();
292        let pair = pairs.next().unwrap();
293        assert_eq!(pairs.next(), None);
294        let t = parse_simple_type(pair);
295        assert_eq!(
296            t,
297            Type {
298                kind: Some(Kind::I64(I64 {
299                    nullability: Nullability::Required as i32,
300                    type_variation_reference: 0,
301                })),
302            }
303        );
304
305        let mut pairs = ExpressionParser::parse(Rule::simple_type, "string?").unwrap();
306        let pair = pairs.next().unwrap();
307        assert_eq!(pairs.next(), None);
308        let t = parse_simple_type(pair);
309        assert_eq!(
310            t,
311            Type {
312                kind: Some(Kind::String(proto::r#type::String {
313                    nullability: Nullability::Nullable as i32,
314                    type_variation_reference: 0,
315                })),
316            }
317        );
318    }
319
320    #[test]
321    fn test_parse_type() {
322        let extensions = SimpleExtensions::default();
323        let mut pairs = ExpressionParser::parse(Rule::r#type, "i64").unwrap();
324        let pair = pairs.next().unwrap();
325        assert_eq!(pairs.next(), None);
326        let t = Type::parse_pair(&extensions, pair).unwrap();
327        assert_eq!(
328            t,
329            Type {
330                kind: Some(Kind::I64(I64 {
331                    nullability: Nullability::Required as i32,
332                    type_variation_reference: 0,
333                }))
334            }
335        );
336    }
337
338    #[test]
339    fn test_parse_list_type() {
340        let extensions = SimpleExtensions::default();
341        let mut pairs = ExpressionParser::parse(Rule::list_type, "list<i64>").unwrap();
342        let pair = pairs.next().unwrap();
343        assert_eq!(pairs.next(), None);
344        let t = parse_list_type(&extensions, pair).unwrap();
345        assert_eq!(
346            t,
347            Type {
348                kind: Some(Kind::List(Box::new(proto::r#type::List {
349                    nullability: Nullability::Required as i32,
350                    r#type: Some(Box::new(Type {
351                        kind: Some(Kind::I64(I64 {
352                            nullability: Nullability::Required as i32,
353                            type_variation_reference: 0,
354                        }))
355                    })),
356                    type_variation_reference: 0,
357                })))
358            }
359        );
360    }
361
362    #[test]
363    fn test_parse_parameters() {
364        let extensions = SimpleExtensions::default();
365        let mut pairs = ExpressionParser::parse(Rule::parameters, "<i64?,string>").unwrap();
366        let pair = pairs.next().unwrap();
367        assert_eq!(pairs.next(), None);
368        let t = parse_parameters(&extensions, pair).unwrap();
369        assert_eq!(
370            t,
371            vec![
372                Parameter {
373                    parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
374                        kind: Some(Kind::I64(proto::r#type::I64 {
375                            nullability: Nullability::Nullable as i32,
376                            type_variation_reference: 0,
377                        })),
378                    })),
379                },
380                Parameter {
381                    parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
382                        kind: Some(Kind::String(proto::r#type::String {
383                            nullability: Nullability::Required as i32,
384                            type_variation_reference: 0,
385                        })),
386                    })),
387                },
388            ]
389        );
390    }
391
392    #[test]
393    fn test_udts() {
394        let mut extensions = SimpleExtensions::default();
395        extensions
396            .add_extension_urn("some_source".to_string(), 4)
397            .unwrap();
398        extensions
399            .add_extension(ExtensionKind::Type, 4, 42, "udt".to_string())
400            .unwrap();
401        let mut pairs = ExpressionParser::parse(Rule::user_defined_type, "udt#42<i64?>").unwrap();
402        let pair = pairs.next().unwrap();
403        assert_eq!(pairs.next(), None);
404
405        let t = parse_user_defined_type(&extensions, pair).unwrap();
406        assert_eq!(
407            t,
408            Type {
409                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
410                    type_reference: 42,
411                    type_variation_reference: 0,
412                    nullability: Nullability::Required as i32,
413                    type_parameters: vec![Parameter {
414                        parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
415                            kind: Some(Kind::I64(proto::r#type::I64 {
416                                nullability: Nullability::Nullable as i32,
417                                type_variation_reference: 0,
418                            })),
419                        })),
420                    }],
421                }))
422            }
423        );
424    }
425}
substrait_explain/parser/types.rs

substrait_explain/parser/
types.rs
