substrait_explain/parser/
relations.rs

1use std::collections::HashMap;
2
3use pest::iterators::Pair;
4use prost::Message;
5use substrait::proto::aggregate_rel::Grouping;
6use substrait::proto::expression::literal::LiteralType;
7use substrait::proto::expression::{Literal, RexType, nested};
8use substrait::proto::extensions::AdvancedExtension;
9use substrait::proto::fetch_rel::{CountMode, OffsetMode};
10use substrait::proto::rel::RelType;
11use substrait::proto::rel_common::{Direct, Emit, EmitKind};
12use substrait::proto::sort_field::{SortDirection, SortKind};
13use substrait::proto::{
14    AggregateRel, CrossRel, Expression, FetchRel, FilterRel, JoinRel, NamedStruct, ProjectRel,
15    ReadRel, Rel, RelCommon, SetRel, SortField, SortRel, Type, aggregate_rel, join_rel, read_rel,
16    set_rel, r#type,
17};
18
19use super::{MessageParseError, ParsePair, Rule, RuleIter, ScopedParsePair, unwrap_single_pair};
20use crate::extensions::any::Any;
21use crate::extensions::registry::ExtensionError;
22use crate::extensions::{AddendumKind, ExtensionArgs, ExtensionRegistry, SimpleExtensions};
23use crate::parser::errors::{ParseContext, ParseError};
24use crate::parser::expressions::{FieldIndex, Name};
25
26/// Parsing context for relations that includes extensions, registry, and optional warning collection
27pub struct RelationParsingContext<'a> {
28    pub registry: &'a ExtensionRegistry,
29    pub line_no: i64,
30    pub line: &'a str,
31}
32
33impl<'a> RelationParsingContext<'a> {
34    /// Resolve extension detail using registry. Any failure is treated as a hard parse error.
35    pub fn resolve_extension_detail(
36        &self,
37        extension_name: &str,
38        extension_args: &ExtensionArgs,
39    ) -> Result<Option<Any>, ParseError> {
40        let detail = self
41            .registry
42            .parse_extension(extension_name, extension_args);
43
44        match detail {
45            Ok(any) => Ok(Some(any)),
46            Err(ExtensionError::NotFound { .. }) => Err(ParseError::UnregisteredExtension {
47                name: extension_name.to_string(),
48                context: ParseContext::new(self.line_no, self.line.to_string()),
49            }),
50            Err(err) => Err(ParseError::ExtensionDetail(
51                ParseContext::new(self.line_no, self.line.to_string()),
52                err,
53            )),
54        }
55    }
56
57    /// Resolve an addendum detail using the registry entry for its kind.
58    /// Any failure is treated as a hard parse error.
59    pub(crate) fn resolve_addendum_detail(
60        &self,
61        kind: AddendumKind,
62        name: &str,
63        args: &ExtensionArgs,
64    ) -> Result<Any, ParseError> {
65        let result = match kind {
66            AddendumKind::Enhancement => self.registry.parse_enhancement(name, args),
67            AddendumKind::Optimization => self.registry.parse_optimization(name, args),
68            AddendumKind::ExtensionTable => self.registry.parse_extension_table(name, args),
69        };
70        result.map_err(|err| match err {
71            ExtensionError::NotFound { .. } => ParseError::UnregisteredExtension {
72                name: name.to_string(),
73                context: ParseContext::new(self.line_no, self.line.to_string()),
74            },
75            err => ParseError::ExtensionDetail(
76                ParseContext::new(self.line_no, self.line.to_string()),
77                err,
78            ),
79        })
80    }
81}
82
83/// A trait for parsing relations with full context for tree building.
84pub trait RelationParsePair: Sized {
85    fn rule() -> Rule;
86    fn message() -> &'static str;
87
88    /// Parse the grammar pair into this relation type and its output field
89    /// count.
90    ///
91    /// Returns `(Self, usize)` where `usize` is the output field count —
92    /// computed during parsing when `input_field_count` is available.
93    fn parse_pair_with_context(
94        extensions: &SimpleExtensions,
95        pair: Pair<Rule>,
96        input_children: Vec<Rel>,
97        input_field_count: usize,
98    ) -> Result<(Self, usize), MessageParseError>;
99
100    /// Consume this parsed relation, apply the advanced extension, and produce
101    /// the final `Rel`.
102    fn into_rel(self, adv_ext: Option<AdvancedExtension>) -> Rel;
103}
104
105pub struct TableName(Vec<String>);
106
107impl ParsePair for TableName {
108    fn rule() -> Rule {
109        Rule::table_name
110    }
111
112    fn message() -> &'static str {
113        "TableName"
114    }
115
116    fn parse_pair(pair: Pair<Rule>) -> Self {
117        assert_eq!(pair.as_rule(), Self::rule());
118        let pairs = pair.into_inner();
119        let mut names = Vec::with_capacity(pairs.len());
120        let mut iter = RuleIter::from(pairs);
121        while let Some(name) = iter.parse_if_next::<Name>() {
122            names.push(name.0);
123        }
124        iter.done();
125        Self(names)
126    }
127}
128
129#[derive(Debug, Clone)]
130pub struct Column {
131    pub name: String,
132    pub typ: Type,
133}
134
135impl ScopedParsePair for Column {
136    fn rule() -> Rule {
137        Rule::named_column
138    }
139
140    fn message() -> &'static str {
141        "Column"
142    }
143
144    fn parse_pair(
145        extensions: &SimpleExtensions,
146        pair: Pair<Rule>,
147    ) -> Result<Self, MessageParseError> {
148        assert_eq!(pair.as_rule(), Self::rule());
149        let mut iter = RuleIter::from(pair.into_inner());
150        let name = iter.parse_next::<Name>().0;
151        let typ = iter.parse_next_scoped(extensions)?;
152        iter.done();
153        Ok(Self { name, typ })
154    }
155}
156
157pub(crate) struct NamedColumnList(pub(crate) Vec<Column>);
158
159impl ScopedParsePair for NamedColumnList {
160    fn rule() -> Rule {
161        Rule::named_column_list
162    }
163
164    fn message() -> &'static str {
165        "NamedColumnList"
166    }
167
168    fn parse_pair(
169        extensions: &SimpleExtensions,
170        pair: Pair<Rule>,
171    ) -> Result<Self, MessageParseError> {
172        assert_eq!(pair.as_rule(), Self::rule());
173        let mut columns = Vec::new();
174        for col in pair.into_inner() {
175            columns.push(Column::parse_pair(extensions, col)?);
176        }
177        Ok(Self(columns))
178    }
179}
180
181/// This is a utility function for extracting a single child from the list of
182/// children, to be used in the RelationParsePair trait. The RelationParsePair
183/// trait passes a Vec of children, because some relations have multiple
184/// children - but most accept exactly one child.
185pub(crate) fn expect_one_child(
186    message: &'static str,
187    pair: &Pair<Rule>,
188    mut input_children: Vec<Rel>,
189) -> Result<Box<Rel>, MessageParseError> {
190    match input_children.len() {
191        0 => Err(MessageParseError::invalid(
192            message,
193            pair.as_span(),
194            format!("{message} missing child"),
195        )),
196        1 => Ok(Box::new(input_children.pop().unwrap())),
197        n => Err(MessageParseError::invalid(
198            message,
199            pair.as_span(),
200            format!("{message} should have 1 input child, got {n}"),
201        )),
202    }
203}
204
205/// Parse a reference list Pair and return an EmitKind::Emit.
206/// Parse a reference list into field indices for emit mapping.
207fn parse_output_mapping(pair: Pair<Rule>) -> Vec<i32> {
208    assert_eq!(pair.as_rule(), Rule::reference_list);
209    pair.into_inner()
210        .map(|p| FieldIndex::parse_pair(p).0)
211        .collect()
212}
213
214/// Build an emit: `Direct` if the mapping is the identity `[0, 1, ..., N-1]`
215/// (where N = `direct_output_count`), otherwise `Emit` with the explicit mapping.
216fn make_emit(output_mapping: Vec<i32>, direct_output_count: usize) -> EmitKind {
217    let is_identity = output_mapping.len() == direct_output_count
218        && output_mapping
219            .iter()
220            .enumerate()
221            .all(|(i, &v)| v == i as i32);
222    if is_identity {
223        EmitKind::Direct(Direct {})
224    } else {
225        EmitKind::Emit(Emit { output_mapping })
226    }
227}
228
229/// Parse a reference list into an emit and output field count.
230fn parse_emit(reference_list: Pair<Rule>, direct_output_count: usize) -> (EmitKind, usize) {
231    let output_mapping = parse_output_mapping(reference_list);
232    let output_count = output_mapping.len();
233    let emit = make_emit(output_mapping, direct_output_count);
234    (emit, output_count)
235}
236
237/// Extracts named arguments from pest pairs with duplicate detection and completeness checking.
238///
239/// Usage: `extractor.pop("limit", Rule::fetch_value).0.pop("offset", Rule::fetch_value).0.done()`
240///
241/// The fluent API ensures all arguments are processed exactly once and none are forgotten.
242pub struct ParsedNamedArgs<'a> {
243    map: HashMap<&'a str, Pair<'a, Rule>>,
244}
245
246impl<'a> ParsedNamedArgs<'a> {
247    pub fn new(
248        pairs: pest::iterators::Pairs<'a, Rule>,
249        rule: Rule,
250    ) -> Result<Self, MessageParseError> {
251        let mut map = HashMap::new();
252        for pair in pairs {
253            assert_eq!(pair.as_rule(), rule);
254            let mut inner = pair.clone().into_inner();
255            let name_pair = inner.next().unwrap();
256            let value_pair = inner.next().unwrap();
257            assert_eq!(inner.next(), None);
258            let name = name_pair.as_str();
259            if map.contains_key(name) {
260                return Err(MessageParseError::invalid(
261                    "NamedArg",
262                    name_pair.as_span(),
263                    format!("Duplicate argument: {name}"),
264                ));
265            }
266            map.insert(name, value_pair);
267        }
268        Ok(Self { map })
269    }
270
271    // Returns the pair if it exists and matches the rule, otherwise None.
272    // Asserts that the rule must match the rule of the pair (and therefore
273    // panics in non-release-mode if not)
274    pub fn pop(mut self, name: &str, rule: Rule) -> (Self, Option<Pair<'a, Rule>>) {
275        let pair = self.map.remove(name).inspect(|pair| {
276            assert_eq!(pair.as_rule(), rule, "Rule mismatch for argument {name}");
277        });
278        (self, pair)
279    }
280
281    // Returns an error if there are any unused arguments.
282    pub fn done(self) -> Result<(), MessageParseError> {
283        if let Some((name, pair)) = self.map.iter().next() {
284            return Err(MessageParseError::invalid(
285                "NamedArgExtractor",
286                // No span available for all unused args; use default.
287                pair.as_span(),
288                format!("Unknown argument: {name}"),
289            ));
290        }
291        Ok(())
292    }
293}
294
295impl RelationParsePair for ReadRel {
296    fn rule() -> Rule {
297        Rule::read_relation
298    }
299
300    fn message() -> &'static str {
301        "ReadRel"
302    }
303
304    fn parse_pair_with_context(
305        extensions: &SimpleExtensions,
306        pair: Pair<Rule>,
307        input_children: Vec<Rel>,
308        input_field_count: usize,
309    ) -> Result<(Self, usize), MessageParseError> {
310        assert_eq!(pair.as_rule(), Self::rule());
311        // ReadRel is a leaf node - it should have no input children and 0 input fields
312        if !input_children.is_empty() {
313            return Err(MessageParseError::invalid(
314                Self::message(),
315                pair.as_span(),
316                "ReadRel should have no input children",
317            ));
318        }
319        if input_field_count != 0 {
320            return Err(MessageParseError::invalid(
321                "ReadRel",
322                pair.as_span(),
323                "ReadRel should have 0 input fields",
324            ));
325        }
326
327        let mut iter = RuleIter::from(pair.into_inner());
328        let table = iter.parse_next::<TableName>().0;
329        let columns = iter.parse_next_scoped::<NamedColumnList>(extensions)?.0;
330        iter.done();
331
332        let output_count = columns.len();
333        Ok((
334            ReadRel {
335                base_schema: Some(build_named_struct(columns)),
336                read_type: Some(read_rel::ReadType::NamedTable(read_rel::NamedTable {
337                    names: table,
338                    advanced_extension: None,
339                })),
340                ..Default::default()
341            },
342            output_count,
343        ))
344    }
345
346    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
347        self.advanced_extension = adv_ext;
348        Rel {
349            rel_type: Some(RelType::Read(Box::new(self))),
350        }
351    }
352}
353
354/// Parsed `Read:Virtual[rows => columns]` relation. Needs a newtype because the
355/// proto type is `ReadRel` (same as `NamedTable`), but the grammar and handling
356/// are different.
357pub(crate) struct VirtualReadRel(ReadRel);
358
359impl RelationParsePair for VirtualReadRel {
360    fn rule() -> Rule {
361        Rule::virtual_read_relation
362    }
363
364    fn message() -> &'static str {
365        "VirtualReadRel"
366    }
367
368    fn parse_pair_with_context(
369        extensions: &SimpleExtensions,
370        pair: Pair<Rule>,
371        input_children: Vec<Rel>,
372        _input_field_count: usize,
373    ) -> Result<(Self, usize), MessageParseError> {
374        assert_eq!(pair.as_rule(), Self::rule());
375        if !input_children.is_empty() {
376            return Err(MessageParseError::invalid(
377                Self::message(),
378                pair.as_span(),
379                "Read:Virtual should have no input children",
380            ));
381        }
382
383        let mut iter = RuleIter::from(pair.into_inner());
384        let args_pair = iter.pop(Rule::virtual_read_args);
385        let columns_pair = iter.pop(Rule::named_column_list);
386        iter.done();
387
388        let rows = parse_virtual_read_args(extensions, args_pair)?;
389        let columns = NamedColumnList::parse_pair(extensions, columns_pair)?.0;
390
391        // TODO: Validate that each row has the same number of expressions as
392        // columns. Currently no parser-side warning mechanism exists, and while
393        // this is an invalid plan, it is constructible as Substrait. Consider
394        // adding once a warning collection path is available.
395        let output_count = columns.len();
396        Ok((
397            VirtualReadRel(ReadRel {
398                base_schema: Some(build_named_struct(columns)),
399                read_type: Some(read_rel::ReadType::VirtualTable(read_rel::VirtualTable {
400                    expressions: rows,
401                    ..Default::default()
402                })),
403                ..Default::default()
404            }),
405            output_count,
406        ))
407    }
408
409    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
410        self.0.advanced_extension = adv_ext;
411        Rel {
412            rel_type: Some(RelType::Read(Box::new(self.0))),
413        }
414    }
415}
416
417/// Parsed `Read:Extension[columns]` relation. Needs a newtype because the
418/// proto type is `ReadRel` (same as `NamedTable` and `VirtualTable`), but the
419/// read detail is supplied by a required `+ Ext:` addendum.
420pub(crate) struct ExtensionReadRel(ReadRel);
421
422impl ExtensionReadRel {
423    pub(crate) fn parse_pair_with_detail(
424        extensions: &SimpleExtensions,
425        pair: Pair<Rule>,
426        input_children: Vec<Rel>,
427        input_field_count: usize,
428        detail: Any,
429        advanced_extension: Option<AdvancedExtension>,
430    ) -> Result<(Rel, usize), MessageParseError> {
431        assert_eq!(pair.as_rule(), Rule::extension_read_relation);
432        if !input_children.is_empty() {
433            return Err(MessageParseError::invalid(
434                "ExtensionReadRel",
435                pair.as_span(),
436                "Read:Extension should have no input children",
437            ));
438        }
439        if input_field_count != 0 {
440            return Err(MessageParseError::invalid(
441                "ExtensionReadRel",
442                pair.as_span(),
443                "Read:Extension should have 0 input fields",
444            ));
445        }
446
447        let mut iter = RuleIter::from(pair.into_inner());
448        let columns = iter.parse_next_scoped::<NamedColumnList>(extensions)?.0;
449        iter.done();
450
451        let output_count = columns.len();
452        let rel = ExtensionReadRel(ReadRel {
453            base_schema: Some(build_named_struct(columns)),
454            read_type: Some(read_rel::ReadType::ExtensionTable(
455                read_rel::ExtensionTable {
456                    detail: Some(detail.into()),
457                },
458            )),
459            advanced_extension,
460            ..Default::default()
461        })
462        .into_rel();
463
464        Ok((rel, output_count))
465    }
466
467    fn into_rel(self) -> Rel {
468        Rel {
469            rel_type: Some(RelType::Read(Box::new(self.0))),
470        }
471    }
472}
473
474/// Build a `NamedStruct` from parsed columns.
475pub(crate) fn build_named_struct(columns: Vec<Column>) -> NamedStruct {
476    let (names, types): (Vec<_>, Vec<_>) = columns.into_iter().map(|c| (c.name, c.typ)).unzip();
477    NamedStruct {
478        names,
479        r#struct: Some(r#type::Struct {
480            types,
481            type_variation_reference: 0,
482            nullability: r#type::Nullability::Required as i32,
483        }),
484    }
485}
486
487/// `Read:Virtual` positional args: either `empty` or a list of row tuples.
488fn parse_virtual_read_args(
489    extensions: &SimpleExtensions,
490    pair: Pair<Rule>,
491) -> Result<Vec<nested::Struct>, MessageParseError> {
492    assert_eq!(pair.as_rule(), Rule::virtual_read_args);
493    let inner = unwrap_single_pair(pair);
494    match inner.as_rule() {
495        Rule::empty => Ok(vec![]),
496        Rule::virtual_row_list => inner
497            .into_inner()
498            .map(|row| parse_virtual_row(extensions, row))
499            .collect(),
500        _ => unreachable!(
501            "Unexpected rule in virtual_read_args: {:?}",
502            inner.as_rule()
503        ),
504    }
505}
506
507/// Parse a single `virtual_row` (`(expr, expr, ...)` or `()`) into a `nested::Struct`.
508fn parse_virtual_row(
509    extensions: &SimpleExtensions,
510    pair: Pair<Rule>,
511) -> Result<nested::Struct, MessageParseError> {
512    assert_eq!(pair.as_rule(), Rule::virtual_row);
513    let fields = match pair.into_inner().next() {
514        Some(expression_list) => {
515            assert_eq!(expression_list.as_rule(), Rule::expression_list);
516            parse_expression_list(extensions, expression_list)?
517        }
518        // Empty virtual row. An unusual but valid case.
519        None => vec![],
520    };
521    Ok(nested::Struct { fields })
522}
523
524impl RelationParsePair for FilterRel {
525    fn rule() -> Rule {
526        Rule::filter_relation
527    }
528
529    fn message() -> &'static str {
530        "FilterRel"
531    }
532
533    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
534        self.advanced_extension = adv_ext;
535        Rel {
536            rel_type: Some(RelType::Filter(Box::new(self))),
537        }
538    }
539
540    fn parse_pair_with_context(
541        extensions: &SimpleExtensions,
542        pair: Pair<Rule>,
543        input_children: Vec<Rel>,
544        input_field_count: usize,
545    ) -> Result<(Self, usize), MessageParseError> {
546        assert_eq!(pair.as_rule(), Self::rule());
547        let input = expect_one_child(Self::message(), &pair, input_children)?;
548        let mut iter = RuleIter::from(pair.into_inner());
549        let condition = iter.parse_next_scoped::<Expression>(extensions)?;
550        // references (which become the emit)
551        let references_pair = iter.pop(Rule::reference_list);
552        iter.done();
553
554        let (emit, output_count) = parse_emit(references_pair, input_field_count);
555        let common = RelCommon {
556            emit_kind: Some(emit),
557            ..Default::default()
558        };
559
560        Ok((
561            FilterRel {
562                input: Some(input),
563                condition: Some(Box::new(condition)),
564                common: Some(common),
565                advanced_extension: None,
566            },
567            output_count,
568        ))
569    }
570}
571
572impl RelationParsePair for ProjectRel {
573    fn rule() -> Rule {
574        Rule::project_relation
575    }
576
577    fn message() -> &'static str {
578        "ProjectRel"
579    }
580
581    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
582        self.advanced_extension = adv_ext;
583        Rel {
584            rel_type: Some(RelType::Project(Box::new(self))),
585        }
586    }
587
588    fn parse_pair_with_context(
589        extensions: &SimpleExtensions,
590        pair: Pair<Rule>,
591        input_children: Vec<Rel>,
592        input_field_count: usize,
593    ) -> Result<(Self, usize), MessageParseError> {
594        assert_eq!(pair.as_rule(), Self::rule());
595        let input = expect_one_child(Self::message(), &pair, input_children)?;
596
597        let arguments_pair = unwrap_single_pair(pair);
598
599        let mut expressions = Vec::new();
600        let mut output_mapping = Vec::new();
601
602        for arg in arguments_pair.into_inner() {
603            let inner_arg = unwrap_single_pair(arg);
604            match inner_arg.as_rule() {
605                Rule::reference => {
606                    let field_index = FieldIndex::parse_pair(inner_arg);
607                    output_mapping.push(field_index.0);
608                }
609                Rule::expression => {
610                    let expr = Expression::parse_pair(extensions, inner_arg)?;
611                    expressions.push(expr);
612                    // Index into the combined schema: [input fields][computed expressions].
613                    output_mapping.push(input_field_count as i32 + (expressions.len() as i32 - 1));
614                }
615                _ => panic!("Unexpected inner argument rule: {:?}", inner_arg.as_rule()),
616            }
617        }
618
619        let output_count = output_mapping.len();
620        let direct_count = input_field_count + expressions.len();
621        let emit = make_emit(output_mapping, direct_count);
622        let common = RelCommon {
623            emit_kind: Some(emit),
624            ..Default::default()
625        };
626
627        Ok((
628            ProjectRel {
629                input: Some(input),
630                expressions,
631                common: Some(common),
632                advanced_extension: None,
633            },
634            output_count,
635        ))
636    }
637}
638
639impl RelationParsePair for AggregateRel {
640    fn rule() -> Rule {
641        Rule::aggregate_relation
642    }
643
644    fn message() -> &'static str {
645        "AggregateRel"
646    }
647
648    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
649        self.advanced_extension = adv_ext;
650        Rel {
651            rel_type: Some(RelType::Aggregate(Box::new(self))),
652        }
653    }
654
655    fn parse_pair_with_context(
656        extensions: &SimpleExtensions,
657        pair: Pair<Rule>,
658        input_children: Vec<Rel>,
659        // Aggregate defines its own output schema (grouping keys + measures),
660        // so the input field count isn't needed for emit construction.
661        _input_field_count: usize,
662    ) -> Result<(Self, usize), MessageParseError> {
663        assert_eq!(pair.as_rule(), Self::rule());
664        let input = expect_one_child(Self::message(), &pair, input_children)?;
665        let mut iter = RuleIter::from(pair.into_inner());
666        let group_by_pair = iter.pop(Rule::aggregate_group_by);
667        let output_pair = iter.pop(Rule::aggregate_output);
668        iter.done();
669
670        let inner = group_by_pair
671            .into_inner()
672            .next()
673            .expect("aggregate_group_by must have one inner item");
674
675        let grouping_sets = parse_grouping_sets(extensions, inner)?;
676        let (groupings, grouping_expressions) = build_grouping_fields(&grouping_sets);
677
678        let (measures, output_mapping) =
679            parse_aggregate_measures(extensions, output_pair, &grouping_expressions)?;
680
681        let output_count = output_mapping.len();
682        let direct_count = grouping_expressions.len() + measures.len();
683        let emit = make_emit(output_mapping, direct_count);
684        let common = RelCommon {
685            emit_kind: Some(emit),
686            ..Default::default()
687        };
688
689        Ok((
690            AggregateRel {
691                input: Some(input),
692                grouping_expressions,
693                groupings,
694                measures,
695                common: Some(common),
696                advanced_extension: None,
697            },
698            output_count,
699        ))
700    }
701}
702
703/// Parses the output section of an aggregate (everything after `=>`).
704///
705/// For example, in `Aggregate[($0, $1), _ => sum($2), $0, count($2)]`,
706/// this parses `sum($2), $0, count($2)`.
707fn parse_aggregate_measures(
708    extensions: &SimpleExtensions,
709    output_pair: Pair<'_, Rule>,
710    grouping_expressions: &[Expression],
711) -> Result<(Vec<aggregate_rel::Measure>, Vec<i32>), MessageParseError> {
712    assert_eq!(output_pair.as_rule(), Rule::aggregate_output);
713    let mut measures = Vec::new();
714    let mut output_mapping = Vec::new();
715
716    for aggregate_output_item in output_pair.into_inner() {
717        let inner_item = unwrap_single_pair(aggregate_output_item);
718        match inner_item.as_rule() {
719            Rule::reference => {
720                let field_index = FieldIndex::parse_pair(inner_item);
721                output_mapping.push(field_index.0);
722            }
723            Rule::aggregate_measure => {
724                let measure = aggregate_rel::Measure::parse_pair(extensions, inner_item)?;
725                output_mapping.push(grouping_expressions.len() as i32 + measures.len() as i32);
726                measures.push(measure);
727            }
728            _ => panic!(
729                "Unexpected inner output item rule: {:?}",
730                inner_item.as_rule()
731            ),
732        }
733    }
734
735    Ok((measures, output_mapping))
736}
737
738/// Parses the grouping section of an aggregate (everything before `=>`).
739///
740/// For example, in `Aggregate[($0, $1), _ => sum($2), $0, count($2)]`,
741/// this parses `($0, $1), _`.
742///
743/// Each inner Vec is one grouping set; an empty vec represents no grouping (global aggregate).
744///
745/// Grammar: `aggregate_group_by = { grouping_set_list | expression_list }`
746fn parse_grouping_sets(
747    extensions: &SimpleExtensions,
748    inner: Pair<'_, Rule>,
749) -> Result<Vec<Vec<Expression>>, MessageParseError> {
750    assert!(
751        matches!(
752            inner.as_rule(),
753            Rule::expression_list | Rule::grouping_set_list
754        ),
755        "Expected expression_list or grouping_set_list, got {:?}",
756        inner.as_rule()
757    );
758    match inner.as_rule() {
759        Rule::expression_list => Ok(vec![parse_expression_list(extensions, inner)?]),
760        Rule::grouping_set_list => inner
761            .into_inner()
762            .map(|pair| parse_grouping_set(extensions, pair))
763            .collect(),
764        _ => unreachable!(
765            "Unexpected rule in aggregate_group_by: {:?}",
766            inner.as_rule()
767        ),
768    }
769}
770
771/// Parses a single grouping set, e.g. `($0, $1)` or `_`.
772///
773/// Grammar: `grouping_set = { ("(" ~ expression_list ~ ")") | empty }`
774fn parse_grouping_set(
775    extensions: &SimpleExtensions,
776    pair: Pair<'_, Rule>,
777) -> Result<Vec<Expression>, MessageParseError> {
778    assert_eq!(pair.as_rule(), Rule::grouping_set);
779    let inner = pair
780        .into_inner()
781        .next()
782        .expect("grouping_set must have one inner item");
783    match inner.as_rule() {
784        Rule::empty => Ok(vec![]),
785        Rule::expression_list => parse_expression_list(extensions, inner),
786        _ => unreachable!("Unexpected item in grouping_set: {:?}", inner.as_rule()),
787    }
788}
789
790/// Grammar: `expression_list = { expression ~ ("," ~ expression)* }`
791pub(crate) fn parse_expression_list(
792    extensions: &SimpleExtensions,
793    pair: Pair<'_, Rule>,
794) -> Result<Vec<Expression>, MessageParseError> {
795    pair.into_inner()
796        .map(|expr_pair| Expression::parse_pair(extensions, expr_pair))
797        .collect()
798}
799
800/// Deduplicates expressions across all sets and produces the AggregateRel's
801/// protobuf fields: a flat deduplicated expression list and per-set Grouping
802/// messages with index references into that list.
803fn build_grouping_fields(expression_sets: &[Vec<Expression>]) -> (Vec<Grouping>, Vec<Expression>) {
804    let mut expressions: Vec<Expression> = Vec::new();
805    let mut seen: HashMap<Vec<u8>, u32> = HashMap::new();
806
807    let groupings = expression_sets
808        .iter()
809        .map(|set| {
810            let expression_references = set
811                .iter()
812                .map(|exp| {
813                    // TODO: use a better key here than encoding to bytes.
814                    // Ideally, substrait-rs would support `PartialEq` and `Hash`,
815                    // but as there isn't an easy way to do that now, we'll skip.
816                    let key = exp.encode_to_vec();
817                    let next_idx = expressions.len() as u32;
818                    *seen.entry(key).or_insert_with(|| {
819                        expressions.push(exp.clone());
820                        next_idx
821                    })
822                })
823                .collect();
824            Grouping {
825                expression_references,
826                #[allow(deprecated)]
827                grouping_expressions: vec![],
828            }
829        })
830        .collect();
831
832    (groupings, expressions)
833}
834
835impl ScopedParsePair for SortField {
836    fn rule() -> Rule {
837        Rule::sort_field
838    }
839
840    fn message() -> &'static str {
841        "SortField"
842    }
843
844    fn parse_pair(
845        _extensions: &SimpleExtensions,
846        pair: Pair<Rule>,
847    ) -> Result<Self, MessageParseError> {
848        assert_eq!(pair.as_rule(), Self::rule());
849        let mut iter = RuleIter::from(pair.into_inner());
850        let reference_pair = iter.pop(Rule::reference);
851        let field_index = FieldIndex::parse_pair(reference_pair);
852        let direction_pair = iter.pop(Rule::sort_direction);
853        // Strip the '&' prefix from enum syntax (e.g., "&AscNullsFirst" ->
854        // "AscNullsFirst") The grammar includes '&' to distinguish enums from
855        // identifiers, but the enum variant names don't include it
856        let direction = match direction_pair.as_str().trim_start_matches('&') {
857            "AscNullsFirst" => SortDirection::AscNullsFirst,
858            "AscNullsLast" => SortDirection::AscNullsLast,
859            "DescNullsFirst" => SortDirection::DescNullsFirst,
860            "DescNullsLast" => SortDirection::DescNullsLast,
861            other => {
862                return Err(MessageParseError::invalid(
863                    "SortDirection",
864                    direction_pair.as_span(),
865                    format!("Unknown sort direction: {other}"),
866                ));
867            }
868        };
869        iter.done();
870        Ok(SortField {
871            expr: Some(Expression {
872                rex_type: Some(substrait::proto::expression::RexType::Selection(Box::new(
873                    field_index.to_field_reference(),
874                ))),
875            }),
876            // TODO: Add support for SortKind::ComparisonFunctionReference
877            sort_kind: Some(SortKind::Direction(direction as i32)),
878        })
879    }
880}
881
882impl RelationParsePair for SortRel {
883    fn rule() -> Rule {
884        Rule::sort_relation
885    }
886
887    fn message() -> &'static str {
888        "SortRel"
889    }
890
891    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
892        self.advanced_extension = adv_ext;
893        Rel {
894            rel_type: Some(RelType::Sort(Box::new(self))),
895        }
896    }
897
898    fn parse_pair_with_context(
899        extensions: &SimpleExtensions,
900        pair: Pair<Rule>,
901        input_children: Vec<Rel>,
902        input_field_count: usize,
903    ) -> Result<(Self, usize), MessageParseError> {
904        assert_eq!(pair.as_rule(), Self::rule());
905        let input = expect_one_child(Self::message(), &pair, input_children)?;
906        let mut iter = RuleIter::from(pair.into_inner());
907        let sort_field_list_pair = iter.pop(Rule::sort_field_list);
908        let reference_list_pair = iter.pop(Rule::reference_list);
909        let mut sorts = Vec::new();
910        for sort_field_pair in sort_field_list_pair.into_inner() {
911            let sort_field = SortField::parse_pair(extensions, sort_field_pair)?;
912            sorts.push(sort_field);
913        }
914        let (emit, output_count) = parse_emit(reference_list_pair, input_field_count);
915        let common = RelCommon {
916            emit_kind: Some(emit),
917            ..Default::default()
918        };
919        iter.done();
920        Ok((
921            SortRel {
922                input: Some(input),
923                sorts,
924                common: Some(common),
925                advanced_extension: None,
926            },
927            output_count,
928        ))
929    }
930}
931
932impl ScopedParsePair for CountMode {
933    fn rule() -> Rule {
934        Rule::fetch_value
935    }
936    fn message() -> &'static str {
937        "CountMode"
938    }
939    fn parse_pair(
940        extensions: &SimpleExtensions,
941        pair: Pair<Rule>,
942    ) -> Result<Self, MessageParseError> {
943        assert_eq!(pair.as_rule(), Self::rule());
944        let mut arg_inner = RuleIter::from(pair.into_inner());
945        let value_pair = if let Some(int_pair) = arg_inner.try_pop(Rule::integer) {
946            int_pair
947        } else {
948            arg_inner.pop(Rule::expression)
949        };
950        match value_pair.as_rule() {
951            Rule::integer => {
952                let value = value_pair.as_str().parse::<i64>().map_err(|e| {
953                    MessageParseError::invalid(
954                        Self::message(),
955                        value_pair.as_span(),
956                        format!("Invalid integer: {e}"),
957                    )
958                })?;
959                if value < 0 {
960                    return Err(MessageParseError::invalid(
961                        Self::message(),
962                        value_pair.as_span(),
963                        format!("Fetch limit must be non-negative, got: {value}"),
964                    ));
965                }
966                Ok(CountMode::CountExpr(i64_literal_expr(value)))
967            }
968            Rule::expression => {
969                let expr = Expression::parse_pair(extensions, value_pair)?;
970                Ok(CountMode::CountExpr(Box::new(expr)))
971            }
972            _ => Err(MessageParseError::invalid(
973                Self::message(),
974                value_pair.as_span(),
975                format!("Unexpected rule for CountMode: {:?}", value_pair.as_rule()),
976            )),
977        }
978    }
979}
980
981fn i64_literal_expr(value: i64) -> Box<Expression> {
982    Box::new(Expression {
983        rex_type: Some(RexType::Literal(Literal {
984            nullable: false,
985            type_variation_reference: 0,
986            literal_type: Some(LiteralType::I64(value)),
987        })),
988    })
989}
990
991impl ScopedParsePair for OffsetMode {
992    fn rule() -> Rule {
993        Rule::fetch_value
994    }
995    fn message() -> &'static str {
996        "OffsetMode"
997    }
998    fn parse_pair(
999        extensions: &SimpleExtensions,
1000        pair: Pair<Rule>,
1001    ) -> Result<Self, MessageParseError> {
1002        assert_eq!(pair.as_rule(), Self::rule());
1003        let mut arg_inner = RuleIter::from(pair.into_inner());
1004        let value_pair = if let Some(int_pair) = arg_inner.try_pop(Rule::integer) {
1005            int_pair
1006        } else {
1007            arg_inner.pop(Rule::expression)
1008        };
1009        match value_pair.as_rule() {
1010            Rule::integer => {
1011                let value = value_pair.as_str().parse::<i64>().map_err(|e| {
1012                    MessageParseError::invalid(
1013                        Self::message(),
1014                        value_pair.as_span(),
1015                        format!("Invalid integer: {e}"),
1016                    )
1017                })?;
1018                if value < 0 {
1019                    return Err(MessageParseError::invalid(
1020                        Self::message(),
1021                        value_pair.as_span(),
1022                        format!("Fetch offset must be non-negative, got: {value}"),
1023                    ));
1024                }
1025                Ok(OffsetMode::OffsetExpr(i64_literal_expr(value)))
1026            }
1027            Rule::expression => {
1028                let expr = Expression::parse_pair(extensions, value_pair)?;
1029                Ok(OffsetMode::OffsetExpr(Box::new(expr)))
1030            }
1031            _ => Err(MessageParseError::invalid(
1032                Self::message(),
1033                value_pair.as_span(),
1034                format!("Unexpected rule for OffsetMode: {:?}", value_pair.as_rule()),
1035            )),
1036        }
1037    }
1038}
1039
1040impl RelationParsePair for FetchRel {
1041    fn rule() -> Rule {
1042        Rule::fetch_relation
1043    }
1044
1045    fn message() -> &'static str {
1046        "FetchRel"
1047    }
1048
1049    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
1050        self.advanced_extension = adv_ext;
1051        Rel {
1052            rel_type: Some(RelType::Fetch(Box::new(self))),
1053        }
1054    }
1055
1056    fn parse_pair_with_context(
1057        extensions: &SimpleExtensions,
1058        pair: Pair<Rule>,
1059        input_children: Vec<Rel>,
1060        input_field_count: usize,
1061    ) -> Result<(Self, usize), MessageParseError> {
1062        assert_eq!(pair.as_rule(), Self::rule());
1063        let input = expect_one_child(Self::message(), &pair, input_children)?;
1064        let mut iter = RuleIter::from(pair.into_inner());
1065
1066        // Extract all pairs before any validation: RuleIter's Drop panics on
1067        // incomplete consumption, so we must exhaust the iterator before any
1068        // early return. Validation runs after iter.done() below.
1069        let (limit_pair, offset_pair) = match iter.try_pop(Rule::fetch_named_arg_list) {
1070            None => {
1071                iter.pop(Rule::empty);
1072                (None, None)
1073            }
1074            Some(fetch_args_pair) => {
1075                let extractor =
1076                    ParsedNamedArgs::new(fetch_args_pair.into_inner(), Rule::fetch_named_arg)?;
1077                let (extractor, limit_pair) = extractor.pop("limit", Rule::fetch_value);
1078                let (extractor, offset_pair) = extractor.pop("offset", Rule::fetch_value);
1079                extractor.done()?;
1080                (limit_pair, offset_pair)
1081            }
1082        };
1083
1084        let reference_list_pair = iter.pop(Rule::reference_list);
1085        let (emit, output_count) = parse_emit(reference_list_pair, input_field_count);
1086        let common = RelCommon {
1087            emit_kind: Some(emit),
1088            ..Default::default()
1089        };
1090        iter.done();
1091
1092        let count_mode = limit_pair
1093            .map(|pair| CountMode::parse_pair(extensions, pair))
1094            .transpose()?;
1095        let offset_mode = offset_pair
1096            .map(|pair| OffsetMode::parse_pair(extensions, pair))
1097            .transpose()?;
1098        Ok((
1099            FetchRel {
1100                input: Some(input),
1101                common: Some(common),
1102                advanced_extension: None,
1103                offset_mode,
1104                count_mode,
1105            },
1106            output_count,
1107        ))
1108    }
1109}
1110
1111impl ParsePair for join_rel::JoinType {
1112    fn rule() -> Rule {
1113        Rule::join_type
1114    }
1115
1116    fn message() -> &'static str {
1117        "JoinType"
1118    }
1119
1120    fn parse_pair(pair: Pair<Rule>) -> Self {
1121        assert_eq!(pair.as_rule(), Self::rule());
1122        let join_type_str = pair.as_str().trim_start_matches('&');
1123        match join_type_str {
1124            "Inner" => join_rel::JoinType::Inner,
1125            "Left" => join_rel::JoinType::Left,
1126            "Right" => join_rel::JoinType::Right,
1127            "Outer" => join_rel::JoinType::Outer,
1128            "LeftSemi" => join_rel::JoinType::LeftSemi,
1129            "RightSemi" => join_rel::JoinType::RightSemi,
1130            "LeftAnti" => join_rel::JoinType::LeftAnti,
1131            "RightAnti" => join_rel::JoinType::RightAnti,
1132            "LeftSingle" => join_rel::JoinType::LeftSingle,
1133            "RightSingle" => join_rel::JoinType::RightSingle,
1134            "LeftMark" => join_rel::JoinType::LeftMark,
1135            "RightMark" => join_rel::JoinType::RightMark,
1136            _ => panic!("Unknown join type: {join_type_str} (this should be caught by grammar)"),
1137        }
1138    }
1139}
1140
1141impl RelationParsePair for JoinRel {
1142    fn rule() -> Rule {
1143        Rule::join_relation
1144    }
1145
1146    fn message() -> &'static str {
1147        "JoinRel"
1148    }
1149
1150    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
1151        self.advanced_extension = adv_ext;
1152        Rel {
1153            rel_type: Some(RelType::Join(Box::new(self))),
1154        }
1155    }
1156
1157    fn parse_pair_with_context(
1158        extensions: &SimpleExtensions,
1159        pair: Pair<Rule>,
1160        input_children: Vec<Rel>,
1161        input_field_count: usize,
1162    ) -> Result<(Self, usize), MessageParseError> {
1163        assert_eq!(pair.as_rule(), Self::rule());
1164
1165        if input_children.len() != 2 {
1166            return Err(MessageParseError::invalid(
1167                Self::message(),
1168                pair.as_span(),
1169                format!(
1170                    "JoinRel should have exactly 2 input children, got {}",
1171                    input_children.len()
1172                ),
1173            ));
1174        }
1175
1176        let mut children_iter = input_children.into_iter();
1177        let left = Box::new(children_iter.next().unwrap());
1178        let right = Box::new(children_iter.next().unwrap());
1179
1180        let mut iter = RuleIter::from(pair.into_inner());
1181        let join_type = iter.parse_next::<join_rel::JoinType>();
1182        let condition = iter.parse_next_scoped::<Expression>(extensions)?;
1183        let reference_list_pair = iter.pop(Rule::reference_list);
1184        iter.done();
1185
1186        // TODO: For semi/anti joins, the direct output width differs from
1187        // left+right — `input_field_count` would misclassify the emit as Direct.
1188        // Revisit when those join types are supported.
1189        let (emit, output_count) = parse_emit(reference_list_pair, input_field_count);
1190        let common = RelCommon {
1191            emit_kind: Some(emit),
1192            ..Default::default()
1193        };
1194
1195        Ok((
1196            JoinRel {
1197                common: Some(common),
1198                left: Some(left),
1199                right: Some(right),
1200                expression: Some(Box::new(condition)),
1201                post_join_filter: None, // not yet represented in the grammar
1202                r#type: join_type as i32,
1203                advanced_extension: None,
1204            },
1205            output_count,
1206        ))
1207    }
1208}
1209
1210impl RelationParsePair for CrossRel {
1211    fn rule() -> Rule {
1212        Rule::cross_relation
1213    }
1214
1215    fn message() -> &'static str {
1216        "CrossRel"
1217    }
1218
1219    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
1220        self.advanced_extension = adv_ext;
1221        Rel {
1222            rel_type: Some(RelType::Cross(Box::new(self))),
1223        }
1224    }
1225
1226    fn parse_pair_with_context(
1227        _extensions: &SimpleExtensions,
1228        pair: Pair<Rule>,
1229        input_children: Vec<Rel>,
1230        input_field_count: usize,
1231    ) -> Result<(Self, usize), MessageParseError> {
1232        assert_eq!(pair.as_rule(), Self::rule());
1233
1234        if input_children.len() != 2 {
1235            return Err(MessageParseError::invalid(
1236                Self::message(),
1237                pair.as_span(),
1238                format!(
1239                    "CrossRel should have exactly 2 input children, got {}",
1240                    input_children.len()
1241                ),
1242            ));
1243        }
1244
1245        let mut children_iter = input_children.into_iter();
1246        let left = Box::new(children_iter.next().unwrap());
1247        let right = Box::new(children_iter.next().unwrap());
1248
1249        let mut iter = RuleIter::from(pair.into_inner());
1250        let reference_list_pair = iter.pop(Rule::reference_list);
1251        iter.done();
1252
1253        let (emit, output_count) = parse_emit(reference_list_pair, input_field_count);
1254        let common = RelCommon {
1255            emit_kind: Some(emit),
1256            ..Default::default()
1257        };
1258
1259        Ok((
1260            CrossRel {
1261                common: Some(common),
1262                left: Some(left),
1263                right: Some(right),
1264                advanced_extension: None,
1265            },
1266            output_count,
1267        ))
1268    }
1269}
1270
1271impl ParsePair for set_rel::SetOp {
1272    fn rule() -> Rule {
1273        Rule::set_op
1274    }
1275
1276    fn message() -> &'static str {
1277        "SetOp"
1278    }
1279
1280    fn parse_pair(pair: Pair<Rule>) -> Self {
1281        assert_eq!(pair.as_rule(), Self::rule());
1282        let set_op_str = pair.as_str().trim_start_matches('&');
1283        match set_op_str {
1284            "MinusPrimary" => set_rel::SetOp::MinusPrimary,
1285            "MinusPrimaryAll" => set_rel::SetOp::MinusPrimaryAll,
1286            "MinusMultiset" => set_rel::SetOp::MinusMultiset,
1287            "IntersectionPrimary" => set_rel::SetOp::IntersectionPrimary,
1288            "IntersectionMultiset" => set_rel::SetOp::IntersectionMultiset,
1289            "IntersectionMultisetAll" => set_rel::SetOp::IntersectionMultisetAll,
1290            "UnionDistinct" => set_rel::SetOp::UnionDistinct,
1291            "UnionAll" => set_rel::SetOp::UnionAll,
1292            _ => panic!("Unknown set op: {set_op_str} (this should be caught by grammar)"),
1293        }
1294    }
1295}
1296
1297/// Parse a `set_relation` pair given the real per-child output widths (not
1298/// just their sum), so mismatched input schemas are always caught
1299pub(crate) fn parse_set_relation_pair(
1300    pair: Pair<Rule>,
1301    input_children: Vec<Rel>,
1302    child_field_counts: &[usize],
1303    advanced_extension: Option<AdvancedExtension>,
1304) -> Result<(Rel, usize), MessageParseError> {
1305    assert_eq!(pair.as_rule(), Rule::set_relation);
1306
1307    if input_children.len() < 2 {
1308        return Err(MessageParseError::invalid(
1309            "SetRel",
1310            pair.as_span(),
1311            format!(
1312                "SetRel should have at least 2 input children, got {}",
1313                input_children.len()
1314            ),
1315        ));
1316    }
1317
1318    // All inputs must share the same output width (Set is a pass-through
1319    // over a common schema, not a concatenation like Join).
1320    let child_width = child_field_counts[0];
1321    if child_field_counts.iter().any(|&w| w != child_width) {
1322        return Err(MessageParseError::invalid(
1323            "SetRel",
1324            pair.as_span(),
1325            format!(
1326                "SetRel inputs must all have the same number of columns, got widths {child_field_counts:?}"
1327            ),
1328        ));
1329    }
1330
1331    let mut iter = RuleIter::from(pair.into_inner());
1332    let op = iter.parse_next::<set_rel::SetOp>();
1333    let reference_list_pair = iter.pop(Rule::reference_list);
1334    iter.done();
1335
1336    let (emit, output_count) = parse_emit(reference_list_pair, child_width);
1337    let common = RelCommon {
1338        emit_kind: Some(emit),
1339        ..Default::default()
1340    };
1341
1342    Ok((
1343        Rel {
1344            rel_type: Some(RelType::Set(SetRel {
1345                common: Some(common),
1346                inputs: input_children,
1347                op: op as i32,
1348                advanced_extension,
1349            })),
1350        },
1351        output_count,
1352    ))
1353}
1354
1355#[cfg(test)]
1356mod tests {
1357    use pest::Parser;
1358
1359    use super::*;
1360    use crate::fixtures::TestContext;
1361    use crate::parser::{ExpressionParser, Rule};
1362
1363    #[test]
1364    fn test_parse_relation() {
1365        // Removed: test_parse_relation for old Relation struct
1366    }
1367
1368    #[test]
1369    fn test_parse_read_relation() {
1370        let extensions = SimpleExtensions::default();
1371        let read = ReadRel::parse_pair_with_context(
1372            &extensions,
1373            parse_exact(Rule::read_relation, "Read[ab.cd.ef => a:i32, b:string?]"),
1374            vec![],
1375            0,
1376        )
1377        .unwrap()
1378        .0;
1379        let names = match &read.read_type {
1380            Some(read_rel::ReadType::NamedTable(table)) => &table.names,
1381            _ => panic!("Expected NamedTable"),
1382        };
1383        assert_eq!(names, &["ab", "cd", "ef"]);
1384        let columns = &read
1385            .base_schema
1386            .as_ref()
1387            .unwrap()
1388            .r#struct
1389            .as_ref()
1390            .unwrap()
1391            .types;
1392        assert_eq!(columns.len(), 2);
1393    }
1394
1395    /// Produces a ReadRel with 3 columns: a:i32, b:string?, c:i64
1396    fn example_read_relation() -> ReadRel {
1397        let extensions = SimpleExtensions::default();
1398        ReadRel::parse_pair_with_context(
1399            &extensions,
1400            parse_exact(
1401                Rule::read_relation,
1402                "Read[ab.cd.ef => a:i32, b:string?, c:i64]",
1403            ),
1404            vec![],
1405            0,
1406        )
1407        .unwrap()
1408        .0
1409    }
1410
1411    #[test]
1412    fn test_parse_filter_relation() {
1413        let extensions = SimpleExtensions::default();
1414        let filter = FilterRel::parse_pair_with_context(
1415            &extensions,
1416            parse_exact(Rule::filter_relation, "Filter[$1 => $0, $1, $2]"),
1417            vec![example_read_relation().into_rel(None)],
1418            3,
1419        )
1420        .unwrap()
1421        .0;
1422        // Identity mapping [0, 1, 2] over 3 inputs → Direct
1423        let emit_kind = filter.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
1424        assert!(
1425            matches!(emit_kind, EmitKind::Direct(_)),
1426            "Expected Direct for identity emit, got {emit_kind:?}"
1427        );
1428    }
1429
1430    #[test]
1431    fn test_parse_project_relation() {
1432        let extensions = SimpleExtensions::default();
1433        let project = ProjectRel::parse_pair_with_context(
1434            &extensions,
1435            parse_exact(Rule::project_relation, "Project[$0, $1, 42]"),
1436            vec![example_read_relation().into_rel(None)],
1437            3,
1438        )
1439        .unwrap()
1440        .0;
1441
1442        // Should have 1 expression (42) and 2 references ($0, $1)
1443        assert_eq!(project.expressions.len(), 1);
1444
1445        let emit_kind = &project.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
1446        let emit = match emit_kind {
1447            EmitKind::Emit(emit) => &emit.output_mapping,
1448            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
1449        };
1450        // Output mapping should be [0, 1, 3]. References are 0-2; expression is 3.
1451        assert_eq!(emit, &[0, 1, 3]);
1452    }
1453
1454    #[test]
1455    fn test_parse_project_relation_complex() {
1456        let extensions = SimpleExtensions::default();
1457        let project = ProjectRel::parse_pair_with_context(
1458            &extensions,
1459            parse_exact(Rule::project_relation, "Project[42, $0, 100, $2, $1]"),
1460            vec![example_read_relation().into_rel(None)],
1461            5, // Assume 5 input fields
1462        )
1463        .unwrap()
1464        .0;
1465
1466        // Should have 2 expressions (42, 100) and 3 references ($0, $2, $1)
1467        assert_eq!(project.expressions.len(), 2);
1468
1469        let emit_kind = &project.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
1470        let emit = match emit_kind {
1471            EmitKind::Emit(emit) => &emit.output_mapping,
1472            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
1473        };
1474        // Direct mapping: [input_fields..., 42, 100] (input fields first, then expressions)
1475        // Output mapping: [5, 0, 6, 2, 1] (to get: 42, $0, 100, $2, $1)
1476        assert_eq!(emit, &[5, 0, 6, 2, 1]);
1477    }
1478
1479    #[test]
1480    fn test_parse_aggregate_relation() {
1481        let extensions = TestContext::new()
1482            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
1483            .with_function(1, 10, "sum")
1484            .with_function(1, 11, "count")
1485            .extensions;
1486
1487        let aggregate = AggregateRel::parse_pair_with_context(
1488            &extensions,
1489            parse_exact(
1490                Rule::aggregate_relation,
1491                "Aggregate[($0, $1), _ => sum($2):i64, $0, count($2):i64]",
1492            ),
1493            vec![example_read_relation().into_rel(None)],
1494            3,
1495        )
1496        .unwrap()
1497        .0;
1498
1499        // Should have 2 group-by sets ($0, $1) and an empty group, and emit 2 measures (sum($2), count($2))
1500        assert_eq!(aggregate.grouping_expressions.len(), 2);
1501        assert_eq!(aggregate.groupings[0].expression_references.len(), 2);
1502        assert_eq!(aggregate.groupings.len(), 2);
1503        assert_eq!(aggregate.measures.len(), 2);
1504
1505        let emit_kind = &aggregate
1506            .common
1507            .as_ref()
1508            .unwrap()
1509            .emit_kind
1510            .as_ref()
1511            .unwrap();
1512        let emit = match emit_kind {
1513            EmitKind::Emit(emit) => &emit.output_mapping,
1514            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
1515        };
1516        // Output mapping should be [2, 0, 3] (measures and group-by fields in order)
1517        // sum($2) -> 2, $0 -> 0, count($2) -> 3
1518        assert_eq!(emit, &[2, 0, 3]);
1519    }
1520
1521    #[test]
1522    fn test_parse_aggregate_relation_maintain_column_order() {
1523        let extensions = TestContext::new()
1524            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
1525            .with_function(1, 10, "sum")
1526            .with_function(1, 11, "count")
1527            .extensions;
1528
1529        let aggregate = AggregateRel::parse_pair_with_context(
1530            &extensions,
1531            parse_exact(
1532                Rule::aggregate_relation,
1533                "Aggregate[$0 => sum($1):i64, $0, count($1):i64]",
1534            ),
1535            vec![example_read_relation().into_rel(None)],
1536            3,
1537        )
1538        .unwrap()
1539        .0;
1540
1541        // Should have 1 group-by field ($0) and 2 measures (sum($1), count($1))
1542        assert_eq!(aggregate.grouping_expressions.len(), 1);
1543        assert_eq!(aggregate.groupings.len(), 1);
1544        assert_eq!(aggregate.measures.len(), 2);
1545
1546        let emit_kind = &aggregate
1547            .common
1548            .as_ref()
1549            .unwrap()
1550            .emit_kind
1551            .as_ref()
1552            .unwrap();
1553        let emit = match emit_kind {
1554            EmitKind::Emit(emit) => &emit.output_mapping,
1555            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
1556        };
1557        // Output mapping should be [1, 0, 2] (grouping fields + measures)
1558        assert_eq!(emit, &[1, 0, 2]);
1559    }
1560
1561    #[test]
1562    fn test_parse_aggregate_relation_simple() {
1563        let extensions = TestContext::new()
1564            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
1565            .with_function(1, 10, "sum")
1566            .extensions;
1567
1568        let aggregate = AggregateRel::parse_pair_with_context(
1569            &extensions,
1570            parse_exact(Rule::aggregate_relation, "Aggregate[$2, $0 => sum($1):i64]"),
1571            vec![example_read_relation().into_rel(None)],
1572            3,
1573        )
1574        .unwrap()
1575        .0;
1576
1577        assert_eq!(aggregate.grouping_expressions.len(), 2);
1578        assert_eq!(aggregate.groupings.len(), 1);
1579        // expression_references must be positions [0, 1], not raw field indices [2, 0]
1580        assert_eq!(aggregate.groupings[0].expression_references, vec![0, 1]);
1581    }
1582
1583    #[test]
1584    fn test_parse_aggregate_relation_global_aggregate() {
1585        let extensions = TestContext::new()
1586            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
1587            .with_function(1, 10, "sum")
1588            .with_function(1, 11, "count")
1589            .extensions;
1590
1591        let aggregate = AggregateRel::parse_pair_with_context(
1592            &extensions,
1593            parse_exact(
1594                Rule::aggregate_relation,
1595                "Aggregate[_ => sum($0):i64, count($1):i64]",
1596            ),
1597            vec![example_read_relation().into_rel(None)],
1598            3,
1599        )
1600        .unwrap()
1601        .0;
1602
1603        // Should have 0 group-by fields and 2 measures
1604        assert_eq!(aggregate.grouping_expressions.len(), 0);
1605        assert_eq!(aggregate.groupings.len(), 1);
1606        assert_eq!(aggregate.groupings[0].expression_references.len(), 0);
1607        assert_eq!(aggregate.measures.len(), 2);
1608
1609        // Identity mapping [0, 1] over 2 outputs (0 grouping + 2 measures) → Direct
1610        let emit_kind = aggregate
1611            .common
1612            .as_ref()
1613            .unwrap()
1614            .emit_kind
1615            .as_ref()
1616            .unwrap();
1617        assert!(
1618            matches!(emit_kind, EmitKind::Direct(_)),
1619            "Expected Direct for identity emit, got {emit_kind:?}"
1620        );
1621    }
1622
1623    #[test]
1624    fn test_parse_aggregate_relation_grouping_sets() {
1625        let extensions = TestContext::new()
1626            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
1627            .with_function(1, 11, "count")
1628            .extensions;
1629
1630        let read_rel = ReadRel::parse_pair_with_context(
1631            &extensions,
1632            parse_exact(
1633                Rule::read_relation,
1634                "Read[ab.cd.ef => a:i32, b:string?, c:i64, d:i64]",
1635            ),
1636            vec![],
1637            0,
1638        )
1639        .unwrap()
1640        .0;
1641
1642        let aggregate = AggregateRel::parse_pair_with_context(
1643            &extensions,
1644            parse_exact(
1645                Rule::aggregate_relation,
1646                "Aggregate[($0, $1, $2), ($2, $0), ($1), _ => $0, $1, $2, count($3):i64]",
1647            ),
1648            vec![read_rel.into_rel(None)],
1649            4,
1650        )
1651        .unwrap()
1652        .0;
1653
1654        assert_eq!(aggregate.grouping_expressions.len(), 3);
1655        assert_eq!(aggregate.groupings.len(), 4);
1656        // ($0, $1, $2) -> [0, 1, 2]
1657        assert_eq!(aggregate.groupings[0].expression_references, vec![0, 1, 2]);
1658        // ($2, $0) -> [2, 0] (reuses indices, different order)
1659        assert_eq!(aggregate.groupings[1].expression_references, vec![2, 0]);
1660        // ($1) -> [1]
1661        assert_eq!(aggregate.groupings[2].expression_references, vec![1]);
1662        // _ -> empty
1663        assert!(aggregate.groupings[3].expression_references.is_empty());
1664        assert_eq!(aggregate.measures.len(), 1);
1665    }
1666
1667    #[test]
1668    fn test_fetch_relation_positive_values() {
1669        let extensions = SimpleExtensions::default();
1670
1671        // Test valid positive values should work
1672        let fetch_rel = FetchRel::parse_pair_with_context(
1673            &extensions,
1674            parse_exact(Rule::fetch_relation, "Fetch[limit=10, offset=5 => $0]"),
1675            vec![example_read_relation().into_rel(None)],
1676            3,
1677        )
1678        .unwrap()
1679        .0;
1680
1681        // Verify the limit and offset values are correct
1682        assert_eq!(
1683            fetch_rel.count_mode,
1684            Some(CountMode::CountExpr(i64_literal_expr(10)))
1685        );
1686        assert_eq!(
1687            fetch_rel.offset_mode,
1688            Some(OffsetMode::OffsetExpr(i64_literal_expr(5)))
1689        );
1690    }
1691
1692    #[test]
1693    fn test_fetch_relation_negative_limit_rejected() {
1694        let extensions = SimpleExtensions::default();
1695
1696        // Test that fetch relations with negative limits are properly rejected
1697        let parsed_result = ExpressionParser::parse(Rule::fetch_relation, "Fetch[limit=-5 => $0]");
1698        if let Ok(mut pairs) = parsed_result {
1699            let pair = pairs.next().unwrap();
1700            if pair.as_str() == "Fetch[limit=-5 => $0]" {
1701                // Full parse succeeded, now test that validation catches the negative value
1702                let result = FetchRel::parse_pair_with_context(
1703                    &extensions,
1704                    pair,
1705                    vec![example_read_relation().into_rel(None)],
1706                    3,
1707                );
1708                assert!(result.is_err());
1709                let error_msg = result.unwrap_err().to_string();
1710                assert!(error_msg.contains("Fetch limit must be non-negative"));
1711            } else {
1712                // If grammar doesn't fully support negative values, that's also acceptable
1713                // since it would prevent negative values at parse time
1714                println!("Grammar prevents negative limit values at parse time");
1715            }
1716        } else {
1717            // Grammar doesn't support negative values in fetch context
1718            println!("Grammar prevents negative limit values at parse time");
1719        }
1720    }
1721
1722    #[test]
1723    fn test_fetch_relation_negative_offset_rejected() {
1724        let extensions = SimpleExtensions::default();
1725
1726        // Test that fetch relations with negative offsets are properly rejected
1727        let parsed_result =
1728            ExpressionParser::parse(Rule::fetch_relation, "Fetch[offset=-10 => $0]");
1729        if let Ok(mut pairs) = parsed_result {
1730            let pair = pairs.next().unwrap();
1731            if pair.as_str() == "Fetch[offset=-10 => $0]" {
1732                // Full parse succeeded, now test that validation catches the negative value
1733                let result = FetchRel::parse_pair_with_context(
1734                    &extensions,
1735                    pair,
1736                    vec![example_read_relation().into_rel(None)],
1737                    3,
1738                );
1739                assert!(result.is_err());
1740                let error_msg = result.unwrap_err().to_string();
1741                assert!(error_msg.contains("Fetch offset must be non-negative"));
1742            } else {
1743                // If grammar doesn't fully support negative values, that's also acceptable
1744                println!("Grammar prevents negative offset values at parse time");
1745            }
1746        } else {
1747            // Grammar doesn't support negative values in fetch context
1748            println!("Grammar prevents negative offset values at parse time");
1749        }
1750    }
1751
1752    #[test]
1753    fn test_parse_join_relation() {
1754        let extensions = TestContext::new()
1755            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
1756            .with_function(1, 10, "eq")
1757            .extensions;
1758
1759        let left_rel = example_read_relation().into_rel(None);
1760        let right_rel = example_read_relation().into_rel(None);
1761
1762        let join = JoinRel::parse_pair_with_context(
1763            &extensions,
1764            parse_exact(
1765                Rule::join_relation,
1766                "Join[&Inner, eq($0, $3):boolean => $0, $1, $3, $4]",
1767            ),
1768            vec![left_rel, right_rel],
1769            6, // left (3) + right (3) = 6 total input fields
1770        )
1771        .unwrap()
1772        .0;
1773
1774        // Should be an Inner join
1775        assert_eq!(join.r#type, join_rel::JoinType::Inner as i32);
1776
1777        // Should have left and right relations
1778        assert!(join.left.is_some());
1779        assert!(join.right.is_some());
1780
1781        // Should have a join condition
1782        assert!(join.expression.is_some());
1783
1784        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
1785        let emit = match emit_kind {
1786            EmitKind::Emit(emit) => &emit.output_mapping,
1787            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
1788        };
1789        // Output mapping should be [0, 1, 3, 4] (selected columns)
1790        assert_eq!(emit, &[0, 1, 3, 4]);
1791    }
1792
1793    #[test]
1794    fn test_parse_join_relation_left_outer() {
1795        let extensions = TestContext::new()
1796            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
1797            .with_function(1, 10, "eq")
1798            .extensions;
1799
1800        let left_rel = example_read_relation().into_rel(None);
1801        let right_rel = example_read_relation().into_rel(None);
1802
1803        let join = JoinRel::parse_pair_with_context(
1804            &extensions,
1805            parse_exact(
1806                Rule::join_relation,
1807                "Join[&Left, eq($0, $3):boolean => $0, $1, $2]",
1808            ),
1809            vec![left_rel, right_rel],
1810            6,
1811        )
1812        .unwrap()
1813        .0;
1814
1815        // Should be a Left join
1816        assert_eq!(join.r#type, join_rel::JoinType::Left as i32);
1817
1818        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
1819        let emit = match emit_kind {
1820            EmitKind::Emit(emit) => &emit.output_mapping,
1821            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
1822        };
1823        // Output mapping should be [0, 1, 2]
1824        assert_eq!(emit, &[0, 1, 2]);
1825    }
1826
1827    #[test]
1828    fn test_parse_join_relation_left_semi() {
1829        let extensions = TestContext::new()
1830            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
1831            .with_function(1, 10, "eq")
1832            .extensions;
1833
1834        let left_rel = example_read_relation().into_rel(None);
1835        let right_rel = example_read_relation().into_rel(None);
1836
1837        let join = JoinRel::parse_pair_with_context(
1838            &extensions,
1839            parse_exact(
1840                Rule::join_relation,
1841                "Join[&LeftSemi, eq($0, $3):boolean => $0, $1]",
1842            ),
1843            vec![left_rel, right_rel],
1844            6,
1845        )
1846        .unwrap()
1847        .0;
1848
1849        // Should be a LeftSemi join
1850        assert_eq!(join.r#type, join_rel::JoinType::LeftSemi as i32);
1851
1852        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
1853        let emit = match emit_kind {
1854            EmitKind::Emit(emit) => &emit.output_mapping,
1855            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
1856        };
1857        // Output mapping should be [0, 1] (only left columns for semi join)
1858        assert_eq!(emit, &[0, 1]);
1859    }
1860
1861    #[test]
1862    fn test_parse_join_relation_requires_two_children() {
1863        let extensions = SimpleExtensions::default();
1864
1865        // Test with 0 children
1866        let result = JoinRel::parse_pair_with_context(
1867            &extensions,
1868            parse_exact(
1869                Rule::join_relation,
1870                "Join[&Inner, eq($0, $1):boolean => $0, $1]",
1871            ),
1872            vec![],
1873            0,
1874        );
1875        assert!(result.is_err());
1876
1877        // Test with 1 child
1878        let result = JoinRel::parse_pair_with_context(
1879            &extensions,
1880            parse_exact(
1881                Rule::join_relation,
1882                "Join[&Inner, eq($0, $1):boolean => $0, $1]",
1883            ),
1884            vec![example_read_relation().into_rel(None)],
1885            3,
1886        );
1887        assert!(result.is_err());
1888
1889        // Test with 3 children
1890        let result = JoinRel::parse_pair_with_context(
1891            &extensions,
1892            parse_exact(
1893                Rule::join_relation,
1894                "Join[&Inner, eq($0, $1):boolean => $0, $1]",
1895            ),
1896            vec![
1897                example_read_relation().into_rel(None),
1898                example_read_relation().into_rel(None),
1899                example_read_relation().into_rel(None),
1900            ],
1901            9,
1902        );
1903        assert!(result.is_err());
1904    }
1905
1906    fn parse_exact(rule: Rule, input: &'_ str) -> pest::iterators::Pair<'_, Rule> {
1907        let mut pairs = ExpressionParser::parse(rule, input).unwrap();
1908        assert_eq!(pairs.as_str(), input);
1909        let pair = pairs.next().unwrap();
1910        assert_eq!(pairs.next(), None);
1911        pair
1912    }
1913}
substrait_explain/parser/relations.rs

substrait_explain/parser/
relations.rs
