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

substrait_explain/parser/
types.rs
