substrait_explain/parser/

relations.rs

use std::collections::HashMap;

use pest::iterators::Pair;
use prost::Message;
use substrait::proto::aggregate_rel::Grouping;
use substrait::proto::expression::literal::LiteralType;
use substrait::proto::expression::{Literal, RexType, nested};
use substrait::proto::extensions::AdvancedExtension;
use substrait::proto::fetch_rel::{CountMode, OffsetMode};
use substrait::proto::rel::RelType;
use substrait::proto::rel_common::{Direct, Emit, EmitKind};
use substrait::proto::sort_field::SortKind;
use substrait::proto::{
    AggregateRel, CrossRel, Expression, FetchRel, FilterRel, JoinRel, NamedStruct, ProjectRel,
    ReadRel, Rel, RelCommon, SetRel, SortField, SortRel, Type, aggregate_rel, join_rel, read_rel,
    set_rel, r#type,
};

use super::{
    MessageParseError, ParsePair, ParsedNamedArgs, Rule, RuleIter, ScopedParsePair,
    sort_direction_from_str, unwrap_single_pair,
};
use crate::extensions::any::Any;
use crate::extensions::registry::ExtensionError;
use crate::extensions::{AddendumKind, ExtensionArgs, ExtensionRegistry, SimpleExtensions};
use crate::parser::errors::{ParseContext, ParseError};
use crate::parser::expressions::{FieldIndex, Name};

/// Parsing context for relations that includes extensions, registry, and optional warning collection
pub struct RelationParsingContext<'a> {
    pub registry: &'a ExtensionRegistry,
    pub line_no: i64,
    pub line: &'a str,
}

impl<'a> RelationParsingContext<'a> {
    /// Resolve extension detail using registry. Any failure is treated as a hard parse error.
    pub fn resolve_extension_detail(
        &self,
        extension_name: &str,
        extension_args: &ExtensionArgs,
    ) -> Result<Option<Any>, ParseError> {
        let detail = self
            .registry
            .parse_extension(extension_name, extension_args);

        match detail {
            Ok(any) => Ok(Some(any)),
            Err(ExtensionError::NotFound { .. }) => Err(ParseError::UnregisteredExtension {
                name: extension_name.to_string(),
                context: ParseContext::new(self.line_no, self.line.to_string()),
            }),
            Err(err) => Err(ParseError::ExtensionDetail(
                ParseContext::new(self.line_no, self.line.to_string()),
                err,
            )),
        }
    }

    /// Resolve an addendum detail using the registry entry for its kind.
    /// Any failure is treated as a hard parse error.
    pub(crate) fn resolve_addendum_detail(
        &self,
        kind: AddendumKind,
        name: &str,
        args: &ExtensionArgs,
    ) -> Result<Any, ParseError> {
        let result = match kind {
            AddendumKind::Enhancement => self.registry.parse_enhancement(name, args),
            AddendumKind::Optimization => self.registry.parse_optimization(name, args),
            AddendumKind::ExtensionTable => self.registry.parse_extension_table(name, args),
        };
        result.map_err(|err| match err {
            ExtensionError::NotFound { .. } => ParseError::UnregisteredExtension {
                name: name.to_string(),
                context: ParseContext::new(self.line_no, self.line.to_string()),
            },
            err => ParseError::ExtensionDetail(
                ParseContext::new(self.line_no, self.line.to_string()),
                err,
            ),
        })
    }
}

/// A trait for parsing relations with full context for tree building.
pub trait RelationParsePair: Sized {
    fn rule() -> Rule;
    fn message() -> &'static str;

    /// Parse the grammar pair into this relation type and its output field
    /// count.
    ///
    /// Returns `(Self, usize)` where `usize` is the output field count —
    /// computed during parsing when `input_field_count` is available.
    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError>;

    /// Consume this parsed relation, apply the advanced extension, and produce
    /// the final `Rel`.
    fn into_rel(self, adv_ext: Option<AdvancedExtension>) -> Rel;
}

pub struct TableName(Vec<String>);

impl ParsePair for TableName {
    fn rule() -> Rule {
        Rule::table_name
    }

    fn message() -> &'static str {
        "TableName"
    }

    fn parse_pair(pair: Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());
        let pairs = pair.into_inner();
        let mut names = Vec::with_capacity(pairs.len());
        let mut iter = RuleIter::from(pairs);
        while let Some(name) = iter.parse_if_next::<Name>() {
            names.push(name.0);
        }
        iter.done();
        Self(names)
    }
}

#[derive(Debug, Clone)]
pub struct Column {
    pub name: String,
    pub typ: Type,
}

impl ScopedParsePair for Column {
    fn rule() -> Rule {
        Rule::named_column
    }

    fn message() -> &'static str {
        "Column"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut iter = RuleIter::from(pair.into_inner());
        let name = iter.parse_next::<Name>().0;
        let typ = iter.parse_next_scoped(extensions)?;
        iter.done();
        Ok(Self { name, typ })
    }
}

pub(crate) struct NamedColumnList(pub(crate) Vec<Column>);

impl ScopedParsePair for NamedColumnList {
    fn rule() -> Rule {
        Rule::named_column_list
    }

    fn message() -> &'static str {
        "NamedColumnList"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut columns = Vec::new();
        for col in pair.into_inner() {
            columns.push(Column::parse_pair(extensions, col)?);
        }
        Ok(Self(columns))
    }
}

/// This is a utility function for extracting a single child from the list of
/// children, to be used in the RelationParsePair trait. The RelationParsePair
/// trait passes a Vec of children, because some relations have multiple
/// children - but most accept exactly one child.
pub(crate) fn expect_one_child(
    message: &'static str,
    pair: &Pair<Rule>,
    mut input_children: Vec<Rel>,
) -> Result<Box<Rel>, MessageParseError> {
    match input_children.len() {
        0 => Err(MessageParseError::invalid(
            message,
            pair.as_span(),
            format!("{message} missing child"),
        )),
        1 => Ok(Box::new(input_children.pop().unwrap())),
        n => Err(MessageParseError::invalid(
            message,
            pair.as_span(),
            format!("{message} should have 1 input child, got {n}"),
        )),
    }
}

/// Parse a reference list Pair and return an EmitKind::Emit.
/// Parse a reference list into field indices for emit mapping.
fn parse_output_mapping(pair: Pair<Rule>) -> Vec<i32> {
    assert_eq!(pair.as_rule(), Rule::reference_list);
    pair.into_inner()
        .map(|p| FieldIndex::parse_pair(p).0)
        .collect()
}

/// Build an emit: `Direct` if the mapping is the identity `[0, 1, ..., N-1]`
/// (where N = `direct_output_count`), otherwise `Emit` with the explicit mapping.
fn make_emit(output_mapping: Vec<i32>, direct_output_count: usize) -> EmitKind {
    let is_identity = output_mapping.len() == direct_output_count
        && output_mapping
            .iter()
            .enumerate()
            .all(|(i, &v)| v == i as i32);
    if is_identity {
        EmitKind::Direct(Direct {})
    } else {
        EmitKind::Emit(Emit { output_mapping })
    }
}

/// A [`RelCommon`] with an explicit `Direct` emit.
pub(crate) fn direct_common() -> RelCommon {
    RelCommon {
        emit_kind: Some(EmitKind::Direct(Direct {})),
        ..Default::default()
    }
}

/// Parse a reference list into an emit and output field count.
fn parse_emit(reference_list: Pair<Rule>, direct_output_count: usize) -> (EmitKind, usize) {
    let output_mapping = parse_output_mapping(reference_list);
    let output_count = output_mapping.len();
    let emit = make_emit(output_mapping, direct_output_count);
    (emit, output_count)
}

fn parse_output(
    extensions: &SimpleExtensions,
    output: Pair<Rule>,
) -> Result<(Vec<Column>, Option<RelCommon>, usize), MessageParseError> {
    assert_eq!(output.as_rule(), Rule::output);
    let output = unwrap_single_pair(output);
    match output.as_rule() {
        Rule::implicit_output => {
            let mut iter = RuleIter::from(output.into_inner());
            let columns = iter.parse_next_scoped::<NamedColumnList>(extensions)?.0;
            iter.done();
            let output_count = columns.len();
            Ok((columns, Some(direct_common()), output_count))
        }
        Rule::direct_output => {
            let mut iter = RuleIter::from(output.into_inner());
            let columns = iter.parse_next_scoped::<NamedColumnList>(extensions)?.0;
            let emit_suffix = iter.pop(Rule::emit_suffix);
            iter.done();

            let direct_output_count = columns.len();
            let (emit, output_count) = parse_emit_suffix(emit_suffix)
                .unwrap_or((EmitKind::Direct(Direct {}), direct_output_count));
            let common = Some(RelCommon {
                emit_kind: Some(emit),
                ..Default::default()
            });
            Ok((columns, common, output_count))
        }
        other => unreachable!("Unexpected rule in output: {other:?}"),
    }
}

fn parse_emit_suffix(suffix: Pair<Rule>) -> Option<(EmitKind, usize)> {
    assert_eq!(suffix.as_rule(), Rule::emit_suffix);
    let reference_list = suffix.into_inner().next()?;
    assert_eq!(reference_list.as_rule(), Rule::reference_list);
    let output_mapping = parse_output_mapping(reference_list);
    let output_count = output_mapping.len();
    Some((EmitKind::Emit(Emit { output_mapping }), output_count))
}

impl RelationParsePair for ReadRel {
    fn rule() -> Rule {
        Rule::read_relation
    }

    fn message() -> &'static str {
        "ReadRel"
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        // ReadRel is a leaf node - it should have no input children and 0 input fields
        if !input_children.is_empty() {
            return Err(MessageParseError::invalid(
                Self::message(),
                pair.as_span(),
                "ReadRel should have no input children",
            ));
        }
        if input_field_count != 0 {
            return Err(MessageParseError::invalid(
                "ReadRel",
                pair.as_span(),
                "ReadRel should have 0 input fields",
            ));
        }

        let mut iter = RuleIter::from(pair.into_inner());
        let table = iter.parse_next::<TableName>().0;
        let output_pair = iter.pop(Rule::output);
        iter.done();

        let (columns, common, output_count) = parse_output(extensions, output_pair)?;

        Ok((
            ReadRel {
                base_schema: Some(build_named_struct(columns)),
                read_type: Some(read_rel::ReadType::NamedTable(read_rel::NamedTable {
                    names: table,
                    advanced_extension: None,
                })),
                common,
                ..Default::default()
            },
            output_count,
        ))
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Read(Box::new(self))),
        }
    }
}

/// Parsed `Read:Virtual[rows => columns]` relation. Needs a newtype because the
/// proto type is `ReadRel` (same as `NamedTable`), but the grammar and handling
/// are different.
pub(crate) struct VirtualReadRel(ReadRel);

impl RelationParsePair for VirtualReadRel {
    fn rule() -> Rule {
        Rule::virtual_read_relation
    }

    fn message() -> &'static str {
        "VirtualReadRel"
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        _input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        if !input_children.is_empty() {
            return Err(MessageParseError::invalid(
                Self::message(),
                pair.as_span(),
                "Read:Virtual should have no input children",
            ));
        }

        let mut iter = RuleIter::from(pair.into_inner());
        let rows_pair = iter.pop(Rule::virtual_read_rows);
        let filter = iter
            .try_pop(Rule::virtual_read_filter)
            .map(|pair| {
                let expression_pair = unwrap_single_pair(pair);
                Expression::parse_pair(extensions, expression_pair).map(Box::new)
            })
            .transpose()?;
        let columns_pair = iter.pop(Rule::named_column_list);
        iter.done();

        let rows = parse_virtual_read_rows(extensions, rows_pair)?;
        let columns = NamedColumnList::parse_pair(extensions, columns_pair)?.0;

        // TODO: Validate that each row has the same number of expressions as
        // columns. Currently no parser-side warning mechanism exists, and while
        // this is an invalid plan, it is constructible as Substrait. Consider
        // adding once a warning collection path is available.
        let output_count = columns.len();
        Ok((
            VirtualReadRel(ReadRel {
                common: Some(direct_common()),
                base_schema: Some(build_named_struct(columns)),
                read_type: Some(read_rel::ReadType::VirtualTable(read_rel::VirtualTable {
                    expressions: rows,
                    ..Default::default()
                })),
                filter,
                ..Default::default()
            }),
            output_count,
        ))
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.0.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Read(Box::new(self.0))),
        }
    }
}

/// Parsed `Read:Extension[columns]` relation. Needs a newtype because the
/// proto type is `ReadRel` (same as `NamedTable` and `VirtualTable`), but the
/// read detail is supplied by a required `+ Ext:` addendum.
pub(crate) struct ExtensionReadRel(ReadRel);

impl ExtensionReadRel {
    pub(crate) fn parse_pair_with_detail(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
        detail: Any,
        advanced_extension: Option<AdvancedExtension>,
    ) -> Result<(Rel, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Rule::extension_read_relation);
        if !input_children.is_empty() {
            return Err(MessageParseError::invalid(
                "ExtensionReadRel",
                pair.as_span(),
                "Read:Extension should have no input children",
            ));
        }
        if input_field_count != 0 {
            return Err(MessageParseError::invalid(
                "ExtensionReadRel",
                pair.as_span(),
                "Read:Extension should have 0 input fields",
            ));
        }

        let mut iter = RuleIter::from(pair.into_inner());
        let columns = iter.parse_next_scoped::<NamedColumnList>(extensions)?.0;
        iter.done();

        let output_count = columns.len();
        let rel = ExtensionReadRel(ReadRel {
            common: Some(direct_common()),
            base_schema: Some(build_named_struct(columns)),
            read_type: Some(read_rel::ReadType::ExtensionTable(
                read_rel::ExtensionTable {
                    detail: Some(detail.into()),
                },
            )),
            advanced_extension,
            ..Default::default()
        })
        .into_rel();

        Ok((rel, output_count))
    }

    fn into_rel(self) -> Rel {
        Rel {
            rel_type: Some(RelType::Read(Box::new(self.0))),
        }
    }
}

/// Build a `NamedStruct` from parsed columns.
pub(crate) fn build_named_struct(columns: Vec<Column>) -> NamedStruct {
    let (names, types): (Vec<_>, Vec<_>) = columns.into_iter().map(|c| (c.name, c.typ)).unzip();
    NamedStruct {
        names,
        r#struct: Some(r#type::Struct {
            types,
            type_variation_reference: 0,
            nullability: r#type::Nullability::Required as i32,
        }),
    }
}

/// `Read:Virtual` rows: either `empty` or a list of row tuples.
fn parse_virtual_read_rows(
    extensions: &SimpleExtensions,
    pair: Pair<Rule>,
) -> Result<Vec<nested::Struct>, MessageParseError> {
    assert_eq!(pair.as_rule(), Rule::virtual_read_rows);
    let inner = unwrap_single_pair(pair);
    match inner.as_rule() {
        Rule::empty => Ok(vec![]),
        Rule::virtual_row_list => inner
            .into_inner()
            .map(|row| parse_virtual_row(extensions, row))
            .collect(),
        _ => unreachable!(
            "Unexpected rule in virtual_read_rows: {:?}",
            inner.as_rule()
        ),
    }
}

/// Parse a single `virtual_row` (`(expr, expr, ...)` or `()`) into a `nested::Struct`.
fn parse_virtual_row(
    extensions: &SimpleExtensions,
    pair: Pair<Rule>,
) -> Result<nested::Struct, MessageParseError> {
    assert_eq!(pair.as_rule(), Rule::virtual_row);
    let fields = match pair.into_inner().next() {
        Some(expression_list) => {
            assert_eq!(expression_list.as_rule(), Rule::expression_list);
            parse_expression_list(extensions, expression_list)?
        }
        // Empty virtual row. An unusual but valid case.
        None => vec![],
    };
    Ok(nested::Struct { fields })
}

impl RelationParsePair for FilterRel {
    fn rule() -> Rule {
        Rule::filter_relation
    }

    fn message() -> &'static str {
        "FilterRel"
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Filter(Box::new(self))),
        }
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let input = expect_one_child(Self::message(), &pair, input_children)?;
        let mut iter = RuleIter::from(pair.into_inner());
        let condition = iter.parse_next_scoped::<Expression>(extensions)?;
        let references_pair = iter.pop(Rule::reference_list);
        iter.done();

        let (emit, output_count) = parse_emit(references_pair, input_field_count);
        let common = RelCommon {
            emit_kind: Some(emit),
            ..Default::default()
        };

        Ok((
            FilterRel {
                input: Some(input),
                condition: Some(Box::new(condition)),
                common: Some(common),
                advanced_extension: None,
            },
            output_count,
        ))
    }
}

impl RelationParsePair for ProjectRel {
    fn rule() -> Rule {
        Rule::project_relation
    }

    fn message() -> &'static str {
        "ProjectRel"
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Project(Box::new(self))),
        }
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let input = expect_one_child(Self::message(), &pair, input_children)?;

        let arguments_pair = unwrap_single_pair(pair);

        let mut expressions = Vec::new();
        let mut output_mapping = Vec::new();

        for arg in arguments_pair.into_inner() {
            let inner_arg = unwrap_single_pair(arg);
            match inner_arg.as_rule() {
                Rule::reference => {
                    let field_index = FieldIndex::parse_pair(inner_arg);
                    output_mapping.push(field_index.0);
                }
                Rule::expression => {
                    let expr = Expression::parse_pair(extensions, inner_arg)?;
                    expressions.push(expr);
                    // Index into the combined schema: [input fields][computed expressions].
                    output_mapping.push(input_field_count as i32 + (expressions.len() as i32 - 1));
                }
                _ => panic!("Unexpected inner argument rule: {:?}", inner_arg.as_rule()),
            }
        }

        let output_count = output_mapping.len();
        let direct_count = input_field_count + expressions.len();
        let emit = make_emit(output_mapping, direct_count);
        let common = RelCommon {
            emit_kind: Some(emit),
            ..Default::default()
        };

        Ok((
            ProjectRel {
                input: Some(input),
                expressions,
                common: Some(common),
                advanced_extension: None,
            },
            output_count,
        ))
    }
}

impl RelationParsePair for AggregateRel {
    fn rule() -> Rule {
        Rule::aggregate_relation
    }

    fn message() -> &'static str {
        "AggregateRel"
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Aggregate(Box::new(self))),
        }
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        // Aggregate defines its own output schema (grouping keys + measures),
        // so the input field count isn't needed for emit construction.
        _input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let input = expect_one_child(Self::message(), &pair, input_children)?;
        let mut iter = RuleIter::from(pair.into_inner());
        let group_by_pair = iter.pop(Rule::aggregate_group_by);
        let output_pair = iter.pop(Rule::aggregate_output);
        iter.done();

        let inner = group_by_pair
            .into_inner()
            .next()
            .expect("aggregate_group_by must have one inner item");

        let grouping_sets = parse_grouping_sets(extensions, inner)?;
        let (groupings, grouping_expressions) = build_grouping_fields(&grouping_sets);

        let (measures, output_mapping) =
            parse_aggregate_output(extensions, output_pair, &grouping_expressions)?;

        let output_count = output_mapping.len();
        let direct_count = grouping_expressions.len() + measures.len();
        let emit = make_emit(output_mapping, direct_count);
        let common = RelCommon {
            emit_kind: Some(emit),
            ..Default::default()
        };

        Ok((
            AggregateRel {
                input: Some(input),
                grouping_expressions,
                groupings,
                measures,
                common: Some(common),
                advanced_extension: None,
            },
            output_count,
        ))
    }
}

/// Encodes an expression to bytes for use as a structural-equality map key.
/// TODO: use a better key here than encoding to bytes. Ideally, substrait-rs
/// would support `PartialEq` and `Hash`, but as there isn't an easy way to do
/// that now, we'll skip.
fn expression_key(expression: &Expression) -> Vec<u8> {
    expression.encode_to_vec()
}

/// Parses the output section of an aggregate (everything after `=>`).
///
/// For example, in `Aggregate[($0, $1), _ => sum($2), $0, count($2)]`,
/// this parses `sum($2), $0, count($2)`.
fn parse_aggregate_output(
    extensions: &SimpleExtensions,
    output_pair: Pair<'_, Rule>,
    grouping_expressions: &[Expression],
) -> Result<(Vec<aggregate_rel::Measure>, Vec<i32>), MessageParseError> {
    assert_eq!(output_pair.as_rule(), Rule::aggregate_output);

    // every output item is either:
    // - an expression which must already be one of `grouping_expressions`,
    // - a new function call, which we assume must be an aggregate measure.
    //
    // While it would probably be best to check if a function call was an
    // aggregate function, the substrait-explain text does not distinguish
    // between aggregate measures and grouping expressions, and the
    // `SimpleExtensions` do not have that information either (neither here in
    // the Registry nor in the actual Protobuf `ExtensionFunction` definition),
    // so we can only check whether it's a previous expression.
    let grouping_positions: HashMap<Vec<u8>, usize> = grouping_expressions
        .iter()
        .enumerate()
        .map(|(index, expression)| (expression_key(expression), index))
        .collect();

    let mut measures = Vec::new();
    let mut output_mapping = Vec::new();

    for output_item in output_pair.into_inner() {
        assert_eq!(output_item.as_rule(), Rule::expression);
        let span = output_item.as_span();
        let inner_item = unwrap_single_pair(output_item.clone());

        if inner_item.as_rule() == Rule::function_call {
            let expression = Expression::parse_pair(extensions, output_item)?;
            if let Some(&index) = grouping_positions.get(&expression_key(&expression)) {
                output_mapping.push(index as i32);
                continue;
            }

            let measure = aggregate_rel::Measure::parse_pair(extensions, inner_item)?;
            output_mapping.push(grouping_expressions.len() as i32 + measures.len() as i32);
            measures.push(measure);
            continue;
        }

        let expression = Expression::parse_pair(extensions, output_item)?;
        match grouping_positions.get(&expression_key(&expression)) {
            Some(&index) => output_mapping.push(index as i32),
            None => {
                return Err(MessageParseError::invalid(
                    "AggregateRel",
                    span,
                    "output expression is not an aggregate measure and does not match any grouping expression",
                ));
            }
        }
    }

    Ok((measures, output_mapping))
}

/// Parses the grouping section of an aggregate (everything before `=>`).
///
/// For example, in `Aggregate[($0, $1), _ => sum($2), $0, count($2)]`,
/// this parses `($0, $1), _`.
///
/// Each inner Vec is one grouping set; an empty vec represents no grouping (global aggregate).
///
/// Grammar: `aggregate_group_by = { grouping_set_list | expression_list }`
fn parse_grouping_sets(
    extensions: &SimpleExtensions,
    inner: Pair<'_, Rule>,
) -> Result<Vec<Vec<Expression>>, MessageParseError> {
    assert!(
        matches!(
            inner.as_rule(),
            Rule::expression_list | Rule::grouping_set_list
        ),
        "Expected expression_list or grouping_set_list, got {:?}",
        inner.as_rule()
    );
    match inner.as_rule() {
        Rule::expression_list => Ok(vec![parse_expression_list(extensions, inner)?]),
        Rule::grouping_set_list => inner
            .into_inner()
            .map(|pair| parse_grouping_set(extensions, pair))
            .collect(),
        _ => unreachable!(
            "Unexpected rule in aggregate_group_by: {:?}",
            inner.as_rule()
        ),
    }
}

/// Parses a single grouping set, e.g. `($0, $1)` or `_`.
///
/// Grammar: `grouping_set = { ("(" ~ expression_list ~ ")") | empty }`
fn parse_grouping_set(
    extensions: &SimpleExtensions,
    pair: Pair<'_, Rule>,
) -> Result<Vec<Expression>, MessageParseError> {
    assert_eq!(pair.as_rule(), Rule::grouping_set);
    let inner = pair
        .into_inner()
        .next()
        .expect("grouping_set must have one inner item");
    match inner.as_rule() {
        Rule::empty => Ok(vec![]),
        Rule::expression_list => parse_expression_list(extensions, inner),
        _ => unreachable!("Unexpected item in grouping_set: {:?}", inner.as_rule()),
    }
}

/// Grammar: `expression_list = { expression ~ ("," ~ expression)* }`
pub(crate) fn parse_expression_list(
    extensions: &SimpleExtensions,
    pair: Pair<'_, Rule>,
) -> Result<Vec<Expression>, MessageParseError> {
    pair.into_inner()
        .map(|expr_pair| Expression::parse_pair(extensions, expr_pair))
        .collect()
}

/// Deduplicates expressions across all sets and produces the AggregateRel's
/// protobuf fields: a flat deduplicated expression list and per-set Grouping
/// messages with index references into that list.
fn build_grouping_fields(expression_sets: &[Vec<Expression>]) -> (Vec<Grouping>, Vec<Expression>) {
    let mut expressions: Vec<Expression> = Vec::new();
    let mut seen: HashMap<Vec<u8>, u32> = HashMap::new();

    let groupings = expression_sets
        .iter()
        .map(|set| {
            let expression_references = set
                .iter()
                .map(|exp| {
                    let key = expression_key(exp);
                    let next_idx = expressions.len() as u32;
                    *seen.entry(key).or_insert_with(|| {
                        expressions.push(exp.clone());
                        next_idx
                    })
                })
                .collect();
            Grouping {
                expression_references,
                #[allow(deprecated)]
                grouping_expressions: vec![],
            }
        })
        .collect();

    (groupings, expressions)
}

impl ScopedParsePair for SortField {
    fn rule() -> Rule {
        Rule::sort_field
    }

    fn message() -> &'static str {
        "SortField"
    }

    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut iter = RuleIter::from(pair.into_inner());
        let expression_pair = iter.pop(Rule::expression);
        let expression = Expression::parse_pair(extensions, expression_pair)?;
        let direction_pair = iter.pop(Rule::sort_direction);
        let direction = sort_direction_from_str(
            direction_pair.as_str().trim_start_matches('&'),
            direction_pair.as_span(),
        )?;
        iter.done();
        Ok(SortField {
            expr: Some(expression),
            // TODO: Add support for SortKind::ComparisonFunctionReference
            sort_kind: Some(SortKind::Direction(direction as i32)),
        })
    }
}

impl RelationParsePair for SortRel {
    fn rule() -> Rule {
        Rule::sort_relation
    }

    fn message() -> &'static str {
        "SortRel"
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Sort(Box::new(self))),
        }
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let input = expect_one_child(Self::message(), &pair, input_children)?;
        let mut iter = RuleIter::from(pair.into_inner());
        let sort_field_list_pair = iter.pop(Rule::sort_field_list);
        let references_pair = iter.pop(Rule::reference_list);
        let mut sorts = Vec::new();
        for sort_field_pair in sort_field_list_pair.into_inner() {
            let sort_field = SortField::parse_pair(extensions, sort_field_pair)?;
            sorts.push(sort_field);
        }
        let (emit, output_count) = parse_emit(references_pair, input_field_count);
        let common = RelCommon {
            emit_kind: Some(emit),
            ..Default::default()
        };
        iter.done();
        Ok((
            SortRel {
                input: Some(input),
                sorts,
                common: Some(common),
                advanced_extension: None,
            },
            output_count,
        ))
    }
}

impl ScopedParsePair for CountMode {
    fn rule() -> Rule {
        Rule::fetch_value
    }
    fn message() -> &'static str {
        "CountMode"
    }
    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut arg_inner = RuleIter::from(pair.into_inner());
        let value_pair = if let Some(int_pair) = arg_inner.try_pop(Rule::integer) {
            int_pair
        } else {
            arg_inner.pop(Rule::expression)
        };
        match value_pair.as_rule() {
            Rule::integer => {
                let value = value_pair.as_str().parse::<i64>().map_err(|e| {
                    MessageParseError::invalid(
                        Self::message(),
                        value_pair.as_span(),
                        format!("Invalid integer: {e}"),
                    )
                })?;
                if value < 0 {
                    return Err(MessageParseError::invalid(
                        Self::message(),
                        value_pair.as_span(),
                        format!("Fetch limit must be non-negative, got: {value}"),
                    ));
                }
                Ok(CountMode::CountExpr(i64_literal_expr(value)))
            }
            Rule::expression => {
                let expr = Expression::parse_pair(extensions, value_pair)?;
                Ok(CountMode::CountExpr(Box::new(expr)))
            }
            _ => Err(MessageParseError::invalid(
                Self::message(),
                value_pair.as_span(),
                format!("Unexpected rule for CountMode: {:?}", value_pair.as_rule()),
            )),
        }
    }
}

fn i64_literal_expr(value: i64) -> Box<Expression> {
    Box::new(Expression {
        rex_type: Some(RexType::Literal(Literal {
            nullable: false,
            type_variation_reference: 0,
            literal_type: Some(LiteralType::I64(value)),
        })),
    })
}

impl ScopedParsePair for OffsetMode {
    fn rule() -> Rule {
        Rule::fetch_value
    }
    fn message() -> &'static str {
        "OffsetMode"
    }
    fn parse_pair(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
    ) -> Result<Self, MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let mut arg_inner = RuleIter::from(pair.into_inner());
        let value_pair = if let Some(int_pair) = arg_inner.try_pop(Rule::integer) {
            int_pair
        } else {
            arg_inner.pop(Rule::expression)
        };
        match value_pair.as_rule() {
            Rule::integer => {
                let value = value_pair.as_str().parse::<i64>().map_err(|e| {
                    MessageParseError::invalid(
                        Self::message(),
                        value_pair.as_span(),
                        format!("Invalid integer: {e}"),
                    )
                })?;
                if value < 0 {
                    return Err(MessageParseError::invalid(
                        Self::message(),
                        value_pair.as_span(),
                        format!("Fetch offset must be non-negative, got: {value}"),
                    ));
                }
                Ok(OffsetMode::OffsetExpr(i64_literal_expr(value)))
            }
            Rule::expression => {
                let expr = Expression::parse_pair(extensions, value_pair)?;
                Ok(OffsetMode::OffsetExpr(Box::new(expr)))
            }
            _ => Err(MessageParseError::invalid(
                Self::message(),
                value_pair.as_span(),
                format!("Unexpected rule for OffsetMode: {:?}", value_pair.as_rule()),
            )),
        }
    }
}

impl RelationParsePair for FetchRel {
    fn rule() -> Rule {
        Rule::fetch_relation
    }

    fn message() -> &'static str {
        "FetchRel"
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Fetch(Box::new(self))),
        }
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());
        let input = expect_one_child(Self::message(), &pair, input_children)?;
        let mut iter = RuleIter::from(pair.into_inner());

        // Extract all pairs before any validation: RuleIter's Drop panics on
        // incomplete consumption, so we must exhaust the iterator before any
        // early return. Validation runs after iter.done() below.
        let (limit_pair, offset_pair) = match iter.try_pop(Rule::fetch_named_arg_list) {
            None => {
                iter.pop(Rule::empty);
                (None, None)
            }
            Some(fetch_args_pair) => {
                let extractor =
                    ParsedNamedArgs::new(fetch_args_pair.into_inner(), Rule::fetch_named_arg)?;
                let (extractor, limit_pair) = extractor.pop("limit", Rule::fetch_value);
                let (extractor, offset_pair) = extractor.pop("offset", Rule::fetch_value);
                extractor.done()?;
                (limit_pair, offset_pair)
            }
        };

        let references_pair = iter.pop(Rule::reference_list);
        let (emit, output_count) = parse_emit(references_pair, input_field_count);
        let common = RelCommon {
            emit_kind: Some(emit),
            ..Default::default()
        };
        iter.done();

        let count_mode = limit_pair
            .map(|pair| CountMode::parse_pair(extensions, pair))
            .transpose()?;
        let offset_mode = offset_pair
            .map(|pair| OffsetMode::parse_pair(extensions, pair))
            .transpose()?;
        Ok((
            FetchRel {
                input: Some(input),
                common: Some(common),
                advanced_extension: None,
                offset_mode,
                count_mode,
            },
            output_count,
        ))
    }
}

impl ParsePair for join_rel::JoinType {
    fn rule() -> Rule {
        Rule::join_type
    }

    fn message() -> &'static str {
        "JoinType"
    }

    fn parse_pair(pair: Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());
        let join_type_str = pair.as_str().trim_start_matches('&');
        match join_type_str {
            "Inner" => join_rel::JoinType::Inner,
            "Left" => join_rel::JoinType::Left,
            "Right" => join_rel::JoinType::Right,
            "Outer" => join_rel::JoinType::Outer,
            "LeftSemi" => join_rel::JoinType::LeftSemi,
            "RightSemi" => join_rel::JoinType::RightSemi,
            "LeftAnti" => join_rel::JoinType::LeftAnti,
            "RightAnti" => join_rel::JoinType::RightAnti,
            "LeftSingle" => join_rel::JoinType::LeftSingle,
            "RightSingle" => join_rel::JoinType::RightSingle,
            "LeftMark" => join_rel::JoinType::LeftMark,
            "RightMark" => join_rel::JoinType::RightMark,
            _ => panic!("Unknown join type: {join_type_str} (this should be caught by grammar)"),
        }
    }
}

impl RelationParsePair for JoinRel {
    fn rule() -> Rule {
        Rule::join_relation
    }

    fn message() -> &'static str {
        "JoinRel"
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Join(Box::new(self))),
        }
    }

    fn parse_pair_with_context(
        extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        if input_children.len() != 2 {
            return Err(MessageParseError::invalid(
                Self::message(),
                pair.as_span(),
                format!(
                    "JoinRel should have exactly 2 input children, got {}",
                    input_children.len()
                ),
            ));
        }

        let mut children_iter = input_children.into_iter();
        let left = Box::new(children_iter.next().unwrap());
        let right = Box::new(children_iter.next().unwrap());

        let mut iter = RuleIter::from(pair.into_inner());
        let join_type = iter.parse_next::<join_rel::JoinType>();
        let condition = iter.parse_next_scoped::<Expression>(extensions)?;
        let post_join_filter = iter
            .try_pop(Rule::join_post_join_filter)
            .map(|pair| {
                let expression_pair = unwrap_single_pair(pair);
                Expression::parse_pair(extensions, expression_pair).map(Box::new)
            })
            .transpose()?;
        let references_pair = iter.pop(Rule::reference_list);
        iter.done();

        // TODO: For semi/anti joins, the direct output width differs from
        // left+right — `input_field_count` would misclassify the emit as Direct.
        // Revisit when those join types are supported.
        let (emit, output_count) = parse_emit(references_pair, input_field_count);
        let common = RelCommon {
            emit_kind: Some(emit),
            ..Default::default()
        };

        Ok((
            JoinRel {
                common: Some(common),
                left: Some(left),
                right: Some(right),
                expression: Some(Box::new(condition)),
                post_join_filter,
                r#type: join_type as i32,
                advanced_extension: None,
            },
            output_count,
        ))
    }
}

impl RelationParsePair for CrossRel {
    fn rule() -> Rule {
        Rule::cross_relation
    }

    fn message() -> &'static str {
        "CrossRel"
    }

    fn into_rel(mut self, adv_ext: Option<AdvancedExtension>) -> Rel {
        self.advanced_extension = adv_ext;
        Rel {
            rel_type: Some(RelType::Cross(Box::new(self))),
        }
    }

    fn parse_pair_with_context(
        _extensions: &SimpleExtensions,
        pair: Pair<Rule>,
        input_children: Vec<Rel>,
        input_field_count: usize,
    ) -> Result<(Self, usize), MessageParseError> {
        assert_eq!(pair.as_rule(), Self::rule());

        if input_children.len() != 2 {
            return Err(MessageParseError::invalid(
                Self::message(),
                pair.as_span(),
                format!(
                    "CrossRel should have exactly 2 input children, got {}",
                    input_children.len()
                ),
            ));
        }

        let mut children_iter = input_children.into_iter();
        let left = Box::new(children_iter.next().unwrap());
        let right = Box::new(children_iter.next().unwrap());

        let mut iter = RuleIter::from(pair.into_inner());
        let reference_list_pair = iter.pop(Rule::reference_list);
        iter.done();

        let (emit, output_count) = parse_emit(reference_list_pair, input_field_count);
        let common = RelCommon {
            emit_kind: Some(emit),
            ..Default::default()
        };

        Ok((
            CrossRel {
                common: Some(common),
                left: Some(left),
                right: Some(right),
                advanced_extension: None,
            },
            output_count,
        ))
    }
}

impl ParsePair for set_rel::SetOp {
    fn rule() -> Rule {
        Rule::set_op
    }

    fn message() -> &'static str {
        "SetOp"
    }

    fn parse_pair(pair: Pair<Rule>) -> Self {
        assert_eq!(pair.as_rule(), Self::rule());
        let set_op_str = pair.as_str().trim_start_matches('&');
        match set_op_str {
            "MinusPrimary" => set_rel::SetOp::MinusPrimary,
            "MinusPrimaryAll" => set_rel::SetOp::MinusPrimaryAll,
            "MinusMultiset" => set_rel::SetOp::MinusMultiset,
            "IntersectionPrimary" => set_rel::SetOp::IntersectionPrimary,
            "IntersectionMultiset" => set_rel::SetOp::IntersectionMultiset,
            "IntersectionMultisetAll" => set_rel::SetOp::IntersectionMultisetAll,
            "UnionDistinct" => set_rel::SetOp::UnionDistinct,
            "UnionAll" => set_rel::SetOp::UnionAll,
            _ => panic!("Unknown set op: {set_op_str} (this should be caught by grammar)"),
        }
    }
}

/// Parse a `set_relation` pair given the real per-child output widths (not
/// just their sum), so mismatched input schemas are always caught
pub(crate) fn parse_set_relation_pair(
    pair: Pair<Rule>,
    input_children: Vec<Rel>,
    child_field_counts: &[usize],
    advanced_extension: Option<AdvancedExtension>,
) -> Result<(Rel, usize), MessageParseError> {
    assert_eq!(pair.as_rule(), Rule::set_relation);

    if input_children.len() < 2 {
        return Err(MessageParseError::invalid(
            "SetRel",
            pair.as_span(),
            format!(
                "SetRel should have at least 2 input children, got {}",
                input_children.len()
            ),
        ));
    }

    // All inputs must share the same output width (Set is a pass-through
    // over a common schema, not a concatenation like Join).
    let child_width = child_field_counts[0];
    if child_field_counts.iter().any(|&w| w != child_width) {
        return Err(MessageParseError::invalid(
            "SetRel",
            pair.as_span(),
            format!(
                "SetRel inputs must all have the same number of columns, got widths {child_field_counts:?}"
            ),
        ));
    }

    let mut iter = RuleIter::from(pair.into_inner());
    let op = iter.parse_next::<set_rel::SetOp>();
    let reference_list_pair = iter.pop(Rule::reference_list);
    iter.done();

    let (emit, output_count) = parse_emit(reference_list_pair, child_width);
    let common = RelCommon {
        emit_kind: Some(emit),
        ..Default::default()
    };

    Ok((
        Rel {
            rel_type: Some(RelType::Set(SetRel {
                common: Some(common),
                inputs: input_children,
                op: op as i32,
                advanced_extension,
            })),
        },
        output_count,
    ))
}

#[cfg(test)]
mod tests {
    use pest::Parser;

    use super::*;
    use crate::fixtures::TestContext;
    use crate::parser::{ExpressionParser, Rule};

    #[test]
    fn test_parse_relation() {
        // Removed: test_parse_relation for old Relation struct
    }

    #[test]
    fn test_parse_read_relation() {
        let extensions = SimpleExtensions::default();
        let read = ReadRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::read_relation, "Read[ab.cd.ef => a:i32, b:string?]"),
            vec![],
            0,
        )
        .unwrap()
        .0;
        let names = match &read.read_type {
            Some(read_rel::ReadType::NamedTable(table)) => &table.names,
            _ => panic!("Expected NamedTable"),
        };
        assert_eq!(names, &["ab", "cd", "ef"]);
        let columns = &read
            .base_schema
            .as_ref()
            .unwrap()
            .r#struct
            .as_ref()
            .unwrap()
            .types;
        assert_eq!(columns.len(), 2);
        assert_eq!(read.common, Some(direct_common()));
    }

    #[test]
    fn test_parse_read_relation_explicit_emit() {
        let extensions = SimpleExtensions::default();
        let read = ReadRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::read_relation,
                "Read[ab.cd.ef +> a:i32, b:string?, c:i64 |> $2, $0]",
            ),
            vec![],
            0,
        )
        .unwrap();
        let (read, output_count) = read;
        assert_eq!(output_count, 2);
        let emit_kind = read.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        match emit_kind {
            EmitKind::Emit(emit) => assert_eq!(emit.output_mapping, vec![2, 0]),
            other => panic!("Expected EmitKind::Emit, got {other:?}"),
        }
    }

    #[test]
    fn test_parse_read_relation_explicit_identity_emit_not_collapsed() {
        // An explicit `|>` with an identity mapping must stay `Emit`, not
        // collapse to `Direct` - that's the whole point of `|>`.
        let extensions = SimpleExtensions::default();
        let read = ReadRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::read_relation,
                "Read[ab.cd.ef +> a:i32, b:string? |> $0, $1]",
            ),
            vec![],
            0,
        )
        .unwrap()
        .0;
        let emit_kind = read.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        assert!(
            matches!(emit_kind, EmitKind::Emit(_)),
            "Expected explicit Emit to survive even with an identity mapping, got {emit_kind:?}"
        );
    }

    #[test]
    fn test_parse_read_relation_explicit_direct() {
        let extensions = SimpleExtensions::default();
        let (read, output_count) = ReadRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::read_relation, "Read[ab.cd.ef +> a:i32, b:string?]"),
            vec![],
            0,
        )
        .unwrap();
        assert_eq!(output_count, 2);
        let emit_kind = read.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        assert!(
            matches!(emit_kind, EmitKind::Direct(_)),
            "Expected +> without |> to produce Direct, got {emit_kind:?}"
        );
    }

    #[test]
    fn test_parse_virtual_read_relation_filter() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
            .with_function(1, 10, "gt")
            .extensions;

        let read = VirtualReadRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::virtual_read_relation,
                "Read:Virtual[(1, 'alice'), filter=gt($0, 0:i64):boolean => id:i64, name:string]",
            ),
            vec![],
            0,
        )
        .unwrap()
        .0
        .0;

        match read.read_type {
            Some(read_rel::ReadType::VirtualTable(table)) => {
                assert_eq!(table.expressions.len(), 1);
            }
            other => panic!("Expected VirtualTable, got {other:?}"),
        }
        assert!(read.filter.is_some());
    }

    #[test]
    fn test_parse_read_relation_rejects_arrow_with_emit() {
        let parsed = ExpressionParser::parse(
            Rule::read_relation,
            "Read[ab.cd.ef => a:i32, b:string? |> $0, $1]",
        );
        assert!(
            parsed.is_err(),
            "Read cannot combine legacy => output with explicit |>"
        );
    }

    /// Produces a ReadRel with 3 columns: a:i32, b:string?, c:i64
    fn example_read_relation() -> ReadRel {
        let extensions = SimpleExtensions::default();
        ReadRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::read_relation,
                "Read[ab.cd.ef => a:i32, b:string?, c:i64]",
            ),
            vec![],
            0,
        )
        .unwrap()
        .0
    }

    #[test]
    fn test_parse_filter_relation() {
        let extensions = SimpleExtensions::default();
        let filter = FilterRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::filter_relation, "Filter[$1 => $0, $1, $2]"),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;
        // Identity mapping [0, 1, 2] over 3 inputs → Direct
        let emit_kind = filter.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        assert!(
            matches!(emit_kind, EmitKind::Direct(_)),
            "Expected Direct for identity emit, got {emit_kind:?}"
        );
    }

    #[test]
    fn test_parse_project_relation() {
        let extensions = SimpleExtensions::default();
        let project = ProjectRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::project_relation, "Project[$0, $1, 42]"),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        // Should have 1 expression (42) and 2 references ($0, $1)
        assert_eq!(project.expressions.len(), 1);

        let emit_kind = &project.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Output mapping should be [0, 1, 3]. References are 0-2; expression is 3.
        assert_eq!(emit, &[0, 1, 3]);
    }

    #[test]
    fn test_parse_project_relation_complex() {
        let extensions = SimpleExtensions::default();
        let project = ProjectRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::project_relation, "Project[42, $0, 100, $2, $1]"),
            vec![example_read_relation().into_rel(None)],
            5, // Assume 5 input fields
        )
        .unwrap()
        .0;

        // Should have 2 expressions (42, 100) and 3 references ($0, $2, $1)
        assert_eq!(project.expressions.len(), 2);

        let emit_kind = &project.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Direct mapping: [input_fields..., 42, 100] (input fields first, then expressions)
        // Output mapping: [5, 0, 6, 2, 1] (to get: 42, $0, 100, $2, $1)
        assert_eq!(emit, &[5, 0, 6, 2, 1]);
    }

    #[test]
    fn test_parse_aggregate_relation() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "sum")
            .with_function(1, 11, "count")
            .extensions;

        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::aggregate_relation,
                "Aggregate[($0, $1), _ => sum($2):i64, $0, count($2):i64]",
            ),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        // Should have 2 group-by sets ($0, $1) and an empty group, and emit 2 measures (sum($2), count($2))
        assert_eq!(aggregate.grouping_expressions.len(), 2);
        assert_eq!(aggregate.groupings[0].expression_references.len(), 2);
        assert_eq!(aggregate.groupings.len(), 2);
        assert_eq!(aggregate.measures.len(), 2);

        let emit_kind = &aggregate
            .common
            .as_ref()
            .unwrap()
            .emit_kind
            .as_ref()
            .unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Output mapping should be [2, 0, 3] (measures and group-by fields in order)
        // sum($2) -> 2, $0 -> 0, count($2) -> 3
        assert_eq!(emit, &[2, 0, 3]);
    }

    #[test]
    fn test_parse_aggregate_relation_maintain_column_order() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "sum")
            .with_function(1, 11, "count")
            .extensions;

        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::aggregate_relation,
                "Aggregate[$0 => sum($1):i64, $0, count($1):i64]",
            ),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        // Should have 1 group-by field ($0) and 2 measures (sum($1), count($1))
        assert_eq!(aggregate.grouping_expressions.len(), 1);
        assert_eq!(aggregate.groupings.len(), 1);
        assert_eq!(aggregate.measures.len(), 2);

        let emit_kind = &aggregate
            .common
            .as_ref()
            .unwrap()
            .emit_kind
            .as_ref()
            .unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Output mapping should be [1, 0, 2] (grouping fields + measures)
        assert_eq!(emit, &[1, 0, 2]);
    }

    #[test]
    fn test_parse_aggregate_relation_simple() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "sum")
            .extensions;

        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::aggregate_relation, "Aggregate[$2, $0 => sum($1):i64]"),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        assert_eq!(aggregate.grouping_expressions.len(), 2);
        assert_eq!(aggregate.groupings.len(), 1);
        // expression_references must be positions [0, 1], not raw field indices [2, 0]
        assert_eq!(aggregate.groupings[0].expression_references, vec![0, 1]);
    }

    #[test]
    fn test_parse_aggregate_relation_output_reference_uses_grouping_position() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "sum")
            .extensions;

        // grouping_expressions ends up as [$2, $0] (grouping's textual order),
        // so `$0`/`$2` in the output must resolve to their grouping
        // positions (1 and 0 respectively), not their literal reference
        // indices. Output order is swapped relative to grouping order so
        // the resulting mapping isn't an identity (which would collapse to
        // EmitKind::Direct and give us nothing to assert on).
        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::aggregate_relation,
                "Aggregate[$2, $0 => $0, $2, sum($1):i64]",
            ),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        assert_eq!(aggregate.grouping_expressions.len(), 2);
        let emit_kind = &aggregate
            .common
            .as_ref()
            .unwrap()
            .emit_kind
            .as_ref()
            .unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // direct schema is [$2, $0, sum($1)]; output is [$0, $2, sum($1)] -> [1, 0, 2]
        assert_eq!(emit, &[1, 0, 2]);
    }

    #[test]
    fn test_parse_aggregate_relation_output_not_in_grouping_expressions_errors() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "sum")
            .extensions;

        // $1 is not part of the grouping expressions ($0) and is not an
        // aggregate measure, so it has no valid slot in the output schema.
        let result = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::aggregate_relation, "Aggregate[$0 => $1, sum($1):i64]"),
            vec![example_read_relation().into_rel(None)],
            3,
        );

        let error = result.unwrap_err();
        assert!(error.to_string().contains(
            "output expression is not an aggregate measure and does not match any grouping expression"
        ));
    }

    #[test]
    fn test_parse_aggregate_relation_output_with_literal_expression() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "sum")
            .extensions;

        // The first grouping expression is a literal, not a reference or a
        // function call, so it can only match the output item by structural
        // equality (encode_to_vec).
        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::aggregate_relation,
                "Aggregate[42, $0 => $0, 42, sum($1):i64]",
            ),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        assert_eq!(aggregate.grouping_expressions.len(), 2);
        let emit_kind = &aggregate
            .common
            .as_ref()
            .unwrap()
            .emit_kind
            .as_ref()
            .unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // direct schema is [42, $0, sum($1)]; output is [$0, 42, sum($1)] -> [1, 0, 2]
        assert_eq!(emit, &[1, 0, 2]);
    }

    #[test]
    fn test_parse_aggregate_relation_function_call_is_not_a_measure() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
            .with_urn(2, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "eq")
            .with_function(2, 11, "sum")
            .extensions;

        // `eq($0, $1)` is both the grouping expression and repeated
        // verbatim in the output, so it must resolve to the grouping
        // position rather than being parsed as a second, duplicate
        // aggregate measure.
        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::aggregate_relation,
                "Aggregate[eq($0, $1):boolean => eq($0, $1):boolean, sum($2):i64]",
            ),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        assert_eq!(aggregate.grouping_expressions.len(), 1);
        assert_eq!(aggregate.measures.len(), 1);
        let emit_kind = &aggregate
            .common
            .as_ref()
            .unwrap()
            .emit_kind
            .as_ref()
            .unwrap();
        assert!(matches!(emit_kind, EmitKind::Direct(_)));
    }

    #[test]
    fn test_parse_aggregate_relation_global_aggregate() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 10, "sum")
            .with_function(1, 11, "count")
            .extensions;

        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::aggregate_relation,
                "Aggregate[_ => sum($0):i64, count($1):i64]",
            ),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        // Should have 0 group-by fields and 2 measures
        assert_eq!(aggregate.grouping_expressions.len(), 0);
        assert_eq!(aggregate.groupings.len(), 1);
        assert_eq!(aggregate.groupings[0].expression_references.len(), 0);
        assert_eq!(aggregate.measures.len(), 2);

        // Identity mapping [0, 1] over 2 outputs (0 grouping + 2 measures) → Direct
        let emit_kind = aggregate
            .common
            .as_ref()
            .unwrap()
            .emit_kind
            .as_ref()
            .unwrap();
        assert!(
            matches!(emit_kind, EmitKind::Direct(_)),
            "Expected Direct for identity emit, got {emit_kind:?}"
        );
    }

    #[test]
    fn test_parse_aggregate_relation_grouping_sets() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_aggregate.yaml")
            .with_function(1, 11, "count")
            .extensions;

        let read_rel = ReadRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::read_relation,
                "Read[ab.cd.ef => a:i32, b:string?, c:i64, d:i64]",
            ),
            vec![],
            0,
        )
        .unwrap()
        .0;

        let aggregate = AggregateRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::aggregate_relation,
                "Aggregate[($0, $1, $2), ($2, $0), ($1), _ => $0, $1, $2, count($3):i64]",
            ),
            vec![read_rel.into_rel(None)],
            4,
        )
        .unwrap()
        .0;

        assert_eq!(aggregate.grouping_expressions.len(), 3);
        assert_eq!(aggregate.groupings.len(), 4);
        // ($0, $1, $2) -> [0, 1, 2]
        assert_eq!(aggregate.groupings[0].expression_references, vec![0, 1, 2]);
        // ($2, $0) -> [2, 0] (reuses indices, different order)
        assert_eq!(aggregate.groupings[1].expression_references, vec![2, 0]);
        // ($1) -> [1]
        assert_eq!(aggregate.groupings[2].expression_references, vec![1]);
        // _ -> empty
        assert!(aggregate.groupings[3].expression_references.is_empty());
        assert_eq!(aggregate.measures.len(), 1);
    }

    #[test]
    fn test_fetch_relation_positive_values() {
        let extensions = SimpleExtensions::default();

        // Test valid positive values should work
        let fetch_rel = FetchRel::parse_pair_with_context(
            &extensions,
            parse_exact(Rule::fetch_relation, "Fetch[limit=10, offset=5 => $0]"),
            vec![example_read_relation().into_rel(None)],
            3,
        )
        .unwrap()
        .0;

        // Verify the limit and offset values are correct
        assert_eq!(
            fetch_rel.count_mode,
            Some(CountMode::CountExpr(i64_literal_expr(10)))
        );
        assert_eq!(
            fetch_rel.offset_mode,
            Some(OffsetMode::OffsetExpr(i64_literal_expr(5)))
        );
    }

    #[test]
    fn test_fetch_relation_negative_limit_rejected() {
        let extensions = SimpleExtensions::default();

        // Test that fetch relations with negative limits are properly rejected
        let parsed_result = ExpressionParser::parse(Rule::fetch_relation, "Fetch[limit=-5 => $0]");
        if let Ok(mut pairs) = parsed_result {
            let pair = pairs.next().unwrap();
            if pair.as_str() == "Fetch[limit=-5 => $0]" {
                // Full parse succeeded, now test that validation catches the negative value
                let result = FetchRel::parse_pair_with_context(
                    &extensions,
                    pair,
                    vec![example_read_relation().into_rel(None)],
                    3,
                );
                assert!(result.is_err());
                let error_msg = result.unwrap_err().to_string();
                assert!(error_msg.contains("Fetch limit must be non-negative"));
            } else {
                // If grammar doesn't fully support negative values, that's also acceptable
                // since it would prevent negative values at parse time
                println!("Grammar prevents negative limit values at parse time");
            }
        } else {
            // Grammar doesn't support negative values in fetch context
            println!("Grammar prevents negative limit values at parse time");
        }
    }

    #[test]
    fn test_fetch_relation_negative_offset_rejected() {
        let extensions = SimpleExtensions::default();

        // Test that fetch relations with negative offsets are properly rejected
        let parsed_result =
            ExpressionParser::parse(Rule::fetch_relation, "Fetch[offset=-10 => $0]");
        if let Ok(mut pairs) = parsed_result {
            let pair = pairs.next().unwrap();
            if pair.as_str() == "Fetch[offset=-10 => $0]" {
                // Full parse succeeded, now test that validation catches the negative value
                let result = FetchRel::parse_pair_with_context(
                    &extensions,
                    pair,
                    vec![example_read_relation().into_rel(None)],
                    3,
                );
                assert!(result.is_err());
                let error_msg = result.unwrap_err().to_string();
                assert!(error_msg.contains("Fetch offset must be non-negative"));
            } else {
                // If grammar doesn't fully support negative values, that's also acceptable
                println!("Grammar prevents negative offset values at parse time");
            }
        } else {
            // Grammar doesn't support negative values in fetch context
            println!("Grammar prevents negative offset values at parse time");
        }
    }

    #[test]
    fn test_parse_join_relation() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
            .with_function(1, 10, "eq")
            .extensions;

        let left_rel = example_read_relation().into_rel(None);
        let right_rel = example_read_relation().into_rel(None);

        let join = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&Inner, eq($0, $3):boolean => $0, $1, $3, $4]",
            ),
            vec![left_rel, right_rel],
            6, // left (3) + right (3) = 6 total input fields
        )
        .unwrap()
        .0;

        // Should be an Inner join
        assert_eq!(join.r#type, join_rel::JoinType::Inner as i32);

        // Should have left and right relations
        assert!(join.left.is_some());
        assert!(join.right.is_some());

        // Should have a join condition
        assert!(join.expression.is_some());
        assert!(join.post_join_filter.is_none());

        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Output mapping should be [0, 1, 3, 4] (selected columns)
        assert_eq!(emit, &[0, 1, 3, 4]);
    }

    #[test]
    fn test_parse_join_relation_post_join_filter() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
            .with_function(1, 10, "eq")
            .with_function(1, 11, "gt")
            .extensions;

        let left_rel = example_read_relation().into_rel(None);
        let right_rel = example_read_relation().into_rel(None);

        let join = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&RightSemi, eq($0, $3):boolean, post_filter=gt($1, 100:i32):boolean => $0, $1]",
            ),
            vec![left_rel, right_rel],
            6,
        )
        .unwrap()
        .0;

        assert_eq!(join.r#type, join_rel::JoinType::RightSemi as i32);
        assert!(join.expression.is_some());
        assert!(join.post_join_filter.is_some());

        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        assert_eq!(emit, &[0, 1]);
    }

    #[test]
    fn test_parse_join_relation_left_outer() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
            .with_function(1, 10, "eq")
            .extensions;

        let left_rel = example_read_relation().into_rel(None);
        let right_rel = example_read_relation().into_rel(None);

        let join = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&Left, eq($0, $3):boolean => $0, $1, $2]",
            ),
            vec![left_rel, right_rel],
            6,
        )
        .unwrap()
        .0;

        // Should be a Left join
        assert_eq!(join.r#type, join_rel::JoinType::Left as i32);

        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Output mapping should be [0, 1, 2]
        assert_eq!(emit, &[0, 1, 2]);
    }

    #[test]
    fn test_parse_join_relation_left_semi() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
            .with_function(1, 10, "eq")
            .extensions;

        let left_rel = example_read_relation().into_rel(None);
        let right_rel = example_read_relation().into_rel(None);

        let join = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&LeftSemi, eq($0, $3):boolean => $0, $1]",
            ),
            vec![left_rel, right_rel],
            6,
        )
        .unwrap()
        .0;

        // Should be a LeftSemi join
        assert_eq!(join.r#type, join_rel::JoinType::LeftSemi as i32);

        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Output mapping should be [0, 1] (only left columns for semi join)
        assert_eq!(emit, &[0, 1]);
    }

    #[test]
    fn test_parse_join_relation_right_semi() {
        let extensions = TestContext::new()
            .with_urn(1, "https://github.com/substrait-io/substrait/blob/main/extensions/functions_comparison.yaml")
            .with_function(1, 10, "eq")
            .extensions;

        let left_rel = example_read_relation().into_rel(None);
        let right_rel = example_read_relation().into_rel(None);

        let join = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&RightSemi, eq($0, $3):boolean => $0, $1]",
            ),
            vec![left_rel, right_rel],
            6,
        )
        .unwrap()
        .0;

        // Should be a RightSemi join
        assert_eq!(join.r#type, join_rel::JoinType::RightSemi as i32);

        let emit_kind = &join.common.as_ref().unwrap().emit_kind.as_ref().unwrap();
        let emit = match emit_kind {
            EmitKind::Emit(emit) => &emit.output_mapping,
            _ => panic!("Expected EmitKind::Emit, got {emit_kind:?}"),
        };
        // Output mapping should be [0, 1] over right-semi direct output columns.
        assert_eq!(emit, &[0, 1]);
    }

    #[test]
    fn test_parse_join_relation_requires_two_children() {
        let extensions = SimpleExtensions::default();

        // Test with 0 children
        let result = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&Inner, eq($0, $1):boolean => $0, $1]",
            ),
            vec![],
            0,
        );
        assert!(result.is_err());

        // Test with 1 child
        let result = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&Inner, eq($0, $1):boolean => $0, $1]",
            ),
            vec![example_read_relation().into_rel(None)],
            3,
        );
        assert!(result.is_err());

        // Test with 3 children
        let result = JoinRel::parse_pair_with_context(
            &extensions,
            parse_exact(
                Rule::join_relation,
                "Join[&Inner, eq($0, $1):boolean => $0, $1]",
            ),
            vec![
                example_read_relation().into_rel(None),
                example_read_relation().into_rel(None),
                example_read_relation().into_rel(None),
            ],
            9,
        );
        assert!(result.is_err());
    }

    fn parse_exact(rule: Rule, input: &'_ str) -> Pair<'_, Rule> {
        let mut pairs = ExpressionParser::parse(rule, input).unwrap();
        assert_eq!(pairs.as_str(), input);
        let pair = pairs.next().unwrap();
        assert_eq!(pairs.next(), None);
        pair
    }
}