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        Rule::precision_timestamp_tz_type
193        | Rule::precision_timestamp_type
194        | Rule::precision_time_type => parse_precision_type(inner),
195        _ => unimplemented!("{:?}", inner.as_rule()),
196    }
197}
198
199fn parse_precision_type(pair: Pair<Rule>) -> Result<Type, MessageParseError> {
200    let rule = pair.as_rule();
201    let mut iter = iter_pairs(pair.into_inner());
202    let nullability = iter.parse_next::<Nullability>();
203    let precision_pair = iter.pop(Rule::integer);
204    let precision_span = precision_pair.as_span();
205    let precision = precision_pair.as_str().parse::<i32>().unwrap();
206    if !(0..=12).contains(&precision) {
207        return Err(MessageParseError::invalid(
208            "precision time type",
209            precision_span,
210            format!("precision must be between 0 and 12, got {precision}"),
211        ));
212    }
213    iter.done();
214    let kind = match rule {
215        Rule::precision_timestamp_type => {
216            Kind::PrecisionTimestamp(proto::r#type::PrecisionTimestamp {
217                precision,
218                nullability: nullability.into(),
219                type_variation_reference: 0,
220            })
221        }
222        Rule::precision_timestamp_tz_type => {
223            Kind::PrecisionTimestampTz(proto::r#type::PrecisionTimestampTz {
224                precision,
225                nullability: nullability.into(),
226                type_variation_reference: 0,
227            })
228        }
229        Rule::precision_time_type => Kind::PrecisionTime(proto::r#type::PrecisionTime {
230            precision,
231            nullability: nullability.into(),
232            type_variation_reference: 0,
233        }),
234        _ => unreachable!("parse_precision_type called with rule {:?}", rule),
235    };
236    Ok(Type { kind: Some(kind) })
237}
238
239fn parse_list_type(
240    extensions: &SimpleExtensions,
241    pair: Pair<Rule>,
242) -> Result<Type, MessageParseError> {
243    assert_eq!(pair.as_rule(), Rule::list_type);
244    let mut iter = iter_pairs(pair.into_inner());
245    let nullability = iter.parse_next::<Nullability>();
246    let inner = iter.parse_next_scoped::<Type>(extensions)?;
247    iter.done();
248
249    Ok(Type {
250        kind: Some(Kind::List(Box::new(proto::r#type::List {
251            nullability: nullability.into(),
252            r#type: Some(Box::new(inner)),
253            type_variation_reference: 0,
254        }))),
255    })
256}
257
258fn parse_parameters(
259    extensions: &SimpleExtensions,
260    pair: Pair<Rule>,
261) -> Result<Vec<Parameter>, MessageParseError> {
262    assert_eq!(pair.as_rule(), Rule::parameters);
263    let mut iter = iter_pairs(pair.into_inner());
264    let mut params = Vec::new();
265    while let Some(param) = iter.parse_if_next_scoped::<Parameter>(extensions) {
266        params.push(param?);
267    }
268    iter.done();
269    Ok(params)
270}
271
272fn parse_user_defined_type(
273    extensions: &SimpleExtensions,
274    pair: Pair<Rule>,
275) -> Result<Type, MessageParseError> {
276    let span = pair.as_span();
277    assert_eq!(pair.as_rule(), Rule::user_defined_type);
278    let mut iter = iter_pairs(pair.into_inner());
279    // TODO: quoted names not yet supported — see grammar TODO at user_defined_type
280    let name = iter.pop(Rule::identifier).as_str().to_string();
281    let anchor = iter
282        .try_pop(Rule::anchor)
283        .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());
284
285    // TODO: Handle urn_anchor; validate that it matches the anchor
286    let _urn_anchor = iter
287        .try_pop(Rule::urn_anchor)
288        .map(|n| unwrap_single_pair(n).as_str().parse::<u32>().unwrap());
289
290    let nullability = iter.parse_next::<Nullability>();
291    let parameters = match iter.try_pop(Rule::parameters) {
292        Some(p) => parse_parameters(extensions, p)?,
293        None => Vec::new(),
294    };
295    iter.done();
296
297    let anchor = get_and_validate_anchor(extensions, ExtensionKind::Type, anchor, &name, span)?;
298
299    Ok(Type {
300        kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
301            type_reference: anchor,
302            nullability: nullability.into(),
303            type_parameters: parameters,
304            type_variation_reference: 0,
305        })),
306    })
307}
308
309impl ScopedParsePair for Type {
310    fn rule() -> Rule {
311        Rule::r#type
312    }
313
314    fn message() -> &'static str {
315        "Type"
316    }
317
318    fn parse_pair(
319        extensions: &SimpleExtensions,
320        pair: Pair<Rule>,
321    ) -> Result<Self, MessageParseError> {
322        assert_eq!(pair.as_rule(), Rule::r#type);
323        let inner = unwrap_single_pair(pair);
324        match inner.as_rule() {
325            Rule::simple_type => Ok(parse_simple_type(inner)),
326            Rule::compound_type => parse_compound_type(extensions, inner),
327            Rule::user_defined_type => parse_user_defined_type(extensions, inner),
328            _ => unreachable!(
329                "Grammar guarantees type can only be simple_type, compound_type, or user_defined_type, got: {:?}",
330                inner.as_rule()
331            ),
332        }
333    }
334}
335
336#[cfg(test)]
337mod tests {
338    use pest::Parser;
339    use substrait::proto::r#type::{I64, Kind, Nullability};
340
341    use super::*;
342    use crate::parser::ExpressionParser;
343
344    #[test]
345    fn test_parse_simple_type() {
346        let mut pairs = ExpressionParser::parse(Rule::simple_type, "i64").unwrap();
347        let pair = pairs.next().unwrap();
348        assert_eq!(pairs.next(), None);
349        let t = parse_simple_type(pair);
350        assert_eq!(
351            t,
352            Type {
353                kind: Some(Kind::I64(I64 {
354                    nullability: Nullability::Required as i32,
355                    type_variation_reference: 0,
356                })),
357            }
358        );
359
360        let mut pairs = ExpressionParser::parse(Rule::simple_type, "string?").unwrap();
361        let pair = pairs.next().unwrap();
362        assert_eq!(pairs.next(), None);
363        let t = parse_simple_type(pair);
364        assert_eq!(
365            t,
366            Type {
367                kind: Some(Kind::String(proto::r#type::String {
368                    nullability: Nullability::Nullable as i32,
369                    type_variation_reference: 0,
370                })),
371            }
372        );
373    }
374
375    #[test]
376    fn test_parse_type() {
377        let extensions = SimpleExtensions::default();
378        let mut pairs = ExpressionParser::parse(Rule::r#type, "i64").unwrap();
379        let pair = pairs.next().unwrap();
380        assert_eq!(pairs.next(), None);
381        let t = Type::parse_pair(&extensions, pair).unwrap();
382        assert_eq!(
383            t,
384            Type {
385                kind: Some(Kind::I64(I64 {
386                    nullability: Nullability::Required as i32,
387                    type_variation_reference: 0,
388                }))
389            }
390        );
391    }
392
393    #[test]
394    fn test_parse_list_type() {
395        let extensions = SimpleExtensions::default();
396        let mut pairs = ExpressionParser::parse(Rule::list_type, "list<i64>").unwrap();
397        let pair = pairs.next().unwrap();
398        assert_eq!(pairs.next(), None);
399        let t = parse_list_type(&extensions, pair).unwrap();
400        assert_eq!(
401            t,
402            Type {
403                kind: Some(Kind::List(Box::new(proto::r#type::List {
404                    nullability: Nullability::Required as i32,
405                    r#type: Some(Box::new(Type {
406                        kind: Some(Kind::I64(I64 {
407                            nullability: Nullability::Required as i32,
408                            type_variation_reference: 0,
409                        }))
410                    })),
411                    type_variation_reference: 0,
412                })))
413            }
414        );
415    }
416
417    #[test]
418    fn test_parse_parameters() {
419        let extensions = SimpleExtensions::default();
420        let mut pairs = ExpressionParser::parse(Rule::parameters, "<i64?,string>").unwrap();
421        let pair = pairs.next().unwrap();
422        assert_eq!(pairs.next(), None);
423        let t = parse_parameters(&extensions, pair).unwrap();
424        assert_eq!(
425            t,
426            vec![
427                Parameter {
428                    parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
429                        kind: Some(Kind::I64(proto::r#type::I64 {
430                            nullability: Nullability::Nullable as i32,
431                            type_variation_reference: 0,
432                        })),
433                    })),
434                },
435                Parameter {
436                    parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
437                        kind: Some(Kind::String(proto::r#type::String {
438                            nullability: Nullability::Required as i32,
439                            type_variation_reference: 0,
440                        })),
441                    })),
442                },
443            ]
444        );
445    }
446
447    #[test]
448    fn test_udts() {
449        let mut extensions = SimpleExtensions::default();
450        extensions
451            .add_extension_urn("some_source".to_string(), 4)
452            .unwrap();
453        extensions
454            .add_extension(ExtensionKind::Type, 4, 42, "udt".to_string())
455            .unwrap();
456        let mut pairs = ExpressionParser::parse(Rule::user_defined_type, "udt#42<i64?>").unwrap();
457        let pair = pairs.next().unwrap();
458        assert_eq!(pairs.next(), None);
459
460        let t = parse_user_defined_type(&extensions, pair).unwrap();
461        assert_eq!(
462            t,
463            Type {
464                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
465                    type_reference: 42,
466                    type_variation_reference: 0,
467                    nullability: Nullability::Required as i32,
468                    type_parameters: vec![Parameter {
469                        parameter: Some(proto::r#type::parameter::Parameter::DataType(Type {
470                            kind: Some(Kind::I64(proto::r#type::I64 {
471                                nullability: Nullability::Nullable as i32,
472                                type_variation_reference: 0,
473                            })),
474                        })),
475                    }],
476                }))
477            }
478        );
479    }
480
481    #[test]
482    fn test_udts_with_u_prefix() {
483        // Extensions declared with "u!" prefix (e.g. "# 7 @ 4: u!json") normalize to
484        // bare name at storage time. Both "u!json" and "json" in the plan resolve to the
485        // same anchor.
486        let mut extensions = SimpleExtensions::default();
487        extensions
488            .add_extension_urn("some_source".to_string(), 4)
489            .unwrap();
490        extensions
491            .add_extension(ExtensionKind::Type, 4, 7, "u!json".to_string())
492            .unwrap();
493
494        // Plan uses u! prefix with explicit anchor
495        let pair = ExpressionParser::parse(Rule::user_defined_type, "u!json#7")
496            .unwrap()
497            .next()
498            .unwrap();
499        let t = parse_user_defined_type(&extensions, pair).unwrap();
500        assert_eq!(
501            t,
502            Type {
503                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
504                    type_reference: 7,
505                    type_variation_reference: 0,
506                    nullability: Nullability::Required as i32,
507                    type_parameters: vec![],
508                }))
509            }
510        );
511
512        // Plan uses bare name — also resolves to the same anchor
513        let pair = ExpressionParser::parse(Rule::user_defined_type, "json?")
514            .unwrap()
515            .next()
516            .unwrap();
517        let t = parse_user_defined_type(&extensions, pair).unwrap();
518        assert_eq!(
519            t,
520            Type {
521                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
522                    type_reference: 7,
523                    type_variation_reference: 0,
524                    nullability: Nullability::Nullable as i32,
525                    type_parameters: vec![],
526                }))
527            }
528        );
529
530        // Plan uses u! prefix with no anchor — anchor-free lookup must strip u! and find the type
531        let pair = ExpressionParser::parse(Rule::user_defined_type, "u!json")
532            .unwrap()
533            .next()
534            .unwrap();
535        let t = parse_user_defined_type(&extensions, pair).unwrap();
536        assert_eq!(
537            t,
538            Type {
539                kind: Some(Kind::UserDefined(proto::r#type::UserDefined {
540                    type_reference: 7,
541                    type_variation_reference: 0,
542                    nullability: Nullability::Required as i32,
543                    type_parameters: vec![],
544                }))
545            }
546        );
547    }
548
549    #[test]
550    fn test_udt_u_prefix_matches_plain_registration() {
551        // Extension registered without "u!" and plan uses "u!json" — must succeed.
552        let mut extensions = SimpleExtensions::default();
553        extensions.add_extension_urn("src".to_string(), 1).unwrap();
554        extensions
555            .add_extension(ExtensionKind::Type, 1, 3, "json".to_string())
556            .unwrap();
557
558        let pair = ExpressionParser::parse(Rule::user_defined_type, "u!json")
559            .unwrap()
560            .next()
561            .unwrap();
562        let t = parse_user_defined_type(&extensions, pair).unwrap();
563        assert_eq!(
564            t.kind,
565            Some(Kind::UserDefined(proto::r#type::UserDefined {
566                type_reference: 3,
567                type_variation_reference: 0,
568                nullability: Nullability::Required as i32,
569                type_parameters: vec![],
570            }))
571        );
572    }
573
574    #[test]
575    fn test_udt_plain_matches_u_prefix_registration() {
576        // Extension registered with "u!" prefix and plan uses bare "json" — must succeed.
577        let mut extensions = SimpleExtensions::default();
578        extensions.add_extension_urn("src".to_string(), 1).unwrap();
579        extensions
580            .add_extension(ExtensionKind::Type, 1, 3, "u!json".to_string())
581            .unwrap();
582
583        let pair = ExpressionParser::parse(Rule::user_defined_type, "json")
584            .unwrap()
585            .next()
586            .unwrap();
587        let t = parse_user_defined_type(&extensions, pair).unwrap();
588        assert_eq!(
589            t.kind,
590            Some(Kind::UserDefined(proto::r#type::UserDefined {
591                type_reference: 3,
592                type_variation_reference: 0,
593                nullability: Nullability::Required as i32,
594                type_parameters: vec![],
595            }))
596        );
597    }
598}
substrait_explain/parser/types.rs

substrait_explain/parser/
types.rs
