1use std::fmt;
8
9use pest::iterators::Pair;
10use substrait::proto::extensions::AdvancedExtension;
11use substrait::proto::{
12 AggregateRel, CrossRel, FetchRel, FilterRel, JoinRel, Plan, PlanRel, ProjectRel, ReadRel, Rel,
13 RelRoot, SortRel, Version, plan_rel,
14};
15
16use crate::extensions::any::Any;
17use crate::extensions::{AddendumKind, ExtensionRegistry, SimpleExtensions, simple};
18use crate::parser::chunks::ChunkCursor;
19use crate::parser::common::{MessageParseError, ParsePair, ScopedParsePair};
20use crate::parser::errors::{ParseContext, ParseError, ParseResult};
21use crate::parser::expressions::Name;
22use crate::parser::extensions::{
23 AddendumInvocation, ExtensionInvocation, ExtensionParseError, ExtensionParser,
24};
25use crate::parser::relations::{
26 ExtensionReadRel, RelationParsingContext, VirtualReadRel, parse_set_relation_pair,
27};
28use crate::parser::{ErrorKind, ExpressionParser, RelationParsePair, Rule, unwrap_single_pair};
29
30pub const PLAN_HEADER: &str = "=== Plan";
31pub const VERSION_HEADER: &str = "=== Version";
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub struct IndentedLine<'a>(pub usize, pub &'a str);
37
38impl<'a> From<&'a str> for IndentedLine<'a> {
39 fn from(line: &'a str) -> Self {
40 let line = line.trim_end();
41 let mut spaces = 0;
42 for c in line.chars() {
43 if c == ' ' {
44 spaces += 1;
45 } else {
46 break;
47 }
48 }
49
50 let indents = spaces / 2;
51
52 let (_, trimmed) = line.split_at(indents * 2);
53
54 IndentedLine(indents, trimmed)
55 }
56}
57
58#[derive(Debug, Clone)]
61pub struct Addendum<'a> {
62 pub pair: Pair<'a, Rule>, pub line_no: i64,
64}
65
66#[derive(Debug, Clone)]
68pub struct RelationNode<'a> {
69 pub pair: Pair<'a, Rule>,
70 pub line_no: i64,
71 pub addenda: Vec<Addendum<'a>>,
72 pub children: Vec<RelationNode<'a>>,
73}
74
75impl<'a> RelationNode<'a> {
76 pub fn context(&self) -> ParseContext {
77 ParseContext {
78 line_no: self.line_no,
79 line: self.pair.as_str().to_string(),
80 }
81 }
82}
83
84#[derive(Debug, Clone)]
90pub enum LineNode<'a> {
91 Relation(RelationNode<'a>),
92 Addendum(Addendum<'a>),
93}
94
95impl<'a> LineNode<'a> {
96 pub fn parse(line: &'a str, line_no: i64) -> Result<Self, ParseError> {
97 let mut pairs: pest::iterators::Pairs<'a, Rule> =
98 <ExpressionParser as pest::Parser<Rule>>::parse(Rule::planNode, line).map_err(|e| {
99 ParseError::Plan(
100 ParseContext {
101 line_no,
102 line: line.to_string(),
103 },
104 MessageParseError::new("planNode", ErrorKind::InvalidValue, Box::new(e)),
105 )
106 })?;
107
108 let outer = pairs.next().unwrap();
109 assert!(pairs.next().is_none()); let inner = unwrap_single_pair(outer);
111
112 Ok(match inner.as_rule() {
113 Rule::addendum => LineNode::Addendum(Addendum {
114 pair: inner,
115 line_no,
116 }),
117 _ => LineNode::Relation(RelationNode {
118 pair: inner,
119 line_no,
120 addenda: Vec::new(),
121 children: Vec::new(),
122 }),
123 })
124 }
125
126 pub fn parse_root(line: &'a str, line_no: i64) -> Result<Self, ParseError> {
128 let mut pairs: pest::iterators::Pairs<'a, Rule> = <ExpressionParser as pest::Parser<
129 Rule,
130 >>::parse(
131 Rule::top_level_relation, line
132 )
133 .map_err(|e| {
134 ParseError::Plan(
135 ParseContext::new(line_no, line.to_string()),
136 MessageParseError::new("top_level_relation", ErrorKind::Syntax, Box::new(e)),
137 )
138 })?;
139
140 let outer = pairs.next().unwrap();
141 assert!(pairs.next().is_none());
142
143 let inner = unwrap_single_pair(outer);
146 let pair = if inner.as_rule() == Rule::planNode {
147 unwrap_single_pair(inner)
148 } else {
149 inner };
151
152 if pair.as_rule() == Rule::addendum {
155 return Ok(LineNode::Addendum(Addendum { pair, line_no }));
156 }
157
158 Ok(LineNode::Relation(RelationNode {
159 pair,
160 line_no,
161 addenda: Vec::new(),
162 children: Vec::new(),
163 }))
164 }
165}
166
167#[derive(Copy, Clone, Debug, PartialEq, Eq)]
168pub enum State {
169 Initial,
171 Version,
174 Extensions,
176 Plan,
178}
179
180impl State {
181 fn section_rank(&self) -> u8 {
186 match self {
187 State::Initial => 0,
188 State::Version => 1,
189 State::Extensions => 2,
190 State::Plan => 3,
191 }
192 }
193}
194
195impl fmt::Display for State {
196 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
197 write!(f, "{self:?}")
198 }
199}
200
201#[derive(Debug, Clone, Default)]
203pub struct TreeBuilder<'a> {
204 current: Option<RelationNode<'a>>,
207 completed: Vec<RelationNode<'a>>,
209}
210
211impl<'a> TreeBuilder<'a> {
212 pub fn get_at_depth(&mut self, depth: usize) -> Option<&mut RelationNode<'a>> {
214 let mut node = self.current.as_mut()?;
215 for _ in 0..depth {
216 node = node.children.last_mut()?;
217 }
218 Some(node)
219 }
220
221 pub fn add_line(&mut self, depth: usize, node: LineNode<'a>) -> Result<(), ParseError> {
222 match node {
223 LineNode::Relation(rel_node) => {
224 if depth == 0 {
225 if let Some(prev) = self.current.take() {
226 self.completed.push(prev);
227 }
228 self.current = Some(rel_node);
229 return Ok(());
230 }
231
232 let parent = match self.get_at_depth(depth - 1) {
233 None => {
234 return Err(ParseError::Plan(
235 rel_node.context(),
236 MessageParseError::invalid(
237 "relation",
238 rel_node.pair.as_span(),
239 format!("No parent found for depth {depth}"),
240 ),
241 ));
242 }
243 Some(parent) => parent,
244 };
245
246 parent.children.push(rel_node);
247 }
248 LineNode::Addendum(addendum) => {
249 let context =
250 ParseContext::new(addendum.line_no, addendum.pair.as_str().to_string());
251 if depth == 0 {
252 return Err(ParseError::ValidationError(
253 context,
254 "addenda (+ Enh: / + Opt: / + Ext:) cannot appear at the top level"
255 .to_string(),
256 ));
257 }
258
259 let parent = match self.get_at_depth(depth - 1) {
260 None => {
261 return Err(ParseError::ValidationError(
262 context,
263 format!("no parent found for addendum at depth {depth}"),
264 ));
265 }
266 Some(parent) => parent,
267 };
268
269 if !parent.children.is_empty() {
270 return Err(ParseError::ValidationError(
271 context,
272 "addenda (+ Enh: / + Opt: / + Ext:) must appear before child relations, \
273 not after"
274 .to_string(),
275 ));
276 }
277
278 parent.addenda.push(addendum);
279 }
280 }
281 Ok(())
282 }
283
284 pub fn finish(&mut self) -> Vec<RelationNode<'a>> {
289 if let Some(node) = self.current.take() {
290 self.completed.push(node);
291 }
292 std::mem::take(&mut self.completed)
293 }
294}
295
296struct ParsedChild {
299 rel: Rel,
300 field_count: usize,
301}
302
303struct RelationContext<'a> {
307 pair: Pair<'a, Rule>,
308 line_no: i64,
309 children: Vec<ParsedChild>,
310 addenda: Addenda<'a>,
311}
312
313impl<'a> RelationContext<'a> {
314 fn input_field_count(&self) -> usize {
316 self.children.iter().map(|c| c.field_count).sum()
317 }
318
319 fn child_field_counts(&self) -> impl Iterator<Item = usize> + '_ {
322 self.children.iter().map(|c| c.field_count)
323 }
324
325 fn as_parse_context(&self) -> ParseContext {
327 ParseContext::new(self.line_no, self.pair.as_str().to_string())
328 }
329}
330
331fn into_rels(children: Vec<ParsedChild>) -> Vec<Rel> {
333 children.into_iter().map(|c| c.rel).collect()
334}
335
336#[derive(Debug, Clone)]
338struct ParsedAddendum<'a> {
339 line_no: i64,
340 line: &'a str,
341 invocation: AddendumInvocation,
342}
343
344impl<'a> ParsedAddendum<'a> {
345 fn parse(extensions: &SimpleExtensions, addendum: Addendum<'a>) -> Result<Self, ParseError> {
346 let line_no = addendum.line_no;
347 let line = addendum.pair.as_str();
348 let invocation = AddendumInvocation::parse_pair(extensions, addendum.pair)
349 .map_err(|e| ParseError::Plan(ParseContext::new(line_no, line.to_string()), e))?;
350 Ok(Self {
351 line_no,
352 line,
353 invocation,
354 })
355 }
356
357 fn context(&self) -> ParseContext {
358 ParseContext::new(self.line_no, self.line.to_string())
359 }
360
361 fn relation_context<'b>(
362 &'b self,
363 registry: &'b ExtensionRegistry,
364 ) -> RelationParsingContext<'b> {
365 RelationParsingContext {
366 registry,
367 line_no: self.line_no,
368 line: self.line,
369 }
370 }
371 fn resolve_detail(&self, registry: &ExtensionRegistry) -> Result<Any, ParseError> {
372 self.relation_context(registry).resolve_addendum_detail(
373 self.invocation.kind,
374 &self.invocation.name,
375 &self.invocation.args,
376 )
377 }
378}
379
380#[derive(Debug, Clone, Default)]
382struct Addenda<'a> {
383 items: Vec<ParsedAddendum<'a>>,
384}
385
386impl<'a> Addenda<'a> {
387 fn parse(
388 extensions: &SimpleExtensions,
389 addenda: Vec<Addendum<'a>>,
390 ) -> Result<Self, ParseError> {
391 let items = addenda
392 .into_iter()
393 .map(|addendum| ParsedAddendum::parse(extensions, addendum))
394 .collect::<Result<Vec<_>, ParseError>>()?;
395 Ok(Self { items })
396 }
397
398 fn first(&self) -> Option<&ParsedAddendum<'a>> {
399 self.items.first()
400 }
401
402 fn reject_all(&self, message: &'static str) -> Result<(), ParseError> {
403 if let Some(addendum) = self.first() {
404 return Err(ParseError::ValidationError(
405 addendum.context(),
406 message.to_string(),
407 ));
408 }
409 Ok(())
410 }
411
412 fn into_standard_advanced_extension(
413 self,
414 registry: &ExtensionRegistry,
415 ) -> Result<Option<AdvancedExtension>, ParseError> {
416 let mut enhancement = None;
417 let mut optimizations = Vec::new();
418
419 for addendum in self.items {
420 match addendum.invocation.kind {
421 AddendumKind::Enhancement => {
422 if enhancement.is_some() {
423 return Err(ParseError::ValidationError(
424 addendum.context(),
425 "at most one enhancement per relation is allowed".to_string(),
426 ));
427 }
428 enhancement = Some(addendum.resolve_detail(registry)?.into());
429 }
430 AddendumKind::Optimization => {
431 optimizations.push(addendum.resolve_detail(registry)?.into());
432 }
433 AddendumKind::ExtensionTable => {
434 return Err(ParseError::ValidationError(
435 addendum.context(),
436 "+ Ext addenda can only be used with Read:Extension".to_string(),
437 ));
438 }
439 }
440 }
441
442 if enhancement.is_none() && optimizations.is_empty() {
443 return Ok(None);
444 }
445
446 Ok(Some(AdvancedExtension {
447 enhancement,
448 optimization: optimizations,
449 }))
450 }
451
452 fn into_extension_read_parts(
453 self,
454 registry: &ExtensionRegistry,
455 relation_context: ParseContext,
456 ) -> Result<(Any, Option<AdvancedExtension>), ParseError> {
457 let mut extension_table = None;
458 let mut advanced_addenda = Vec::new();
459
460 for addendum in self.items {
461 match addendum.invocation.kind {
462 AddendumKind::ExtensionTable => {
463 if extension_table.is_some() {
464 return Err(ParseError::ValidationError(
465 addendum.context(),
466 "Read:Extension allows exactly one + Ext addendum".to_string(),
467 ));
468 }
469 extension_table = Some(addendum);
470 }
471 AddendumKind::Enhancement | AddendumKind::Optimization => {
472 advanced_addenda.push(addendum);
473 }
474 }
475 }
476
477 let extension_table = extension_table.ok_or_else(|| {
478 ParseError::ValidationError(
479 relation_context,
480 "Read:Extension requires exactly one + Ext addendum".to_string(),
481 )
482 })?;
483
484 let detail = extension_table.resolve_detail(registry)?;
485 let advanced_extension = Addenda {
486 items: advanced_addenda,
487 }
488 .into_standard_advanced_extension(registry)?;
489
490 Ok((detail, advanced_extension))
491 }
492}
493
494#[derive(Debug, Clone, Default)]
496pub struct RelationParser<'a> {
497 tree: TreeBuilder<'a>,
498}
499
500impl<'a> RelationParser<'a> {
501 fn parse_relation(
506 &self,
507 extensions: &SimpleExtensions,
508 registry: &ExtensionRegistry,
509 ctx: RelationContext,
510 ) -> Result<(Rel, usize), ParseError> {
511 match ctx.pair.as_rule() {
512 Rule::extension_read_relation => {
513 self.parse_extension_read_relation(extensions, registry, ctx)
514 }
515 Rule::virtual_read_relation => {
516 self.parse_rel::<VirtualReadRel>(extensions, registry, ctx)
517 }
518 Rule::read_relation => self.parse_rel::<ReadRel>(extensions, registry, ctx),
519 Rule::filter_relation => self.parse_rel::<FilterRel>(extensions, registry, ctx),
520 Rule::project_relation => self.parse_rel::<ProjectRel>(extensions, registry, ctx),
521 Rule::aggregate_relation => self.parse_rel::<AggregateRel>(extensions, registry, ctx),
522 Rule::sort_relation => self.parse_rel::<SortRel>(extensions, registry, ctx),
523 Rule::fetch_relation => self.parse_rel::<FetchRel>(extensions, registry, ctx),
524 Rule::join_relation => self.parse_rel::<JoinRel>(extensions, registry, ctx),
525 Rule::set_relation => self.parse_set_relation(registry, ctx),
526 Rule::cross_relation => self.parse_rel::<CrossRel>(extensions, registry, ctx),
527 Rule::extension_relation => self.parse_extension_relation(extensions, registry, ctx),
528 _ => unreachable!("unhandled relation rule: {:?}", ctx.pair.as_rule()),
529 }
530 }
531
532 fn parse_rel<T: RelationParsePair>(
537 &self,
538 extensions: &SimpleExtensions,
539 registry: &ExtensionRegistry,
540 ctx: RelationContext,
541 ) -> Result<(Rel, usize), ParseError> {
542 let RelationContext {
543 pair,
544 line_no,
545 children,
546 addenda,
547 } = ctx;
548 assert_eq!(pair.as_rule(), T::rule());
549 let line = pair.as_str();
550 let advanced_extension = addenda.into_standard_advanced_extension(registry)?;
551 let input_field_count: usize = children.iter().map(|c| c.field_count).sum();
552 let input_children = into_rels(children);
553
554 match T::parse_pair_with_context(extensions, pair, input_children, input_field_count) {
558 Ok((parsed, count)) => Ok((parsed.into_rel(advanced_extension), count)),
559 Err(e) => Err(ParseError::Plan(
560 ParseContext::new(line_no, line.to_string()),
561 e,
562 )),
563 }
564 }
565
566 fn parse_set_relation(
570 &self,
571 registry: &ExtensionRegistry,
572 ctx: RelationContext,
573 ) -> Result<(Rel, usize), ParseError> {
574 assert_eq!(ctx.pair.as_rule(), Rule::set_relation);
575 let context = ctx.as_parse_context();
576 let child_field_counts: Vec<usize> = ctx.child_field_counts().collect();
577 let RelationContext {
578 pair,
579 children,
580 addenda,
581 ..
582 } = ctx;
583 let advanced_extension = addenda.into_standard_advanced_extension(registry)?;
584 let input_children = into_rels(children);
585
586 parse_set_relation_pair(
587 pair,
588 input_children,
589 &child_field_counts,
590 advanced_extension,
591 )
592 .map_err(|e| ParseError::Plan(context, e))
593 }
594
595 fn parse_extension_relation(
599 &self,
600 extensions: &SimpleExtensions,
601 registry: &ExtensionRegistry,
602 ctx: RelationContext,
603 ) -> Result<(Rel, usize), ParseError> {
604 assert_eq!(ctx.pair.as_rule(), Rule::extension_relation);
605 let line_no = ctx.line_no;
606 let line = ctx.pair.as_str().to_string();
607 let pair_span = ctx.pair.as_span();
608
609 ctx.addenda
610 .reject_all("extension relations do not support addenda (+ Enh / + Opt / + Ext)")?;
611
612 let ExtensionInvocation {
613 relation_kind,
614 name,
615 args: extension_args,
616 } = ExtensionInvocation::parse_pair(extensions, ctx.pair.clone())
617 .map_err(|e| ParseError::Plan(ParseContext::new(line_no, line.clone()), e))?;
618
619 let child_count = ctx.children.len();
620 relation_kind
621 .validate_child_count(child_count)
622 .map_err(|e| {
623 ParseError::Plan(
624 ParseContext::new(line_no, line.to_string()),
625 MessageParseError::invalid("extension_relation", pair_span, e),
626 )
627 })?;
628
629 let context = RelationParsingContext {
630 registry,
631 line_no,
632 line: &line,
633 };
634
635 let detail = context.resolve_extension_detail(&name, &extension_args)?;
636 let output_column_count = extension_args.output_columns.len();
637
638 let rel = relation_kind.create_rel(detail, into_rels(ctx.children));
639
640 Ok((rel, output_column_count))
641 }
642
643 fn parse_extension_read_relation(
646 &self,
647 extensions: &SimpleExtensions,
648 registry: &ExtensionRegistry,
649 ctx: RelationContext,
650 ) -> Result<(Rel, usize), ParseError> {
651 assert_eq!(ctx.pair.as_rule(), Rule::extension_read_relation);
652 let context = ParseContext::new(ctx.line_no, ctx.pair.as_str().to_string());
653 let input_field_count = ctx.input_field_count();
654 let (detail, advanced_extension) = ctx
655 .addenda
656 .into_extension_read_parts(registry, context.clone())?;
657
658 ExtensionReadRel::parse_pair_with_detail(
659 extensions,
660 ctx.pair,
661 into_rels(ctx.children),
662 input_field_count,
663 detail,
664 advanced_extension,
665 )
666 .map_err(|e| ParseError::Plan(context, e))
667 }
668
669 fn build_rel(
674 &self,
675 extensions: &SimpleExtensions,
676 registry: &ExtensionRegistry,
677 node: RelationNode,
678 ) -> Result<(Rel, usize), ParseError> {
679 let mut children: Vec<ParsedChild> = Vec::new();
680 for child in node.children {
681 let (rel, field_count) = self.build_rel(extensions, registry, child)?;
682 children.push(ParsedChild { rel, field_count });
683 }
684
685 let addenda = Addenda::parse(extensions, node.addenda)?;
686
687 self.parse_relation(
688 extensions,
689 registry,
690 RelationContext {
691 pair: node.pair,
692 line_no: node.line_no,
693 children,
694 addenda,
695 },
696 )
697 }
698
699 fn build_plan_rel(
701 &self,
702 extensions: &SimpleExtensions,
703 registry: &ExtensionRegistry,
704 node: RelationNode,
705 ) -> Result<PlanRel, ParseError> {
706 if node.pair.as_rule() != Rule::root_relation {
708 let (rel, _) = self.build_rel(extensions, registry, node)?;
709 return Ok(PlanRel {
710 rel_type: Some(plan_rel::RelType::Rel(rel)),
711 });
712 }
713
714 if !node.addenda.is_empty() {
716 let first = &node.addenda[0];
717 let context = ParseContext::new(first.line_no, first.pair.as_str().to_string());
718 return Err(ParseError::ValidationError(
719 context,
720 "addenda (+ Enh: / + Opt: / + Ext:) are not supported on Root relations"
721 .to_string(),
722 ));
723 }
724
725 let context = node.context();
727 let span = node.pair.as_span();
728
729 let column_names_pair = unwrap_single_pair(node.pair);
731 assert_eq!(column_names_pair.as_rule(), Rule::root_name_list);
732
733 let names: Vec<String> = column_names_pair
734 .into_inner()
735 .map(|name_pair| {
736 assert_eq!(name_pair.as_rule(), Rule::name);
737 Name::parse_pair(name_pair).0
738 })
739 .collect();
740
741 let mut children = node.children;
742 let child = match children.len() {
743 1 => {
744 let (rel, _) = self.build_rel(extensions, registry, children.pop().unwrap())?;
745 rel
746 }
747 n => {
748 return Err(ParseError::Plan(
749 context,
750 MessageParseError::invalid(
751 "root_relation",
752 span,
753 format!("Root relation must have exactly one child, found {n}"),
754 ),
755 ));
756 }
757 };
758
759 Ok(PlanRel {
760 rel_type: Some(plan_rel::RelType::Root(RelRoot {
761 names,
762 input: Some(child),
763 })),
764 })
765 }
766
767 fn build(
769 mut self,
770 extensions: &SimpleExtensions,
771 registry: &ExtensionRegistry,
772 ) -> Result<Vec<PlanRel>, ParseError> {
773 let nodes = self.tree.finish();
774 nodes
775 .into_iter()
776 .map(|n| self.build_plan_rel(extensions, registry, n))
777 .collect::<Result<Vec<PlanRel>, ParseError>>()
778 }
779}
780
781#[derive(Debug)]
891pub struct Parser<'a> {
892 line_no: i64,
893 state: State,
894 cursor: Option<ChunkCursor<'a>>,
898 version: Option<Version>,
900 extension_parser: ExtensionParser,
901 extension_registry: ExtensionRegistry,
902 relation_parser: RelationParser<'a>,
903}
904impl<'a> Default for Parser<'a> {
905 fn default() -> Self {
906 Self::new()
907 }
908}
909
910impl<'a> Parser<'a> {
911 pub fn parse(input: &str) -> ParseResult {
940 Self::new().parse_plan(input)
941 }
942
943 pub fn new() -> Self {
945 Self {
946 line_no: 1,
947 state: State::Initial,
948 cursor: None,
949 version: None,
950 extension_parser: ExtensionParser::default(),
951 extension_registry: ExtensionRegistry::new(),
952 relation_parser: RelationParser::default(),
953 }
954 }
955
956 pub fn with_extension_registry(mut self, registry: ExtensionRegistry) -> Self {
958 self.extension_registry = registry;
959 self
960 }
961
962 pub fn parse_plan(mut self, input: &'a str) -> ParseResult {
964 self.cursor = ChunkCursor::new(input, 1);
965 while self.cursor.is_some() {
966 let (chunk, line_no) = self.next_chunk();
967
968 if chunk.trim().is_empty() {
969 continue;
970 }
971
972 self.line_no = line_no;
973 self.parse_line(chunk)?;
982 }
983
984 let plan = self.build_plan()?;
985 Ok(plan)
986 }
987
988 fn next_chunk(&mut self) -> (&'a str, i64) {
994 let mut c = self
995 .cursor
996 .take()
997 .expect("next_chunk called with no cursor");
998
999 let first = c
1001 .peek_line()
1002 .expect("a non-exhausted cursor always yields a line");
1003 c.merge(first);
1004
1005 if matches!(self.state, State::Plan) {
1006 let base = IndentedLine::from(first.as_str()).0;
1007 while let Some(line) = c.peek_line() {
1008 let IndentedLine(depth, body) = IndentedLine::from(line.as_str());
1009 if depth == base + 1 && body.starts_with("- ") {
1010 c.merge(line);
1011 } else {
1012 break;
1013 }
1014 }
1015 }
1016
1017 let line_no = c.start_line_no();
1018 let (chunk, rest) = c.next();
1019 self.cursor = rest;
1020 (chunk.trim_end_matches(['\r', '\n']), line_no)
1021 }
1022
1023 fn parse_line(&mut self, line: &'a str) -> Result<(), ParseError> {
1025 let indented_line = IndentedLine::from(line);
1026 let line_no = self.line_no;
1027 let ctx = || ParseContext {
1028 line_no,
1029 line: line.to_string(),
1030 };
1031
1032 if let IndentedLine(0, l) = indented_line
1038 && l.starts_with("===")
1039 {
1040 let next = self.parse_section_header(l)?;
1041 if next.section_rank() <= self.state.section_rank() {
1042 return Err(ParseError::ValidationError(
1043 ctx(),
1044 format!(
1045 "section {next} is out of order after {current}; sections must appear \
1046 in the order Version, Extensions, Plan",
1047 current = self.state,
1048 ),
1049 ));
1050 }
1051 self.state = next;
1052 return Ok(());
1053 }
1054
1055 match self.state {
1056 State::Initial => match indented_line {
1057 IndentedLine(0, l) if l.trim().is_empty() => Ok(()),
1058 IndentedLine(n, l) => Err(ParseError::Initial(
1059 ParseContext::new(n as i64, l.to_string()),
1060 MessageParseError::invalid(
1061 "initial",
1062 pest::Span::new(l, 0, l.len()).expect("Invalid span?!"),
1063 format!("Unknown initial line: {l:?}"),
1064 ),
1065 )),
1066 },
1067 State::Version => self.parse_version(indented_line),
1068 State::Extensions => self
1069 .parse_extensions(indented_line)
1070 .map_err(|e| ParseError::Extension(ctx(), e)),
1071 State::Plan => {
1072 let IndentedLine(depth, line_str) = indented_line;
1073
1074 let node = if depth == 0 {
1076 LineNode::parse_root(line_str, line_no)?
1077 } else {
1078 LineNode::parse(line_str, line_no)?
1079 };
1080
1081 self.relation_parser.tree.add_line(depth, node)
1082 }
1083 }
1084 }
1085
1086 fn parse_section_header(&mut self, line: &str) -> Result<State, ParseError> {
1091 if line == simple::EXTENSIONS_HEADER {
1092 return Ok(State::Extensions);
1093 }
1094 if line == PLAN_HEADER {
1095 return Ok(State::Plan);
1096 }
1097 if let Some(rest) = line.strip_prefix(VERSION_HEADER)
1098 && (rest.is_empty() || rest.starts_with(' '))
1099 {
1100 self.version = Some(self.parse_version_header(line)?);
1101 return Ok(State::Version);
1102 }
1103
1104 Err(ParseError::Initial(
1105 ParseContext::new(self.line_no, line.to_string()),
1106 MessageParseError::invalid(
1107 "initial",
1108 pest::Span::new(line, 0, line.len()).expect("Invalid span?!"),
1109 format!("Unknown section header: {line:?}"),
1110 ),
1111 ))
1112 }
1113
1114 fn parse_version_header(&self, line: &str) -> Result<Version, ParseError> {
1117 let ctx = || ParseContext::new(self.line_no, line.to_string());
1118 let rest = line
1119 .strip_prefix(VERSION_HEADER)
1120 .expect("version header prefix checked by caller")
1121 .trim();
1122
1123 let mut numbers = rest.split('.');
1124 let mut next_number = |field: &str| -> Result<u32, ParseError> {
1125 let part = numbers.next().filter(|p| !p.is_empty()).ok_or_else(|| {
1126 ParseError::ValidationError(
1127 ctx(),
1128 format!("version header is missing the {field} number, expected '{VERSION_HEADER} <major>.<minor>.<patch>'"),
1129 )
1130 })?;
1131 part.parse::<u32>().map_err(|e| {
1132 ParseError::ValidationError(
1133 ctx(),
1134 format!("invalid {field} version number {part:?}: {e}"),
1135 )
1136 })
1137 };
1138
1139 let major_number = next_number("major")?;
1140 let minor_number = next_number("minor")?;
1141 let patch_number = next_number("patch")?;
1142 if numbers.next().is_some() {
1143 return Err(ParseError::ValidationError(
1144 ctx(),
1145 format!(
1146 "version {rest:?} has too many components, expected '<major>.<minor>.<patch>'"
1147 ),
1148 ));
1149 }
1150
1151 Ok(Version {
1152 major_number,
1153 minor_number,
1154 patch_number,
1155 ..Default::default()
1156 })
1157 }
1158
1159 fn parse_version(&mut self, line: IndentedLine) -> Result<(), ParseError> {
1163 match line {
1164 IndentedLine(0, l) if l.trim().is_empty() => Ok(()),
1165 IndentedLine(1, l) => self.parse_version_detail(l),
1166 IndentedLine(_, l) => Err(ParseError::ValidationError(
1167 ParseContext::new(self.line_no, l.to_string()),
1168 format!(
1169 "unexpected line in version section: {l:?}; expected an indented 'producer:' \
1170 or 'git_hash:' line, or a section header"
1171 ),
1172 )),
1173 }
1174 }
1175
1176 fn parse_version_detail(&mut self, line: &str) -> Result<(), ParseError> {
1179 let ctx = || ParseContext::new(self.line_no, line.to_string());
1180 let (key, value) = line.split_once(':').ok_or_else(|| {
1181 ParseError::ValidationError(
1182 ctx(),
1183 format!("expected 'producer:' or 'git_hash:' in version section, found {line:?}"),
1184 )
1185 })?;
1186 let value = value.trim().to_string();
1187 let version = self
1188 .version
1189 .as_mut()
1190 .expect("version is set on entry to the version section");
1191 match key.trim() {
1192 "producer" => version.producer = value,
1193 "git_hash" => version.git_hash = value,
1194 other => {
1195 return Err(ParseError::ValidationError(
1196 ctx(),
1197 format!("unknown version field {other:?}, expected 'producer' or 'git_hash'"),
1198 ));
1199 }
1200 }
1201 Ok(())
1202 }
1203
1204 fn parse_extensions(&mut self, line: IndentedLine<'_>) -> Result<(), ExtensionParseError> {
1208 self.extension_parser.parse_line(line)
1209 }
1210
1211 fn build_plan(self) -> Result<Plan, ParseError> {
1213 let Parser {
1214 version,
1215 relation_parser,
1216 extension_parser,
1217 extension_registry,
1218 ..
1219 } = self;
1220
1221 let extensions = extension_parser.extensions();
1222
1223 let root_relations = relation_parser.build(extensions, &extension_registry)?;
1225
1226 Ok(Plan {
1228 version,
1229 extension_urns: extensions.to_extension_urns(),
1230 extensions: extensions.to_extension_declarations(),
1231 relations: root_relations,
1232 ..Default::default()
1233 })
1234 }
1235}
1236
1237#[cfg(test)]
1238mod tests {
1239 use substrait::proto::extensions::simple_extension_declaration::MappingType;
1240 use substrait::proto::rel::RelType;
1241
1242 use super::*;
1243 use crate::extensions::simple::ExtensionKind;
1244 use crate::parser::extensions::ExpectedExtensionLine;
1245
1246 #[test]
1247 fn test_parse_basic_block() {
1248 let mut expected_extensions = SimpleExtensions::new();
1249 expected_extensions
1250 .add_extension_urn("/urn/common".to_string(), 1)
1251 .unwrap();
1252 expected_extensions
1253 .add_extension_urn("/urn/specific_funcs".to_string(), 2)
1254 .unwrap();
1255 expected_extensions
1256 .add_extension(ExtensionKind::Function, 1, 10, "func_a".to_string())
1257 .unwrap();
1258 expected_extensions
1259 .add_extension(ExtensionKind::Function, 2, 11, "func_b_special".to_string())
1260 .unwrap();
1261 expected_extensions
1262 .add_extension(ExtensionKind::Type, 1, 20, "SomeType".to_string())
1263 .unwrap();
1264 expected_extensions
1265 .add_extension(ExtensionKind::TypeVariation, 2, 30, "VarX".to_string())
1266 .unwrap();
1267
1268 let mut parser = ExtensionParser::default();
1269 let input_block = r#"
1270URNs:
1271 @ 1: /urn/common
1272 @ 2: /urn/specific_funcs
1273Functions:
1274 # 10 @ 1: func_a
1275 # 11 @ 2: func_b_special
1276Types:
1277 # 20 @ 1: SomeType
1278Type Variations:
1279 # 30 @ 2: VarX
1280"#;
1281
1282 for line_str in input_block.trim().lines() {
1283 parser
1284 .parse_line(IndentedLine::from(line_str))
1285 .unwrap_or_else(|e| panic!("Failed to parse line \'{line_str}\': {e:?}"));
1286 }
1287
1288 assert_eq!(*parser.extensions(), expected_extensions);
1289
1290 let extensions_str = parser.extensions().to_string(" ");
1291 let expected_str = format!(
1294 "{}\n{}",
1295 simple::EXTENSIONS_HEADER,
1296 input_block.trim_start()
1297 );
1298 assert_eq!(extensions_str.trim(), expected_str.trim());
1299 assert_eq!(
1302 parser.state(),
1303 ExpectedExtensionLine::ExtensionDeclarations(ExtensionKind::TypeVariation)
1304 );
1305
1306 parser.parse_line(IndentedLine(0, "")).unwrap();
1308 assert_eq!(parser.state(), ExpectedExtensionLine::Extensions);
1309 }
1310
1311 #[test]
1313 fn test_parse_complete_extension_block() {
1314 let mut parser = ExtensionParser::default();
1315 let input_block = r#"
1316URNs:
1317 @ 1: /urn/common
1318 @ 2: /urn/specific_funcs
1319 @ 3: /urn/types_lib
1320 @ 4: /urn/variations_lib
1321Functions:
1322 # 10 @ 1: func_a
1323 # 11 @ 2: func_b_special
1324 # 12 @ 1: func_c_common
1325Types:
1326 # 20 @ 1: CommonType
1327 # 21 @ 3: LibraryType
1328 # 22 @ 1: AnotherCommonType
1329Type Variations:
1330 # 30 @ 4: VarX
1331 # 31 @ 4: VarY
1332"#;
1333
1334 for line_str in input_block.trim().lines() {
1335 parser
1336 .parse_line(IndentedLine::from(line_str))
1337 .unwrap_or_else(|e| panic!("Failed to parse line \'{line_str}\': {e:?}"));
1338 }
1339
1340 let extensions_str = parser.extensions().to_string(" ");
1341 let expected_str = format!(
1344 "{}\n{}",
1345 simple::EXTENSIONS_HEADER,
1346 input_block.trim_start()
1347 );
1348 assert_eq!(extensions_str.trim(), expected_str.trim());
1349 }
1350
1351 #[test]
1352 fn test_parse_relation_tree() {
1353 let plan = r#"=== Plan
1355Project[$0, $1, 42, 84]
1356 Filter[$2 => $0, $1]
1357 Read[my.table => a:i32, b:string?, c:boolean]
1358"#;
1359 let mut parser = Parser::default();
1360 for line in plan.lines() {
1361 parser.parse_line(line).unwrap();
1362 }
1363
1364 let plan = parser.build_plan().unwrap();
1366
1367 let root_rel = &plan.relations[0].rel_type;
1368 let first_rel = match root_rel {
1369 Some(plan_rel::RelType::Rel(rel)) => rel,
1370 _ => panic!("Expected Rel type, got {root_rel:?}"),
1371 };
1372 let project = match &first_rel.rel_type {
1374 Some(RelType::Project(p)) => p,
1375 other => panic!("Expected Project at root, got {other:?}"),
1376 };
1377
1378 assert!(project.input.is_some());
1380 let filter_input = project.input.as_ref().unwrap();
1381
1382 match &filter_input.rel_type {
1384 Some(RelType::Filter(_)) => {
1385 match &filter_input.rel_type {
1386 Some(RelType::Filter(filter)) => {
1387 assert!(filter.input.is_some());
1388 let read_input = filter.input.as_ref().unwrap();
1389
1390 match &read_input.rel_type {
1392 Some(RelType::Read(_)) => {}
1393 other => panic!("Expected Read relation, got {other:?}"),
1394 }
1395 }
1396 other => panic!("Expected Filter relation, got {other:?}"),
1397 }
1398 }
1399 other => panic!("Expected Filter relation, got {other:?}"),
1400 }
1401 }
1402
1403 #[test]
1404 fn test_parse_root_relation() {
1405 let plan = r#"=== Plan
1407Root[result]
1408 Project[$0, $1]
1409 Read[my.table => a:i32, b:string?]
1410"#;
1411 let mut parser = Parser::default();
1412 for line in plan.lines() {
1413 parser.parse_line(line).unwrap();
1414 }
1415
1416 let plan = parser.build_plan().unwrap();
1417
1418 assert_eq!(plan.relations.len(), 1);
1420
1421 let root_rel = &plan.relations[0].rel_type;
1422 let rel_root = match root_rel {
1423 Some(plan_rel::RelType::Root(rel_root)) => rel_root,
1424 other => panic!("Expected Root type, got {other:?}"),
1425 };
1426
1427 assert_eq!(rel_root.names, vec!["result"]);
1429
1430 let project_input = match &rel_root.input {
1432 Some(rel) => rel,
1433 None => panic!("Root should have an input"),
1434 };
1435
1436 let project = match &project_input.rel_type {
1437 Some(RelType::Project(p)) => p,
1438 other => panic!("Expected Project as root input, got {other:?}"),
1439 };
1440
1441 let read_input = match &project.input {
1443 Some(rel) => rel,
1444 None => panic!("Project should have an input"),
1445 };
1446
1447 match &read_input.rel_type {
1448 Some(RelType::Read(_)) => {}
1449 other => panic!("Expected Read relation, got {other:?}"),
1450 }
1451 }
1452
1453 #[test]
1454 fn test_parse_root_relation_no_names() {
1455 let plan = r#"=== Plan
1457Root[]
1458 Project[$0, $1]
1459 Read[my.table => a:i32, b:string?]
1460"#;
1461 let mut parser = Parser::default();
1462 for line in plan.lines() {
1463 parser.parse_line(line).unwrap();
1464 }
1465
1466 let plan = parser.build_plan().unwrap();
1467
1468 let root_rel = &plan.relations[0].rel_type;
1469 let rel_root = match root_rel {
1470 Some(plan_rel::RelType::Root(rel_root)) => rel_root,
1471 other => panic!("Expected Root type, got {other:?}"),
1472 };
1473
1474 assert_eq!(rel_root.names, Vec::<String>::new());
1476 }
1477
1478 #[test]
1479 fn test_parse_full_plan() {
1480 let input = r#"
1482=== Extensions
1483URNs:
1484 @ 1: /urn/common
1485 @ 2: /urn/specific_funcs
1486Functions:
1487 # 10 @ 1: func_a
1488 # 11 @ 2: func_b_special
1489Types:
1490 # 20 @ 1: SomeType
1491Type Variations:
1492 # 30 @ 2: VarX
1493
1494=== Plan
1495Project[$0, $1, 42, 84]
1496 Filter[$2 => $0, $1]
1497 Read[my.table => a:i32, b:string?, c:boolean]
1498"#;
1499
1500 let plan = Parser::parse(input).unwrap();
1501
1502 assert_eq!(plan.extension_urns.len(), 2);
1504 assert_eq!(plan.extensions.len(), 4);
1505 assert_eq!(plan.relations.len(), 1);
1506
1507 let urn1 = &plan.extension_urns[0];
1509 assert_eq!(urn1.extension_urn_anchor, 1);
1510 assert_eq!(urn1.urn, "/urn/common");
1511
1512 let urn2 = &plan.extension_urns[1];
1513 assert_eq!(urn2.extension_urn_anchor, 2);
1514 assert_eq!(urn2.urn, "/urn/specific_funcs");
1515
1516 let func1 = &plan.extensions[0];
1518 match &func1.mapping_type {
1519 Some(MappingType::ExtensionFunction(f)) => {
1520 assert_eq!(f.function_anchor, 10);
1521 assert_eq!(f.extension_urn_reference, 1);
1522 assert_eq!(f.name, "func_a");
1523 }
1524 other => panic!("Expected ExtensionFunction, got {other:?}"),
1525 }
1526
1527 let func2 = &plan.extensions[1];
1528 match &func2.mapping_type {
1529 Some(MappingType::ExtensionFunction(f)) => {
1530 assert_eq!(f.function_anchor, 11);
1531 assert_eq!(f.extension_urn_reference, 2);
1532 assert_eq!(f.name, "func_b_special");
1533 }
1534 other => panic!("Expected ExtensionFunction, got {other:?}"),
1535 }
1536
1537 let type1 = &plan.extensions[2];
1538 match &type1.mapping_type {
1539 Some(MappingType::ExtensionType(t)) => {
1540 assert_eq!(t.type_anchor, 20);
1541 assert_eq!(t.extension_urn_reference, 1);
1542 assert_eq!(t.name, "SomeType");
1543 }
1544 other => panic!("Expected ExtensionType, got {other:?}"),
1545 }
1546
1547 let var1 = &plan.extensions[3];
1548 match &var1.mapping_type {
1549 Some(MappingType::ExtensionTypeVariation(v)) => {
1550 assert_eq!(v.type_variation_anchor, 30);
1551 assert_eq!(v.extension_urn_reference, 2);
1552 assert_eq!(v.name, "VarX");
1553 }
1554 other => panic!("Expected ExtensionTypeVariation, got {other:?}"),
1555 }
1556
1557 let root_rel = &plan.relations[0];
1559 match &root_rel.rel_type {
1560 Some(plan_rel::RelType::Rel(rel)) => {
1561 match &rel.rel_type {
1562 Some(RelType::Project(project)) => {
1563 assert_eq!(project.expressions.len(), 2); assert!(project.input.is_some()); let filter_input = project.input.as_ref().unwrap();
1569 match &filter_input.rel_type {
1570 Some(RelType::Filter(filter)) => {
1571 assert!(filter.input.is_some()); let read_input = filter.input.as_ref().unwrap();
1575 match &read_input.rel_type {
1576 Some(RelType::Read(read)) => {
1577 let schema = read.base_schema.as_ref().unwrap();
1579 assert_eq!(schema.names.len(), 3);
1580 assert_eq!(schema.names[0], "a");
1581 assert_eq!(schema.names[1], "b");
1582 assert_eq!(schema.names[2], "c");
1583
1584 let struct_ = schema.r#struct.as_ref().unwrap();
1585 assert_eq!(struct_.types.len(), 3);
1586 }
1587 other => panic!("Expected Read relation, got {other:?}"),
1588 }
1589 }
1590 other => panic!("Expected Filter relation, got {other:?}"),
1591 }
1592 }
1593 other => panic!("Expected Project relation, got {other:?}"),
1594 }
1595 }
1596 other => panic!("Expected Rel type, got {other:?}"),
1597 }
1598 }
1599
1600 #[test]
1603 fn test_section_headers_out_of_order_rejected() {
1604 let input = r#"
1605=== Plan
1606Root[result]
1607 Read[t => a:i32]
1608
1609=== Extensions
1610URNs:
1611 @ 1: /urn/common
1612"#;
1613 let err = Parser::parse(input).unwrap_err();
1614 let msg = err.to_string();
1615 assert!(
1616 msg.contains("out of order"),
1617 "expected an out-of-order error, got: {msg}"
1618 );
1619 }
1620
1621 #[test]
1622 fn test_repeated_section_header_rejected() {
1623 let input = r#"
1624=== Extensions
1625URNs:
1626 @ 1: /urn/common
1627
1628=== Extensions
1629URNs:
1630 @ 2: /urn/other
1631
1632=== Plan
1633Root[result]
1634 Read[t => a:i32]
1635"#;
1636 let err = Parser::parse(input).unwrap_err();
1637 let msg = err.to_string();
1638 assert!(
1639 msg.contains("out of order"),
1640 "expected an out-of-order error, got: {msg}"
1641 );
1642 }
1643
1644 #[test]
1645 fn test_version_after_plan_rejected() {
1646 let input = r#"
1647=== Plan
1648Root[result]
1649 Read[t => a:i32]
1650
1651=== Version 1.0.0
1652"#;
1653 let err = Parser::parse(input).unwrap_err();
1654 let msg = err.to_string();
1655 assert!(
1656 msg.contains("out of order"),
1657 "expected an out-of-order error, got: {msg}"
1658 );
1659 }
1660
1661 #[test]
1664 fn test_all_sections_in_order_accepted() {
1665 let input = r#"
1666=== Version 1.2.3
1667=== Extensions
1668URNs:
1669 @ 1: /urn/common
1670
1671=== Plan
1672Root[result]
1673 Read[t => a:i32]
1674"#;
1675 let plan = Parser::parse(input).unwrap();
1676 assert_eq!(plan.relations.len(), 1);
1677 assert!(plan.version.is_some());
1678 }
1679}